kosong.tooling

  1from abc import ABC, abstractmethod
  2from asyncio import Future
  3from typing import Any, ClassVar, Protocol, Self, cast, override, runtime_checkable
  4
  5import jsonschema
  6import pydantic
  7from pydantic import BaseModel, GetCoreSchemaHandler, model_validator
  8from pydantic.json_schema import GenerateJsonSchema
  9from pydantic_core import core_schema
 10
 11from kosong.message import ContentPart, ToolCall
 12from kosong.utils.jsonschema import deref_json_schema
 13from kosong.utils.typing import JsonType
 14
 15type ParametersType = dict[str, Any]
 16
 17
 18class Tool(BaseModel):
 19    """The definition of a tool that can be recognized by the model."""
 20
 21    name: str
 22    """The name of the tool."""
 23
 24    description: str
 25    """The description of the tool."""
 26
 27    parameters: ParametersType
 28    """The parameters of the tool, in JSON Schema format."""
 29
 30    @model_validator(mode="after")
 31    def _validate_parameters(self) -> Self:
 32        jsonschema.validate(self.parameters, jsonschema.Draft202012Validator.META_SCHEMA)
 33        return self
 34
 35
 36class DisplayBlock(BaseModel, ABC):
 37    """
 38    A block of content to be displayed to the user.
 39
 40    Similar to `ContentPart`, but scoped to user-facing UI.
 41    `ContentPart` is for model-facing message content; `DisplayBlock` is for tool/UI extensions.
 42
 43    Unlike `ContentPart`, Kosong users may directly subclass `DisplayBlock` to define custom
 44    display blocks for their applications.
 45    """
 46
 47    __display_block_registry: ClassVar[dict[str, type["DisplayBlock"]]] = {}
 48
 49    type: str
 50    ...  # to be added by subclasses
 51
 52    def __init_subclass__(cls, **kwargs: Any) -> None:
 53        super().__init_subclass__(**kwargs)
 54
 55        invalid_subclass_error_msg = (
 56            f"DisplayBlock subclass {cls.__name__} must have a `type` field of type `str`"
 57        )
 58
 59        type_value = getattr(cls, "type", None)
 60        if type_value is None or not isinstance(type_value, str):
 61            raise ValueError(invalid_subclass_error_msg)
 62
 63        cls.__display_block_registry[type_value] = cls
 64
 65    @classmethod
 66    def __get_pydantic_core_schema__(
 67        cls, source_type: Any, handler: GetCoreSchemaHandler
 68    ) -> core_schema.CoreSchema:
 69        # If we're dealing with the base DisplayBlock class, use custom validation
 70        if cls.__name__ == "DisplayBlock":
 71
 72            def validate_display_block(value: Any) -> Any:
 73                # if it's already an instance of a DisplayBlock subclass, return it
 74                if hasattr(value, "__class__") and issubclass(value.__class__, cls):
 75                    return value
 76
 77                # if it's a dict with a type field, dispatch to the appropriate subclass
 78                if isinstance(value, dict) and "type" in value:
 79                    type_value: Any | None = cast(dict[str, Any], value).get("type")
 80                    if not isinstance(type_value, str):
 81                        raise ValueError(f"Cannot validate {value} as DisplayBlock")
 82                    target_class = cls.__display_block_registry.get(type_value)
 83                    if target_class is None:
 84                        data = {k: v for k, v in cast(dict[str, Any], value).items() if k != "type"}
 85                        return UnknownDisplayBlock.model_validate(
 86                            {"type": type_value, "data": data}
 87                        )
 88                    return target_class.model_validate(value)
 89
 90                raise ValueError(f"Cannot validate {value} as DisplayBlock")
 91
 92            return core_schema.no_info_plain_validator_function(validate_display_block)
 93
 94        # for subclasses, use the default schema
 95        return handler(source_type)
 96
 97
 98class UnknownDisplayBlock(DisplayBlock):
 99    """Fallback display block for unknown types."""
100
101    type: str = "unknown"
102    data: JsonType
103
104
105class BriefDisplayBlock(DisplayBlock):
106    """A brief display block with plain string content."""
107
108    type: str = "brief"
109    text: str
110
111
112class ToolReturnValue(BaseModel):
113    """The return type of a callable tool."""
114
115    is_error: bool
116    """Whether the tool call resulted in an error."""
117
118    # For model
119    output: str | list[ContentPart]
120    """The output content returned by the tool."""
121    message: str
122    """An explanatory message to be given to the model."""
123
124    # For user
125    display: list[DisplayBlock]
126    """The content blocks to be displayed to the user."""
127
128    # For debugging/testing
129    extras: dict[str, JsonType] | None = None
130
131    @property
132    def brief(self) -> str:
133        """Get the brief display block data, if any."""
134        for block in self.display:
135            if isinstance(block, BriefDisplayBlock):
136                return block.text
137        return ""
138
139
140class ToolOk(ToolReturnValue):
141    """Subclass of `ToolReturnValue` representing a successful tool call."""
142
143    def __init__(
144        self,
145        *,
146        output: str | ContentPart | list[ContentPart],
147        message: str = "",
148        brief: str = "",
149    ) -> None:
150        super().__init__(
151            is_error=False,
152            output=([output] if isinstance(output, ContentPart) else output),
153            message=message,
154            display=[BriefDisplayBlock(text=brief)] if brief else [],
155        )
156
157
158class ToolError(ToolReturnValue):
159    """Subclass of `ToolReturnValue` representing a failed tool call."""
160
161    def __init__(
162        self, *, message: str, brief: str, output: str | ContentPart | list[ContentPart] = ""
163    ):
164        super().__init__(
165            is_error=True,
166            output=([output] if isinstance(output, ContentPart) else output),
167            message=message,
168            display=[BriefDisplayBlock(text=brief)] if brief else [],
169        )
170
171
172class CallableTool(Tool, ABC):
173    """
174    The abstract base class of tools that can be called as callables.
175
176    The tool will be called with the arguments provided in the `ToolCall`.
177    If the arguments are given as a JSON array, it will be unpacked into positional arguments.
178    If the arguments are given as a JSON object, it will be unpacked into keyword arguments.
179    Otherwise, the arguments will be passed as a single argument.
180    """
181
182    @property
183    def base(self) -> Tool:
184        """The base tool definition."""
185        return self
186
187    async def call(self, arguments: JsonType) -> ToolReturnValue:
188        from kosong.tooling.error import ToolValidateError
189
190        try:
191            jsonschema.validate(arguments, self.parameters)
192        except jsonschema.ValidationError as e:
193            return ToolValidateError(str(e))
194
195        if isinstance(arguments, list):
196            ret = await self.__call__(*arguments)
197        elif isinstance(arguments, dict):
198            ret = await self.__call__(**arguments)
199        else:
200            ret = await self.__call__(arguments)
201        if not isinstance(ret, ToolReturnValue):  # type: ignore[reportUnnecessaryIsInstance]
202            # let's do not trust the return type of the tool
203            ret = ToolError(
204                message=f"Invalid return type: {type(ret)}",
205                brief="Invalid return type",
206            )
207        return ret
208
209    @abstractmethod
210    async def __call__(self, *args: Any, **kwargs: Any) -> ToolReturnValue:
211        """
212        @public
213
214        The implementation of the callable tool.
215        """
216        ...
217
218
219class _GenerateJsonSchemaNoTitles(GenerateJsonSchema):
220    """Custom JSON schema generator that omits titles."""
221
222    @override
223    def field_title_should_be_set(self, schema) -> bool:  # type: ignore[reportMissingParameterType]
224        return False
225
226    @override
227    def _update_class_schema(self, json_schema, cls, config) -> None:  # type: ignore[reportMissingParameterType]
228        super()._update_class_schema(json_schema, cls, config)
229        json_schema.pop("title", None)
230
231
232class CallableTool2[Params: BaseModel](ABC):
233    """
234    The abstract base class of tools that can be called as callables, with typed parameters.
235
236    The tool will be called with the arguments provided in the `ToolCall`.
237    The arguments must be a JSON object, and will be validated by Pydantic to the `Params` type.
238    """
239
240    name: str
241    """The name of the tool."""
242    description: str
243    """The description of the tool."""
244    params: type[Params]
245    """The Pydantic model type of the tool parameters."""
246
247    def __init__(
248        self,
249        name: str | None = None,
250        description: str | None = None,
251        params: type[Params] | None = None,
252    ) -> None:
253        cls = self.__class__
254
255        self.name = name or getattr(cls, "name", "")
256        if not self.name:
257            raise ValueError(
258                "Tool name must be provided either as class variable or constructor argument"
259            )
260        if not isinstance(self.name, str):  # type: ignore[reportUnnecessaryIsInstance]
261            raise ValueError("Tool name must be a string")
262
263        self.description = description or getattr(cls, "description", "")
264        if not self.description:
265            raise ValueError(
266                "Tool description must be provided either as class variable or constructor argument"
267            )
268        if not isinstance(self.description, str):  # type: ignore[reportUnnecessaryIsInstance]
269            raise ValueError("Tool description must be a string")
270
271        self.params = params or getattr(cls, "params", None)  # type: ignore
272        if not self.params:
273            raise ValueError(
274                "Tool param must be provided either as class variable or constructor argument"
275            )
276        if not isinstance(self.params, type) or not issubclass(self.params, BaseModel):  # type: ignore[reportUnnecessaryIsInstance]
277            raise ValueError("Tool params must be a subclass of pydantic.BaseModel")
278
279        self._base = Tool(
280            name=self.name,
281            description=self.description,
282            parameters=deref_json_schema(
283                self.params.model_json_schema(schema_generator=_GenerateJsonSchemaNoTitles)
284            ),
285        )
286
287    @property
288    def base(self) -> Tool:
289        """The base tool definition."""
290        return self._base
291
292    async def call(self, arguments: JsonType) -> ToolReturnValue:
293        from kosong.tooling.error import ToolValidateError
294
295        try:
296            params = self.params.model_validate(arguments)
297        except pydantic.ValidationError as e:
298            return ToolValidateError(str(e))
299
300        ret = await self.__call__(params)
301        if not isinstance(ret, ToolReturnValue):  # type: ignore[reportUnnecessaryIsInstance]
302            # let's do not trust the return type of the tool
303            ret = ToolError(
304                message=f"Invalid return type: {type(ret)}",
305                brief="Invalid return type",
306            )
307        return ret
308
309    @abstractmethod
310    async def __call__(self, params: Params) -> ToolReturnValue:
311        """
312        @public
313
314        The implementation of the callable tool.
315        """
316        ...
317
318
319class ToolResult(BaseModel):
320    """The result of a tool call."""
321
322    tool_call_id: str
323    """The ID of the tool call."""
324    return_value: ToolReturnValue
325    """The actual return value of the tool call."""
326
327
328ToolResultFuture = Future[ToolResult]
329type HandleResult = ToolResultFuture | ToolResult
330
331
332@runtime_checkable
333class Toolset(Protocol):
334    """
335    The interface of toolsets that can register tools and handle tool calls.
336    """
337
338    @property
339    def tools(self) -> list[Tool]:
340        """The list of tool definitions registered in this toolset."""
341        ...
342
343    def handle(self, tool_call: ToolCall) -> HandleResult:
344        """
345        Handle a tool call.
346        The result of the tool call, or the async future of the result, should be returned.
347        The result should be a `ToolReturnValue`.
348
349        This method MUST NOT do any blocking operations because it will be called during
350        consuming the chat response stream.
351        This method MUST NOT raise any exception except for `asyncio.CancelledError`. Any other
352        error should be returned as a `ToolReturnValue` with `is_error=True`.
353        """
354        ...