Spaces:
Paused
Paused
File size: 5,197 Bytes
9792ea7 | 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 | # -*- coding: utf-8 -*-
"""The tool module utils."""
import inspect
from typing import Any, Dict, Callable
from docstring_parser import parse
from pydantic import Field, create_model, ConfigDict
def _remove_title_field(schema: dict) -> dict:
"""Remove the title field from the JSON schema to avoid
misleading the LLM."""
# The top level title field
if "title" in schema:
schema.pop("title")
# properties
if "properties" in schema:
for prop in schema["properties"].values():
if isinstance(prop, dict):
_remove_title_field(prop)
# items
if "items" in schema and isinstance(schema["items"], dict):
_remove_title_field(schema["items"])
# additionalProperties
if "additionalProperties" in schema and isinstance(
schema["additionalProperties"],
dict,
):
_remove_title_field(schema["additionalProperties"])
# $defs — referenced sub-schemas, e.g. Pydantic models used as parameter
# types generate "$defs": {"SubModel": {"title": "SubModel", ...}}.
# These titles are auto-generated noise just like property titles, and
# should be removed for the same reason.
if "$defs" in schema and isinstance(schema["$defs"], dict):
for def_schema in schema["$defs"].values():
if isinstance(def_schema, dict):
_remove_title_field(def_schema)
return schema
def _extract_func_description(docstring: str) -> str:
"""Extract the function description from the docstring.
Args:
docstring (`str`):
The docstring to extract the function description from.
Returns:
`str`:
The extracted function description.
"""
parsed_docstring = parse(docstring or "")
descriptions = []
if parsed_docstring.short_description is not None:
descriptions.append(parsed_docstring.short_description)
if parsed_docstring.long_description is not None:
descriptions.append(parsed_docstring.long_description)
return "\n".join(descriptions)
def _extract_input_schema(
tool_func: Callable,
include_var_positional: bool = False,
include_var_keyword: bool = False,
) -> dict:
"""Extract input schema from the tool function's docstring
Args:
tool_func (`ToolFunction`):
The tool function to extract the JSON schema from.
include_var_positional (`bool`):
Whether to include variable positional arguments in the JSON
schema.
include_var_keyword (`bool`):
Whether to include variable keyword arguments in the JSON schema.
Returns:
`dict`:
The extracted input JSON schema.
"""
docstring = parse(tool_func.__doc__ or "")
params_docstring = {_.arg_name: _.description for _ in docstring.params}
# Create a dynamic model with the function signature
fields = {}
for name, param in inspect.signature(tool_func).parameters.items():
# Skip the `self` and `cls` parameters
if name in ["self", "cls"]:
continue
# Handle `**kwargs`
if param.kind == inspect.Parameter.VAR_KEYWORD:
if not include_var_keyword:
continue
fields[name] = (
Dict[str, Any]
if param.annotation == inspect.Parameter.empty
else Dict[str, param.annotation], # type: ignore
Field(
description=params_docstring.get(
f"**{name}",
params_docstring.get(name, None),
),
default={}
if param.default is param.empty
else param.default,
),
)
elif param.kind == inspect.Parameter.VAR_POSITIONAL:
if not include_var_positional:
continue
fields[name] = (
list[Any]
if param.annotation == inspect.Parameter.empty
else list[param.annotation], # type: ignore
Field(
description=params_docstring.get(
f"*{name}",
params_docstring.get(name, None),
),
default=[]
if param.default is param.empty
else param.default,
),
)
else:
fields[name] = (
Any
if param.annotation == inspect.Parameter.empty
else param.annotation,
Field(
description=params_docstring.get(name, None),
default=...
if param.default is param.empty
else param.default,
),
)
base_model = create_model(
"_StructuredOutputDynamicClass",
__config__=ConfigDict(arbitrary_types_allowed=True),
**fields,
)
params_json_schema = base_model.model_json_schema()
# Remove the title from the json schema
_remove_title_field(params_json_schema)
return params_json_schema
|