First pass at an interface for surfacing tool calls

Author: cpsievert

chatlas/__init__.py

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

chatlas/_chat.py

@@ -39,7 +39,7 @@
 )
 from ._logging import log_tool_error
 from ._provider import Provider
-from ._tools import Tool
+from ._tools import Stringable, Tool, ToolResult
 from ._turn import Turn, user_turn
 from ._typing_extensions import TypedDict
 from ._utils import html_escape, wrap_async
@@ -96,6 +96,9 @@ def __init__(
             "rich_console": {},
             "css_styles": {},
         }
+        self._on_tool_request_default: Optional[
+            Callable[[ContentToolRequest], Stringable]
+        ] = None
 
     def get_turns(
         self,
@@ -658,7 +661,7 @@ def stream(
             kwargs=kwargs,
         )
 
-        def wrapper() -> Generator[str, None, None]:
+        def wrapper() -> Generator[Stringable, None, None]:
             with display:
                 for chunk in generator:
                     yield chunk
@@ -695,7 +698,7 @@ async def stream_async(
 
         display = self._markdown_display(echo=echo)
 
-        async def wrapper() -> AsyncGenerator[str, None]:
+        async def wrapper() -> AsyncGenerator[Stringable, None]:
             with display:
                 async for chunk in self._chat_impl_async(
                     turn,
@@ -831,6 +834,7 @@ def register_tool(
         self,
         func: Callable[..., Any] | Callable[..., Awaitable[Any]],
         *,
+        on_request: Optional[Callable[[ContentToolRequest], Stringable]] = None,
         model: Optional[type[BaseModel]] = None,
     ):
         """
@@ -900,16 +904,49 @@ def add(a: int, b: int) -> int:
         ----------
         func
             The function to be invoked when the tool is called.
+        on_request
+            A callable that will be passed a :class:`~chatlas.ContentToolRequest`
+            when the tool is requested. If defined, and the callable returns a
+            stringable object, that value will be yielded to the chat as a part
+            of the response.
         model
             A Pydantic model that describes the input parameters for the function.
             If not provided, the model will be inferred from the function's type hints.
             The primary reason why you might want to provide a model in
             Note that the name and docstring of the model takes precedence over the
             name and docstring of the function.
         """
-        tool = Tool(func, model=model)
+        tool = Tool(func, on_request=on_request, model=model)
         self._tools[tool.name] = tool
 
+    def on_tool_request(
+        self,
+        func: Callable[[ContentToolRequest], Stringable],
+    ):
+        """
+        Register a default function to be invoked when a tool is requested.
+
+        This function will be invoked if a tool is requested that does not have
+        a specific `on_request` function defined.
+
+        Parameters
+        ----------
+        func
+            A callable that will be passed a :class:`~chatlas.ContentToolRequest`
+            when the tool is requested. If defined, and the callable returns a
+            stringable object, that value will be yielded to the chat as a part
+            of the response.
+        """
+        self._on_tool_request_default = func
+
+    def _on_tool_request(self, req: ContentToolRequest) -> Stringable | None:
+        tool_def = self._tools.get(req.name, None)
+        if tool_def and tool_def.on_request:
+            return tool_def.on_request(req)
+        if self._on_tool_request_default:
+            return self._on_tool_request_default(req)
+        return None
+
     def export(
         self,
         filename: str | Path,
@@ -1040,7 +1077,7 @@ def _chat_impl(
         display: MarkdownDisplay,
         stream: bool,
         kwargs: Optional[SubmitInputArgsT] = None,
-    ) -> Generator[str, None, None]:
+    ) -> Generator[Stringable, None, None]:
         user_turn_result: Turn | None = user_turn
         while user_turn_result is not None:
             for chunk in self._submit_turns(
@@ -1051,7 +1088,24 @@ def _chat_impl(
                 kwargs=kwargs,
             ):
                 yield chunk
-            user_turn_result = self._invoke_tools()
+
+            turn = self.get_last_turn(role="assistant")
+            assert turn is not None
+            user_turn_result = None
+
+            results: list[ContentToolResult] = []
+            for x in turn.contents:
+                if isinstance(x, ContentToolRequest):
+                    req = self._on_tool_request(x)
+                    if req is not None:
+                        yield req
+                    result, output = self._invoke_tool_request(x)
+                    if output is not None:
+                        yield output
+                    results.append(result)
+
+            if results:
+                user_turn_result = Turn("user", results)
 
     async def _chat_impl_async(
         self,
@@ -1060,7 +1114,7 @@ async def _chat_impl_async(
         display: MarkdownDisplay,
         stream: bool,
         kwargs: Optional[SubmitInputArgsT] = None,
-    ) -> AsyncGenerator[str, None]:
+    ) -> AsyncGenerator[Stringable, None]:
         user_turn_result: Turn | None = user_turn
         while user_turn_result is not None:
             async for chunk in self._submit_turns_async(
@@ -1071,7 +1125,24 @@ async def _chat_impl_async(
                 kwargs=kwargs,
             ):
                 yield chunk
-            user_turn_result = await self._invoke_tools_async()
+
+            turn = self.get_last_turn(role="assistant")
+            assert turn is not None
+            user_turn_result = None
+
+            results: list[ContentToolResult] = []
+            for x in turn.contents:
+                if isinstance(x, ContentToolRequest):
+                    req = self._on_tool_request(x)
+                    if req is not None:
+                        yield req
+                    result, output = await self._invoke_tool_request_async(x)
+                    if output is not None:
+                        yield output
+                    results.append(result)
+
+            if results:
+                user_turn_result = Turn("user", results)
 
     def _submit_turns(
         self,
@@ -1085,7 +1156,7 @@ def _submit_turns(
         if any(x._is_async for x in self._tools.values()):
             raise ValueError("Cannot use async tools in a synchronous chat")
 
-        def emit(text: str | Content):
+        def emit(text: Stringable):
             display.update(str(text))
 
         emit("<br>\n\n")
@@ -1148,7 +1219,7 @@ async def _submit_turns_async(
         data_model: type[BaseModel] | None = None,
         kwargs: Optional[SubmitInputArgsT] = None,
     ) -> AsyncGenerator[str, None]:
-        def emit(text: str | Content):
+        def emit(text: Stringable):
             display.update(str(text))
 
         emit("<br>\n\n")
@@ -1202,88 +1273,62 @@ def emit(text: str | Content):
 
         self._turns.extend([user_turn, turn])
 
-    def _invoke_tools(self) -> Turn | None:
-        turn = self.get_last_turn()
-        if turn is None:
-            return None
-
-        results: list[ContentToolResult] = []
-        for x in turn.contents:
-            if isinstance(x, ContentToolRequest):
-                tool_def = self._tools.get(x.name, None)
-                func = tool_def.func if tool_def is not None else None
-                results.append(self._invoke_tool(func, x.arguments, x.id))
-
-        if not results:
-            return None
+    def _invoke_tool_request(
+        self, x: ContentToolRequest
+    ) -> tuple[ContentToolResult, Stringable]:
+        tool_def = self._tools.get(x.name, None)
+        func = tool_def.func if tool_def is not None else None
 
-        return Turn("user", results)
-
-    async def _invoke_tools_async(self) -> Turn | None:
-        turn = self.get_last_turn()
-        if turn is None:
-            return None
-
-        results: list[ContentToolResult] = []
-        for x in turn.contents:
-            if isinstance(x, ContentToolRequest):
-                tool_def = self._tools.get(x.name, None)
-                func = None
-                if tool_def:
-                    if tool_def._is_async:
-                        func = tool_def.func
-                    else:
-                        func = wrap_async(tool_def.func)
-                results.append(await self._invoke_tool_async(func, x.arguments, x.id))
-
-        if not results:
-            return None
-
-        return Turn("user", results)
-
-    @staticmethod
-    def _invoke_tool(
-        func: Callable[..., Any] | None,
-        arguments: object,
-        id_: str,
-    ) -> ContentToolResult:
         if func is None:
-            return ContentToolResult(id_, value=None, error="Unknown tool")
+            return ContentToolResult(x.id, value=None, error="Unknown tool"), None
 
         name = func.__name__
 
         try:
-            if isinstance(arguments, dict):
-                result = func(**arguments)
+            if isinstance(x.arguments, dict):
+                result = func(**x.arguments)
             else:
-                result = func(arguments)
+                result = func(x.arguments)
 
-            return ContentToolResult(id_, value=result, error=None, name=name)
+            value, output = (result, None)
+            if isinstance(result, ToolResult):
+                value, output = (result.assistant, result.output)
+
+            return ContentToolResult(x.id, value=value, error=None, name=name), output
         except Exception as e:
-            log_tool_error(name, str(arguments), e)
-            return ContentToolResult(id_, value=None, error=str(e), name=name)
+            log_tool_error(name, str(x.arguments), e)
+            return ContentToolResult(x.id, value=None, error=str(e), name=name), None
+
+    async def _invoke_tool_request_async(
+        self, x: ContentToolRequest
+    ) -> tuple[ContentToolResult, Stringable]:
+        tool_def = self._tools.get(x.name, None)
+        func = None
+        if tool_def:
+            if tool_def._is_async:
+                func = tool_def.func
+            else:
+                func = wrap_async(tool_def.func)
 
-    @staticmethod
-    async def _invoke_tool_async(
-        func: Callable[..., Awaitable[Any]] | None,
-        arguments: object,
-        id_: str,
-    ) -> ContentToolResult:
         if func is None:
-            return ContentToolResult(id_, value=None, error="Unknown tool")
+            return ContentToolResult(x.id, value=None, error="Unknown tool"), None
 
         name = func.__name__
 
         try:
-            if isinstance(arguments, dict):
-                result = await func(**arguments)
+            if isinstance(x.arguments, dict):
+                result = await func(**x.arguments)
             else:
-                result = await func(arguments)
+                result = await func(x.arguments)
+
+            value, output = (result, None)
+            if isinstance(result, ToolResult):
+                value, output = (result.assistant, result.output)
 
-            return ContentToolResult(id_, value=result, error=None, name=name)
+            return ContentToolResult(x.id, value=value, error=None, name=name), output
         except Exception as e:
-            log_tool_error(func.__name__, str(arguments), e)
-            return ContentToolResult(id_, value=None, error=str(e), name=name)
+            log_tool_error(func.__name__, str(x.arguments), e)
+            return ContentToolResult(x.id, value=None, error=str(e), name=name), None
 
     def _markdown_display(
         self, echo: Literal["text", "all", "none"]
@@ -1378,15 +1423,15 @@ class ChatResponse:
         still be retrieved (via the `content` attribute).
     """
 
-    def __init__(self, generator: Generator[str, None]):
+    def __init__(self, generator: Generator[Stringable, None]):
         self._generator = generator
         self.content: str = ""
 
     def __iter__(self) -> Iterator[str]:
         return self
 
     def __next__(self) -> str:
-        chunk = next(self._generator)
+        chunk = str(next(self._generator))
         self.content += chunk  # Keep track of accumulated content
         return chunk
 
@@ -1430,15 +1475,15 @@ class ChatResponseAsync:
         still be retrieved (via the `content` attribute).
     """
 
-    def __init__(self, generator: AsyncGenerator[str, None]):
+    def __init__(self, generator: AsyncGenerator[Stringable, None]):
         self._generator = generator
         self.content: str = ""
 
     def __aiter__(self) -> AsyncIterator[str]:
         return self
 
     async def __anext__(self) -> str:
-        chunk = await self._generator.__anext__()
+        chunk = str(await self._generator.__anext__())
         self.content += chunk  # Keep track of accumulated content
         return chunk