feat: Add ToolRejectError and tool request/result callbacks (#101)

* feat: Add ToolRejectError and tool request/result callbacks

* Add to API reference and test for console output

Author: cpsievert

CHANGELOG.md

@@ -9,9 +9,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [UNRELEASED]
 
+### New features
+
+* New `.on_tool_request()` and `.on_tool_result()` methods register callbacks that fire when a tool is requested or produces a result. These callbacks can be used to implement custom logging or other actions when tools are called, without modifying the tool function (#101).
+* New `ToolRejectError` exception can be thrown from tool request/result callbacks or from within a tool function itself to prevent the tool from executing. Moreover, this exception will provide some context for the the LLM to know that the tool didn't produce a result because it was rejected. (#101)
+
 ### Improvements
 
-* The `CHATLAS_LOG` environment variable nows enable logs for the relevant model provider. It now also supports a lovel of `debug` in addition to `info`. (#97)
+* The `CHATLAS_LOG` environment variable now enables logs for the relevant model provider. It now also supports a level of `debug` in addition to `info`. (#97)
 * `ChatSnowflake()` now supports tool calling. (#98)
 * `Chat` instances can now be deep copied, which is useful for forking the chat session. (#96)
 

chatlas/__init__.py

@@ -16,7 +16,7 @@
 from ._provider import Provider
 from ._snowflake import ChatSnowflake
 from ._tokens import token_usage
-from ._tools import Tool
+from ._tools import Tool, ToolRejectError
 from ._turn import Turn
 
 try:
@@ -51,6 +51,7 @@
     "Provider",
     "token_usage",
     "Tool",
+    "ToolRejectError",
     "Turn",
     "types",
 )

chatlas/_callbacks.py

@@ -0,0 +1,56 @@
+from collections import OrderedDict
+from typing import Any, Callable
+
+from ._utils import is_async_callable
+
+
+class CallbackManager:
+    def __init__(self) -> None:
+        self._callbacks: dict[str, Callable[..., Any]] = OrderedDict()
+        self._id: int = 1
+
+    def add(self, callback: Callable[..., Any]) -> Callable[[], None]:
+        callback_id = self._next_id()
+        self._callbacks[callback_id] = callback
+
+        def _rm_callback() -> None:
+            self._callbacks.pop(callback_id, None)
+
+        return _rm_callback
+
+    def invoke(self, *args: Any, **kwargs: Any) -> None:
+        if not self._callbacks:
+            return
+
+        # Invoke in reverse insertion order
+        for callback_id in reversed(list(self._callbacks.keys())):
+            callback = self._callbacks[callback_id]
+            if is_async_callable(callback):
+                raise RuntimeError(
+                    "Can't use async callbacks with `.chat()`/`.stream()`."
+                    "Async callbacks can only be used with `.chat_async()`/`.stream_async()`."
+                )
+            callback(*args, **kwargs)
+
+    async def invoke_async(self, *args: Any, **kwargs: Any) -> None:
+        if not self._callbacks:
+            return
+
+        # Invoke in reverse insertion order
+        for callback_id in reversed(list(self._callbacks.keys())):
+            callback = self._callbacks[callback_id]
+            if is_async_callable(callback):
+                await callback(*args, **kwargs)
+            else:
+                callback(*args, **kwargs)
+
+    def count(self) -> int:
+        return len(self._callbacks)
+
+    def get_callbacks(self) -> list[Callable[..., Any]]:
+        return list(self._callbacks.values())
+
+    def _next_id(self) -> str:
+        current_id = self._id
+        self._id += 1
+        return str(current_id)

chatlas/_chat.py

@@ -26,6 +26,7 @@
 
 from pydantic import BaseModel
 
+from ._callbacks import CallbackManager
 from ._content import (
     Content,
     ContentJson,
@@ -42,7 +43,7 @@
 )
 from ._logging import log_tool_error
 from ._provider import Provider
-from ._tools import Tool
+from ._tools import Tool, ToolRejectError
 from ._turn import Turn, user_turn
 from ._typing_extensions import TypedDict
 from ._utils import html_escape, wrap_async
@@ -96,6 +97,8 @@ def __init__(
         self.provider = provider
         self._turns: list[Turn] = list(turns or [])
         self._tools: dict[str, Tool] = {}
+        self._on_tool_request_callbacks = CallbackManager()
+        self._on_tool_result_callbacks = CallbackManager()
         self._current_display: Optional[MarkdownDisplay] = None
         self._echo_options: EchoDisplayOptions = {
             "rich_markdown": {},
@@ -988,6 +991,53 @@ def add(a: int, b: int) -> int:
         tool = Tool(func, model=model)
         self._tools[tool.name] = tool
 
+    def on_tool_request(self, callback: Callable[[ContentToolRequest], None]):
+        """
+        Register a callback for a tool request event.
+
+        A tool request event occurs when the assistant requests a tool to be
+        called on its behalf. Before invoking the tool, `on_tool_request`
+        handlers are called with the relevant `ContentToolRequest` object. This
+        is useful if you want to handle tool requests in a custom way, such as
+        requiring logging them or requiring user approval before invoking the
+        tool
+
+        Parameters
+        ----------
+        callback
+            A function to be called when a tool request event occurs.
+            This function must have a single argument, which will be the
+            tool request (i.e., a `ContentToolRequest` object).
+
+        Returns
+        -------
+        A callable that can be used to remove the callback later.
+        """
+        return self._on_tool_request_callbacks.add(callback)
+
+    def on_tool_result(self, callback: Callable[[ContentToolResult], None]):
+        """
+        Register a callback for a tool result event.
+
+        A tool result event occurs when a tool has been invoked and the
+        result is ready to be provided to the assistant. After the tool
+        has been invoked, `on_tool_result` handlers are called with the
+        relevant `ContentToolResult` object. This is useful if you want to
+        handle tool results in a custom way such as logging them.
+
+        Parameters
+        ----------
+        callback
+            A function to be called when a tool result event occurs.
+            This function must have a single argument, which will be the
+            tool result (i.e., a `ContentToolResult` object).
+
+        Returns
+        -------
+        A callable that can be used to remove the callback later.
+        """
+        return self._on_tool_result_callbacks.add(callback)
+
     @property
     def current_display(self) -> Optional[MarkdownDisplay]:
         """
@@ -1418,28 +1468,43 @@ def _invoke_tool(self, x: ContentToolRequest) -> ContentToolResult:
             e = RuntimeError(f"Unknown tool: {x.name}")
             return ContentToolResult(value=None, error=e, request=x)
 
-        args = x.arguments
-
+        # First, invoke the request callbacks. If a ToolRejectError is raised,
+        # treat it like a tool failure (i.e., gracefully handle it).
+        result: ContentToolResult | None = None
         try:
-            if isinstance(args, dict):
-                result = func(**args)
-            else:
-                result = func(args)
+            self._on_tool_request_callbacks.invoke(x)
+        except ToolRejectError as e:
+            result = ContentToolResult(value=None, error=e, request=x)
+
+        # Invoke the tool (if it hasn't been rejected).
+        if result is None:
+            try:
+                if isinstance(x.arguments, dict):
+                    res = func(**x.arguments)
+                else:
+                    res = func(x.arguments)
 
-            if not isinstance(result, ContentToolResult):
-                result = ContentToolResult(value=result)
+                if isinstance(res, ContentToolResult):
+                    result = res
+                else:
+                    result = ContentToolResult(value=res)
+
+                result.request = x
+            except Exception as e:
+                result = ContentToolResult(value=None, error=e, request=x)
 
-            result.request = x
-            return result
-        except Exception as e:
+        # If we've captured an error, notify and log it.
+        if result.error:
             warnings.warn(
                 f"Calling tool '{x.name}' led to an error.",
                 ToolFailureWarning,
                 stacklevel=2,
             )
             traceback.print_exc()
-            log_tool_error(x.name, str(args), e)
-            return ContentToolResult(value=None, error=e, request=x)
+            log_tool_error(x.name, str(x.arguments), result.error)
+
+        self._on_tool_result_callbacks.invoke(result)
+        return result
 
     async def _invoke_tool_async(self, x: ContentToolRequest) -> ContentToolResult:
         tool_def = self._tools.get(x.name, None)
@@ -1454,28 +1519,43 @@ async def _invoke_tool_async(self, x: ContentToolRequest) -> ContentToolResult:
             e = RuntimeError(f"Unknown tool: {x.name}")
             return ContentToolResult(value=None, error=e, request=x)
 
-        args = x.arguments
-
+        # First, invoke the request callbacks. If a ToolRejectError is raised,
+        # treat it like a tool failure (i.e., gracefully handle it).
+        result: ContentToolResult | None = None
         try:
-            if isinstance(args, dict):
-                result = await func(**args)
-            else:
-                result = await func(args)
+            await self._on_tool_request_callbacks.invoke_async(x)
+        except ToolRejectError as e:
+            result = ContentToolResult(value=None, error=e, request=x)
+
+        # Invoke the tool (if it hasn't been rejected).
+        if result is None:
+            try:
+                if isinstance(x.arguments, dict):
+                    res = await func(**x.arguments)
+                else:
+                    res = await func(x.arguments)
 
-            if not isinstance(result, ContentToolResult):
-                result = ContentToolResult(value=result)
+                if isinstance(res, ContentToolResult):
+                    result = res
+                else:
+                    result = ContentToolResult(value=res)
 
-            result.request = x
-            return result
-        except Exception as e:
+                result.request = x
+            except Exception as e:
+                result = ContentToolResult(value=None, error=e, request=x)
+
+        # If we've captured an error, notify and log it.
+        if result.error:
             warnings.warn(
                 f"Calling tool '{x.name}' led to an error.",
                 ToolFailureWarning,
                 stacklevel=2,
             )
             traceback.print_exc()
-            log_tool_error(x.name, str(args), e)
-            return ContentToolResult(value=None, error=e, request=x)
+            log_tool_error(x.name, str(x.arguments), result.error)
+
+        await self._on_tool_result_callbacks.invoke_async(result)
+        return result
 
     def _markdown_display(self, echo: EchoOptions) -> ChatMarkdownDisplay:
         """

chatlas/_tools.py

@@ -8,7 +8,10 @@
 
 from . import _utils
 
-__all__ = ("Tool",)
+__all__ = (
+    "Tool",
+    "ToolRejectError",
+)
 
 if TYPE_CHECKING:
     from openai.types.chat import ChatCompletionToolParam
@@ -47,6 +50,61 @@ def __init__(
         self.name = self.schema["function"]["name"]
 
 
+class ToolRejectError(Exception):
+    """
+    Error to represent a tool call being rejected.
+
+    This error is meant to be raised when an end user has chosen to deny a tool
+    call. It can be raised in a tool function or in a `.on_tool_request()`
+    callback registered via a :class:`~chatlas.Chat`. When used in the callback,
+    the tool call is rejected before the tool function is invoked.
+
+    Parameters
+    ----------
+    reason
+        A string describing the reason for rejecting the tool call. This will be
+        included in the error message passed to the LLM. In addition to the
+        reason, the error message will also include "Tool call rejected." to
+        indicate that the tool call was not processed.
+
+    Raises
+    -------
+    ToolRejectError
+        An error with a message informing the LLM that the tool call was
+        rejected (and the reason why).
+
+    Examples
+    --------
+    >>> import os
+    >>> import chatlas as ctl
+    >>>
+    >>> chat = ctl.ChatOpenAI()
+    >>>
+    >>> def list_files():
+    ...     "List files in the user's current directory"
+    ...     while True:
+    ...         allow = input(
+    ...             "Would you like to allow access to your current directory? (yes/no): "
+    ...         )
+    ...         if allow.lower() == "yes":
+    ...             return os.listdir(".")
+    ...         elif allow.lower() == "no":
+    ...             raise ctl.ToolRejectError(
+    ...                 "The user has chosen to disallow the tool call."
+    ...             )
+    ...         else:
+    ...             print("Please answer with 'yes' or 'no'.")
+    >>>
+    >>> chat.register_tool(list_files)
+    >>> chat.chat("What files are available in my current directory?")
+    """
+
+    def __init__(self, reason: str = "The user has chosen to disallow the tool call."):
+        message = f"Tool call rejected. {reason}"
+        super().__init__(message)
+        self.message = message
+
+
 def func_to_schema(
     func: Callable[..., Any] | Callable[..., Awaitable[Any]],
     model: Optional[type[BaseModel]] = None,

docs/_quarto.yml

@@ -113,6 +113,7 @@ quartodoc:
       desc: Add context to python function before registering it as a tool.
       contents:
         - Tool
+        - ToolRejectError
     - title: Turns
       desc: A provider-agnostic representation of content generated during an assistant/user turn.
       contents: