Cleanup

Author: cpsievert

chatlas/_chat.py

@@ -1880,7 +1880,7 @@ async def cleanup_mcp_tools(self, names: Optional[Sequence[str]] = None):
 
     def register_tool(
         self,
-        func: Callable[..., Any] | Callable[..., Awaitable[Any]] | Tool | "ToolBuiltIn",
+        func: Callable[..., Any] | Callable[..., Awaitable[Any]] | Tool | ToolBuiltIn,
         *,
         force: bool = False,
         name: Optional[str] = None,
@@ -1974,11 +1974,7 @@ def add(a: int, b: int) -> int:
         ValueError
             If a tool with the same name already exists and `force` is `False`.
         """
-        if isinstance(func, ToolBuiltIn):
-            # ToolBuiltIn objects are stored directly without conversion
-            tool = func
-            tool_name = tool.name
-        elif isinstance(func, Tool):
+        if isinstance(func, Tool):
             name = name or func.name
             annotations = annotations or func.annotations
             if model is not None:
@@ -1987,17 +1983,20 @@ def add(a: int, b: int) -> int:
                 )
             func = func.func
             tool = Tool.from_func(func, name=name, model=model, annotations=annotations)
-            tool_name = tool.name
         else:
-            tool = Tool.from_func(func, name=name, model=model, annotations=annotations)
-            tool_name = tool.name
+            if isinstance(func, ToolBuiltIn):
+                tool = func
+            else:
+                tool = Tool.from_func(
+                    func, name=name, model=model, annotations=annotations
+                )
 
-        if tool_name in self._tools and not force:
+        if tool.name in self._tools and not force:
             raise ValueError(
-                f"Tool with name '{tool_name}' is already registered. "
+                f"Tool with name '{tool.name}' is already registered. "
                 "Set `force=True` to overwrite it."
             )
-        self._tools[tool_name] = tool
+        self._tools[tool.name] = tool
 
     def get_tools(self) -> list[Tool | ToolBuiltIn]:
         """
@@ -2530,7 +2529,7 @@ def _submit_turns(
         data_model: type[BaseModel] | None = None,
         kwargs: Optional[SubmitInputArgsT] = None,
     ) -> Generator[str, None, None]:
-        if any(hasattr(x, "_is_async") and x._is_async for x in self._tools.values()):
+        if any(isinstance(x, Tool) and x._is_async for x in self._tools.values()):
             raise ValueError("Cannot use async tools in a synchronous chat")
 
         def emit(text: str | Content):
@@ -2700,7 +2699,6 @@ def _invoke_tool(self, request: ContentToolRequest):
             return
 
         if isinstance(tool, ToolBuiltIn):
-            # Built-in tools are handled by the provider, not invoked directly
             yield self._handle_tool_error_result(
                 request,
                 error=RuntimeError(
@@ -2710,8 +2708,6 @@ def _invoke_tool(self, request: ContentToolRequest):
             )
             return
 
-        func = tool.func
-
         # 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
@@ -2723,9 +2719,9 @@ def _invoke_tool(self, request: ContentToolRequest):
 
         try:
             if isinstance(request.arguments, dict):
-                res = func(**request.arguments)
+                res = tool.func(**request.arguments)
             else:
-                res = func(request.arguments)
+                res = tool.func(request.arguments)
 
             # Normalize res as a generator of results.
             if not inspect.isgenerator(res):
@@ -2760,7 +2756,6 @@ async def _invoke_tool_async(self, request: ContentToolRequest):
             return
 
         if isinstance(tool, ToolBuiltIn):
-            # Built-in tools are handled by the provider, not invoked directly
             yield self._handle_tool_error_result(
                 request,
                 error=RuntimeError(
@@ -2770,11 +2765,6 @@ async def _invoke_tool_async(self, request: ContentToolRequest):
             )
             return
 
-        if tool._is_async:
-            func = tool.func
-        else:
-            func = wrap_async(tool.func)
-
         # 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
@@ -2784,6 +2774,11 @@ async def _invoke_tool_async(self, request: ContentToolRequest):
             yield self._handle_tool_error_result(request, e)
             return
 
+        if tool._is_async:
+            func = tool.func
+        else:
+            func = wrap_async(tool.func)
+
         # Invoke the tool (if it hasn't been rejected).
         try:
             if isinstance(request.arguments, dict):

chatlas/_content.py

@@ -109,14 +109,7 @@ def from_tool(cls, tool: "Tool | ToolBuiltIn") -> "ToolInfo":
         from ._tools import ToolBuiltIn
 
         if isinstance(tool, ToolBuiltIn):
-            # For built-in tools, extract info from the definition
-            defn = tool.definition
-            return cls(
-                name=tool.name,
-                description=defn.get("description", ""),
-                parameters=defn.get("parameters", {}),
-                annotations=None,
-            )
+            return cls(name=tool.name, description=tool.name, parameters={})
         else:
             # For regular tools, extract from schema
             func_schema = tool.schema["function"]

chatlas/_mcp_manager.py

@@ -23,7 +23,7 @@ class SessionInfo(ABC):
 
     # Primary derived attributes
     session: ClientSession | None = None
-    tools: dict[str, Tool] = field(default_factory=dict)
+    tools: dict[str, Tool | ToolBuiltIn] = field(default_factory=dict)
 
     # Background task management
     ready_event: asyncio.Event = field(default_factory=asyncio.Event)
@@ -74,7 +74,7 @@ async def request_tools(self) -> None:
             tool_names = tool_names.difference(exclude)
 
         # Apply namespace and convert to chatlas.Tool instances
-        self_tools: dict[str, Tool] = {}
+        self_tools: dict[str, Tool | ToolBuiltIn] = {}
         for tool in response.tools:
             if tool.name not in tool_names:
                 continue

chatlas/_provider_anthropic.py

@@ -15,7 +15,6 @@
 )
 
 import orjson
-from openai.types.chat import ChatCompletionToolParam
 from pydantic import BaseModel
 
 from ._chat import Chat
@@ -38,7 +37,7 @@
     StandardModelParams,
 )
 from ._tokens import get_token_pricing
-from ._tools import Tool, basemodel_to_param_schema
+from ._tools import Tool, ToolBuiltIn, basemodel_to_param_schema
 from ._turn import AssistantTurn, SystemTurn, Turn, UserTurn, user_turn
 from ._utils import split_http_client_kwargs
 
@@ -48,7 +47,7 @@
         MessageParam,
         RawMessageStreamEvent,
         TextBlock,
-        ToolParam,
+        ToolUnionParam,
         ToolUseBlock,
     )
     from anthropic.types.cache_control_ephemeral_param import CacheControlEphemeralParam
@@ -304,7 +303,7 @@ def chat_perform(
         *,
         stream: Literal[False],
         turns: list[Turn],
-        tools: dict[str, Tool],
+        tools: dict[str, Tool | ToolBuiltIn],
         data_model: Optional[type[BaseModel]] = None,
         kwargs: Optional["SubmitInputArgs"] = None,
     ): ...
@@ -315,7 +314,7 @@ def chat_perform(
         *,
         stream: Literal[True],
         turns: list[Turn],
-        tools: dict[str, Tool],
+        tools: dict[str, Tool | ToolBuiltIn],
         data_model: Optional[type[BaseModel]] = None,
         kwargs: Optional["SubmitInputArgs"] = None,
     ): ...
@@ -325,7 +324,7 @@ def chat_perform(
         *,
         stream: bool,
         turns: list[Turn],
-        tools: dict[str, Tool],
+        tools: dict[str, Tool | ToolBuiltIn],
         data_model: Optional[type[BaseModel]] = None,
         kwargs: Optional["SubmitInputArgs"] = None,
     ):
@@ -338,7 +337,7 @@ async def chat_perform_async(
         *,
         stream: Literal[False],
         turns: list[Turn],
-        tools: dict[str, Tool],
+        tools: dict[str, Tool | ToolBuiltIn],
         data_model: Optional[type[BaseModel]] = None,
         kwargs: Optional["SubmitInputArgs"] = None,
     ): ...
@@ -349,7 +348,7 @@ async def chat_perform_async(
         *,
         stream: Literal[True],
         turns: list[Turn],
-        tools: dict[str, Tool],
+        tools: dict[str, Tool | ToolBuiltIn],
         data_model: Optional[type[BaseModel]] = None,
         kwargs: Optional["SubmitInputArgs"] = None,
     ): ...
@@ -359,7 +358,7 @@ async def chat_perform_async(
         *,
         stream: bool,
         turns: list[Turn],
-        tools: dict[str, Tool],
+        tools: dict[str, Tool | ToolBuiltIn],
         data_model: Optional[type[BaseModel]] = None,
         kwargs: Optional["SubmitInputArgs"] = None,
     ):
@@ -370,12 +369,12 @@ def _chat_perform_args(
         self,
         stream: bool,
         turns: list[Turn],
-        tools: dict[str, Tool],
+        tools: dict[str, Tool | ToolBuiltIn],
         data_model: Optional[type[BaseModel]] = None,
         kwargs: Optional["SubmitInputArgs"] = None,
     ) -> "SubmitInputArgs":
         tool_schemas = [
-            self._anthropic_tool_schema(tool.schema) for tool in tools.values()
+            self._anthropic_tool_schema(tool) for tool in tools.values()
         ]
 
         # If data extraction is requested, add a "mock" tool with parameters inferred from the data model
@@ -395,7 +394,7 @@ def _structured_tool_call(**kwargs: Any):
                 },
             }
 
-            tool_schemas.append(self._anthropic_tool_schema(data_model_tool.schema))
+            tool_schemas.append(self._anthropic_tool_schema(data_model_tool))
 
             if stream:
                 stream = False
@@ -497,7 +496,7 @@ def value_tokens(self, completion):
     def token_count(
         self,
         *args: Content | str,
-        tools: dict[str, Tool],
+        tools: dict[str, Tool | ToolBuiltIn],
         data_model: Optional[type[BaseModel]],
     ) -> int:
         kwargs = self._token_count_args(
@@ -511,7 +510,7 @@ def token_count(
     async def token_count_async(
         self,
         *args: Content | str,
-        tools: dict[str, Tool],
+        tools: dict[str, Tool | ToolBuiltIn],
         data_model: Optional[type[BaseModel]],
     ) -> int:
         kwargs = self._token_count_args(
@@ -525,7 +524,7 @@ async def token_count_async(
     def _token_count_args(
         self,
         *args: Content | str,
-        tools: dict[str, Tool],
+        tools: dict[str, Tool | ToolBuiltIn],
         data_model: Optional[type[BaseModel]],
     ) -> dict[str, Any]:
         turn = user_turn(*args)
@@ -655,11 +654,14 @@ def _as_content_block(content: Content) -> "ContentBlockParam":
         raise ValueError(f"Unknown content type: {type(content)}")
 
     @staticmethod
-    def _anthropic_tool_schema(schema: "ChatCompletionToolParam") -> "ToolParam":
-        fn = schema["function"]
+    def _anthropic_tool_schema(tool: "Tool | ToolBuiltIn") -> "ToolUnionParam":
+        if isinstance(tool, ToolBuiltIn):
+            return tool.definition  # type: ignore
+
+        fn = tool.schema["function"]
         name = fn["name"]
 
-        res: "ToolParam" = {
+        res: "ToolUnionParam" = {
             "name": name,
             "input_schema": {
                 "type": "object",