types: allow items and defs for tools (#501)
Some checks failed
test / test (3.10) (push) Has been cancelled
test / test (3.11) (push) Has been cancelled
test / test (3.12) (push) Has been cancelled
test / test (3.13) (push) Has been cancelled
test / test (3.8) (push) Has been cancelled
test / test (3.9) (push) Has been cancelled
test / lint (push) Has been cancelled

This commit is contained in:
Parth Sareen 2025-04-16 13:33:51 -07:00 committed by GitHub
parent e33ceb7188
commit 65f94b4fba
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 19 additions and 4 deletions

View File

@ -307,15 +307,19 @@ class Tool(SubscriptableBaseModel):
description: Optional[str] = None
class Parameters(SubscriptableBaseModel):
model_config = ConfigDict(populate_by_name=True)
type: Optional[Literal['object']] = 'object'
defs: Optional[Any] = Field(None, alias='$defs')
items: Optional[Any] = None
required: Optional[Sequence[str]] = None
class Property(SubscriptableBaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
type: Optional[Union[str, Sequence[str]]] = None
items: Optional[Any] = None
description: Optional[str] = None
enum: Optional[Sequence] = None
enum: Optional[Sequence[Any]] = None
properties: Optional[Mapping[str, Property]] = None
@ -325,6 +329,15 @@ class Tool(SubscriptableBaseModel):
class ChatRequest(BaseGenerateRequest):
@model_serializer(mode='wrap')
def serialize_model(self, nxt):
output = nxt(self)
if 'tools' in output and output['tools']:
for tool in output['tools']:
if 'function' in tool and 'parameters' in tool['function'] and 'defs' in tool['function']['parameters']:
tool['function']['parameters']['$defs'] = tool['function']['parameters'].pop('defs')
return output
messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None
'Messages to chat with.'

View File

@ -187,7 +187,7 @@ def test_function_with_only_description():
tool = convert_function_to_tool(only_description).model_dump()
assert tool['function']['description'] == 'A function with only a description.'
assert tool['function']['parameters'] == {'type': 'object', 'properties': {}, 'required': None}
assert tool['function']['parameters'] == {'type': 'object', 'defs': None, 'items': None, 'required': None, 'properties': {}}
def only_description_with_args(x: int, y: int):
"""
@ -199,9 +199,11 @@ def test_function_with_only_description():
assert tool['function']['description'] == 'A function with only a description.'
assert tool['function']['parameters'] == {
'type': 'object',
'defs': None,
'items': None,
'properties': {
'x': {'type': 'integer', 'description': '', 'enum': None},
'y': {'type': 'integer', 'description': '', 'enum': None},
'x': {'type': 'integer', 'description': '', 'enum': None, 'items': None},
'y': {'type': 'integer', 'description': '', 'enum': None, 'items': None},
},
'required': ['x', 'y'],
}