kosong.tooling.simple

  1import asyncio
  2import inspect
  3import json
  4from collections.abc import Iterable
  5from typing import TYPE_CHECKING, Any, Self
  6
  7from kosong.message import ToolCall
  8from kosong.tooling import (
  9    CallableTool,
 10    CallableTool2,
 11    HandleResult,
 12    Tool,
 13    ToolResult,
 14    ToolReturnValue,
 15    Toolset,
 16)
 17from kosong.tooling.error import (
 18    ToolNotFoundError,
 19    ToolParseError,
 20    ToolRuntimeError,
 21)
 22from kosong.utils.typing import JsonType
 23
 24if TYPE_CHECKING:
 25
 26    def type_check(
 27        simple: "SimpleToolset",
 28    ):
 29        _: Toolset = simple
 30
 31
 32type ToolType = CallableTool | CallableTool2[Any]
 33"""The tool type that can be added to the `SimpleToolset`."""
 34
 35
 36class SimpleToolset:
 37    """A simple toolset that can handle tool calls concurrently."""
 38
 39    _tool_dict: dict[str, ToolType]
 40
 41    def __init__(self, tools: Iterable[ToolType] | None = None):
 42        """Initialize the simple toolset with an optional iterable of tools."""
 43        self._tool_dict = {}
 44        if tools:
 45            for tool in tools:
 46                self += tool
 47
 48    def __iadd__(self, tool: ToolType) -> Self:
 49        """
 50        @public
 51        Add a tool to the toolset.
 52        """
 53        return_annotation = inspect.signature(tool.__call__).return_annotation
 54
 55        # Check if the return annotation is ToolReturnValue
 56        # Supports both actual type and string annotation (when using
 57        # `from __future__ import annotations`)
 58        if return_annotation is ToolReturnValue:
 59            pass
 60        elif isinstance(return_annotation, str):
 61            # String annotation - check if it matches ToolReturnValue
 62            # Accept any suffix of the full module path, e.g.:
 63            #   "ToolReturnValue", "tooling.ToolReturnValue", "kosong.tooling.ToolReturnValue"
 64            full_name = f"{ToolReturnValue.__module__}.ToolReturnValue"
 65            full_parts = full_name.split(".")
 66            if not any(
 67                return_annotation == ".".join(full_parts[i:]) for i in range(len(full_parts))
 68            ):
 69                raise TypeError(
 70                    f"Expected tool `{tool.name}` to return `ToolReturnValue`, "
 71                    f"but got `{return_annotation}`"
 72                )
 73        else:
 74            raise TypeError(
 75                f"Expected tool `{tool.name}` to return `ToolReturnValue`, "
 76                f"but got `{return_annotation}`"
 77            )
 78
 79        self._tool_dict[tool.name] = tool
 80        return self
 81
 82    def __add__(self, tool: ToolType) -> "SimpleToolset":
 83        """
 84        @public
 85        Return a new toolset with the given tool added.
 86        """
 87        new_toolset = SimpleToolset()
 88        new_toolset._tool_dict = self._tool_dict.copy()
 89        new_toolset += tool
 90        return new_toolset
 91
 92    def add(self, tool: ToolType) -> None:
 93        """
 94        @public
 95        Add a tool to the toolset.
 96        """
 97        self += tool
 98
 99    def remove(self, tool_name: str) -> None:
100        """
101        @public
102        Remove a tool from the toolset.
103        """
104        if tool_name not in self._tool_dict:
105            raise KeyError(f"Tool `{tool_name}` not found in the toolset.")
106        del self._tool_dict[tool_name]
107
108    @property
109    def tools(self) -> list[Tool]:
110        return [tool.base for tool in self._tool_dict.values()]
111
112    def handle(self, tool_call: ToolCall) -> HandleResult:
113        if tool_call.function.name not in self._tool_dict:
114            return ToolResult(
115                tool_call_id=tool_call.id,
116                return_value=ToolNotFoundError(tool_call.function.name),
117            )
118
119        tool = self._tool_dict[tool_call.function.name]
120
121        try:
122            arguments: JsonType = json.loads(tool_call.function.arguments or "{}", strict=False)
123        except json.JSONDecodeError as e:
124            return ToolResult(tool_call_id=tool_call.id, return_value=ToolParseError(str(e)))
125
126        async def _call():
127            try:
128                ret = await tool.call(arguments)
129                return ToolResult(tool_call_id=tool_call.id, return_value=ret)
130            except Exception as e:
131                return ToolResult(tool_call_id=tool_call.id, return_value=ToolRuntimeError(str(e)))
132
133        return asyncio.create_task(_call())