File size: 11,350 Bytes
e062359 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 | from __future__ import annotations
import logging
from abc import ABC, abstractmethod
from typing import Any, Union, Generic, TypeVar, Callable, Iterable, Coroutine, cast, overload
from inspect import iscoroutinefunction
from typing_extensions import TypeAlias, override
import pydantic
import docstring_parser
from pydantic import BaseModel
from ... import _compat
from ..._utils import is_dict
from ..._compat import cached_property
from ..._models import TypeAdapter
from ...types.beta import BetaToolParam, BetaToolUnionParam
from ..._utils._utils import CallableT
from ...types.tool_param import InputSchema
from ...types.beta.beta_tool_result_block_param import Content as BetaContent
log = logging.getLogger(__name__)
BetaFunctionToolResultType: TypeAlias = Union[str, Iterable[BetaContent]]
Function = Callable[..., BetaFunctionToolResultType]
FunctionT = TypeVar("FunctionT", bound=Function)
AsyncFunction = Callable[..., Coroutine[Any, Any, BetaFunctionToolResultType]]
AsyncFunctionT = TypeVar("AsyncFunctionT", bound=AsyncFunction)
class BetaBuiltinFunctionTool(ABC):
@abstractmethod
def to_dict(self) -> BetaToolUnionParam: ...
@abstractmethod
def call(self, input: object) -> BetaFunctionToolResultType: ...
@property
def name(self) -> str:
raw = self.to_dict()
if "mcp_server_name" in raw:
return raw["mcp_server_name"]
return raw["name"]
class BetaAsyncBuiltinFunctionTool(ABC):
@abstractmethod
def to_dict(self) -> BetaToolUnionParam: ...
@abstractmethod
async def call(self, input: object) -> BetaFunctionToolResultType: ...
@property
def name(self) -> str:
raw = self.to_dict()
if "mcp_server_name" in raw:
return raw["mcp_server_name"]
return raw["name"]
class BaseFunctionTool(Generic[CallableT]):
func: CallableT
"""The function this tool is wrapping"""
name: str
"""The name of the tool that will be sent to the API"""
description: str
input_schema: InputSchema
def __init__(
self,
func: CallableT,
*,
name: str | None = None,
description: str | None = None,
input_schema: InputSchema | type[BaseModel] | None = None,
defer_loading: bool | None = None,
) -> None:
if _compat.PYDANTIC_V1:
raise RuntimeError("Tool functions are only supported with Pydantic v2")
self.func = func
self._func_with_validate = pydantic.validate_call(func)
self.name = name or func.__name__
self._defer_loading = defer_loading
self.description = description or self._get_description_from_docstring()
if input_schema is not None:
if isinstance(input_schema, type):
self.input_schema: InputSchema = input_schema.model_json_schema()
else:
self.input_schema = input_schema
else:
self.input_schema = self._create_schema_from_function()
@property
def __call__(self) -> CallableT:
return self.func
def to_dict(self) -> BetaToolParam:
defn: BetaToolParam = {
"name": self.name,
"description": self.description,
"input_schema": self.input_schema,
}
if self._defer_loading is not None:
defn["defer_loading"] = self._defer_loading
return defn
@cached_property
def _parsed_docstring(self) -> docstring_parser.Docstring:
return docstring_parser.parse(self.func.__doc__ or "")
def _get_description_from_docstring(self) -> str:
"""Extract description from parsed docstring."""
if self._parsed_docstring.short_description:
description = self._parsed_docstring.short_description
if self._parsed_docstring.long_description:
description += f"\n\n{self._parsed_docstring.long_description}"
return description
return ""
def _create_schema_from_function(self) -> InputSchema:
"""Create JSON schema from function signature using pydantic."""
from pydantic_core import CoreSchema
from pydantic.json_schema import JsonSchemaValue, GenerateJsonSchema
from pydantic_core.core_schema import ArgumentsParameter
class CustomGenerateJsonSchema(GenerateJsonSchema):
def __init__(self, *, func: Callable[..., Any], parsed_docstring: Any) -> None:
super().__init__()
self._func = func
self._parsed_docstring = parsed_docstring
def __call__(self, *_args: Any, **_kwds: Any) -> "CustomGenerateJsonSchema": # noqa: ARG002
return self
@override
def kw_arguments_schema(
self,
arguments: "list[ArgumentsParameter]",
var_kwargs_schema: CoreSchema | None,
) -> JsonSchemaValue:
schema = super().kw_arguments_schema(arguments, var_kwargs_schema)
if schema.get("type") != "object":
return schema
properties = schema.get("properties")
if not properties or not is_dict(properties):
return schema
# Add parameter descriptions from docstring
for param in self._parsed_docstring.params:
prop_schema = properties.get(param.arg_name)
if not prop_schema or not is_dict(prop_schema):
continue
if param.description and "description" not in prop_schema:
prop_schema["description"] = param.description
return schema
schema_generator = CustomGenerateJsonSchema(func=self.func, parsed_docstring=self._parsed_docstring)
return self._adapter.json_schema(schema_generator=schema_generator) # type: ignore
@cached_property
def _adapter(self) -> TypeAdapter[Any]:
return TypeAdapter(self._func_with_validate)
class BetaFunctionTool(BaseFunctionTool[FunctionT]):
def call(self, input: object) -> BetaFunctionToolResultType:
if iscoroutinefunction(self.func):
raise RuntimeError("Cannot call a coroutine function synchronously. Use `@async_tool` instead.")
if not is_dict(input):
raise TypeError(f"Input must be a dictionary, got {type(input).__name__}")
try:
return self._func_with_validate(**cast(Any, input))
except pydantic.ValidationError as e:
raise ValueError(f"Invalid arguments for function {self.name}") from e
class BetaAsyncFunctionTool(BaseFunctionTool[AsyncFunctionT]):
async def call(self, input: object) -> BetaFunctionToolResultType:
if not iscoroutinefunction(self.func):
raise RuntimeError("Cannot call a synchronous function asynchronously. Use `@tool` instead.")
if not is_dict(input):
raise TypeError(f"Input must be a dictionary, got {type(input).__name__}")
try:
return await self._func_with_validate(**cast(Any, input))
except pydantic.ValidationError as e:
raise ValueError(f"Invalid arguments for function {self.name}") from e
@overload
def beta_tool(func: FunctionT) -> BetaFunctionTool[FunctionT]: ...
@overload
def beta_tool(
func: FunctionT,
*,
name: str | None = None,
description: str | None = None,
input_schema: InputSchema | type[BaseModel] | None = None,
) -> BetaFunctionTool[FunctionT]: ...
@overload
def beta_tool(
*,
name: str | None = None,
description: str | None = None,
input_schema: InputSchema | type[BaseModel] | None = None,
defer_loading: bool | None = None,
) -> Callable[[FunctionT], BetaFunctionTool[FunctionT]]: ...
def beta_tool(
func: FunctionT | None = None,
*,
name: str | None = None,
description: str | None = None,
input_schema: InputSchema | type[BaseModel] | None = None,
defer_loading: bool | None = None,
) -> BetaFunctionTool[FunctionT] | Callable[[FunctionT], BetaFunctionTool[FunctionT]]:
"""Create a FunctionTool from a function with automatic schema inference.
Can be used as a decorator with or without parentheses:
@function_tool
def my_func(x: int) -> str: ...
@function_tool()
def my_func(x: int) -> str: ...
@function_tool(name="custom_name")
def my_func(x: int) -> str: ...
"""
if _compat.PYDANTIC_V1:
raise RuntimeError("Tool functions are only supported with Pydantic v2")
if func is not None:
# @beta_tool called without parentheses
return BetaFunctionTool(
func=func, name=name, description=description, input_schema=input_schema, defer_loading=defer_loading
)
# @beta_tool()
def decorator(func: FunctionT) -> BetaFunctionTool[FunctionT]:
return BetaFunctionTool(
func=func, name=name, description=description, input_schema=input_schema, defer_loading=defer_loading
)
return decorator
@overload
def beta_async_tool(func: AsyncFunctionT) -> BetaAsyncFunctionTool[AsyncFunctionT]: ...
@overload
def beta_async_tool(
func: AsyncFunctionT,
*,
name: str | None = None,
description: str | None = None,
input_schema: InputSchema | type[BaseModel] | None = None,
defer_loading: bool | None = None,
) -> BetaAsyncFunctionTool[AsyncFunctionT]: ...
@overload
def beta_async_tool(
*,
name: str | None = None,
description: str | None = None,
input_schema: InputSchema | type[BaseModel] | None = None,
defer_loading: bool | None = None,
) -> Callable[[AsyncFunctionT], BetaAsyncFunctionTool[AsyncFunctionT]]: ...
def beta_async_tool(
func: AsyncFunctionT | None = None,
*,
name: str | None = None,
description: str | None = None,
input_schema: InputSchema | type[BaseModel] | None = None,
defer_loading: bool | None = None,
) -> BetaAsyncFunctionTool[AsyncFunctionT] | Callable[[AsyncFunctionT], BetaAsyncFunctionTool[AsyncFunctionT]]:
"""Create an AsyncFunctionTool from a function with automatic schema inference.
Can be used as a decorator with or without parentheses:
@async_tool
async def my_func(x: int) -> str: ...
@async_tool()
async def my_func(x: int) -> str: ...
@async_tool(name="custom_name")
async def my_func(x: int) -> str: ...
"""
if _compat.PYDANTIC_V1:
raise RuntimeError("Tool functions are only supported with Pydantic v2")
if func is not None:
# @beta_async_tool called without parentheses
return BetaAsyncFunctionTool(
func=func,
name=name,
description=description,
input_schema=input_schema,
defer_loading=defer_loading,
)
# @beta_async_tool()
def decorator(func: AsyncFunctionT) -> BetaAsyncFunctionTool[AsyncFunctionT]:
return BetaAsyncFunctionTool(
func=func,
name=name,
description=description,
input_schema=input_schema,
defer_loading=defer_loading,
)
return decorator
BetaRunnableTool = Union[BetaFunctionTool[Any], BetaBuiltinFunctionTool]
BetaAsyncRunnableTool = Union[BetaAsyncFunctionTool[Any], BetaAsyncBuiltinFunctionTool]
|