Skip to content

Move description onto tool when appropriate #228

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Dec 12, 2024
Merged

Move description onto tool when appropriate #228

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 1 addition & 3 deletions docs/agents.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -479,7 +479,6 @@ def print_schema(messages: list[Message], info: AgentInfo) -> ModelAnyResponse:
print(tool.parameters_json_schema)
"""
{
'description': 'Get me foobar.',
'properties': {
'a': {'description': 'apple pie', 'title': 'A', 'type': 'integer'},
'b': {'description': 'banana cake', 'title': 'B', 'type': 'string'},
Expand Down Expand Up @@ -540,9 +539,8 @@ print(test_model.agent_model_function_tools)
[
ToolDefinition(
name='foobar',
description='',
description='This is a Foobar',
parameters_json_schema={
'description': 'This is a Foobar',
'properties': {
'x': {'title': 'X', 'type': 'integer'},
'y': {'title': 'Y', 'type': 'string'},
Expand Down
10 changes: 5 additions & 5 deletions pydantic_ai_slim/pydantic_ai/_pydantic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,14 +138,14 @@ def function_schema(function: Callable[..., Any], takes_ctx: bool) -> FunctionSc
json_schema = GenerateJsonSchema().generate(schema)

# workaround for https://github.com/pydantic/pydantic/issues/10785
# if we build a custom TypeDict schema (matches when `single_arg_name` is None), we manually set
# if we build a custom TypeDict schema (matches when `single_arg_name is None`), we manually set
# `additionalProperties` in the JSON Schema
if single_arg_name is None:
json_schema['additionalProperties'] = bool(var_kwargs_schema)

# instead of passing `description` through in core_schema, we just add it here
if description:
json_schema = {'description': description} | json_schema
elif not description:
# if the tool description is not set, and we have a single parameter, take the description from that
# and set it on the tool
description = json_schema.pop('description', None)

return FunctionSchema(
description=description,
Expand Down
128 changes: 79 additions & 49 deletions tests/test_tools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
from dataclasses import dataclass
from typing import Annotated, Any, Callable, Union

import pydantic_core
import pytest
from inline_snapshot import snapshot
from pydantic import BaseModel, Field
Expand Down Expand Up @@ -73,7 +74,7 @@ async def google_style_docstring(foo: int, bar: str) -> str: # pragma: no cover
async def get_json_schema(_messages: list[Message], info: AgentInfo) -> ModelAnyResponse:
assert len(info.function_tools) == 1
r = info.function_tools[0]
return ModelTextResponse(json.dumps(r.parameters_json_schema))
return ModelTextResponse(pydantic_core.to_json(r).decode())


def test_docstring_google(set_event_loop: None):
Expand All @@ -84,18 +85,25 @@ def test_docstring_google(set_event_loop: None):
json_schema = json.loads(result.data)
assert json_schema == snapshot(
{
'name': 'google_style_docstring',
'description': 'Do foobar stuff, a lot.',
'additionalProperties': False,
'properties': {
'foo': {'description': 'The foo thing.', 'title': 'Foo', 'type': 'integer'},
'bar': {'description': 'The bar thing.', 'title': 'Bar', 'type': 'string'},
'parameters_json_schema': {
'properties': {
'foo': {'description': 'The foo thing.', 'title': 'Foo', 'type': 'integer'},
'bar': {'description': 'The bar thing.', 'title': 'Bar', 'type': 'string'},
},
'required': ['foo', 'bar'],
'type': 'object',
'additionalProperties': False,
},
'required': ['foo', 'bar'],
'type': 'object',
'outer_typed_dict_key': None,
}
)
# description should be the first key
assert next(iter(json_schema)) == 'description'
keys = list(json_schema.keys())
# name should be the first key
assert keys[0] == 'name'
# description should be the second key
assert keys[1] == 'description'


def sphinx_style_docstring(foo: int, /) -> str: # pragma: no cover
Expand All @@ -115,13 +123,15 @@ def test_docstring_sphinx(set_event_loop: None):
json_schema = json.loads(result.data)
assert json_schema == snapshot(
{
'name': 'sphinx_style_docstring',
'description': 'Sphinx style docstring.',
'additionalProperties': False,
'properties': {
'foo': {'description': 'The foo thing.', 'title': 'Foo', 'type': 'integer'},
'parameters_json_schema': {
'properties': {'foo': {'description': 'The foo thing.', 'title': 'Foo', 'type': 'integer'}},
'required': ['foo'],
'type': 'object',
'additionalProperties': False,
},
'required': ['foo'],
'type': 'object',
'outer_typed_dict_key': None,
}
)

Expand All @@ -147,14 +157,18 @@ def test_docstring_numpy(set_event_loop: None):
json_schema = json.loads(result.data)
assert json_schema == snapshot(
{
'name': 'numpy_style_docstring',
'description': 'Numpy style docstring.',
'additionalProperties': False,
'properties': {
'foo': {'description': 'The foo thing.', 'title': 'Foo', 'type': 'integer'},
'bar': {'description': 'The bar thing.', 'title': 'Bar', 'type': 'string'},
'parameters_json_schema': {
'properties': {
'foo': {'description': 'The foo thing.', 'title': 'Foo', 'type': 'integer'},
'bar': {'description': 'The bar thing.', 'title': 'Bar', 'type': 'string'},
},
'required': ['foo', 'bar'],
'type': 'object',
'additionalProperties': False,
},
'required': ['foo', 'bar'],
'type': 'object',
'outer_typed_dict_key': None,
}
)

Expand All @@ -172,10 +186,10 @@ def test_docstring_unknown(set_event_loop: None):
json_schema = json.loads(result.data)
assert json_schema == snapshot(
{
'name': 'unknown_docstring',
'description': 'Unknown style docstring.',
'additionalProperties': True,
'properties': {},
'type': 'object',
'parameters_json_schema': {'properties': {}, 'type': 'object', 'additionalProperties': True},
'outer_typed_dict_key': None,
}
)

Expand All @@ -201,13 +215,16 @@ def test_docstring_google_no_body(set_event_loop: None):
json_schema = json.loads(result.data)
assert json_schema == snapshot(
{
'additionalProperties': False,
'name': 'google_style_docstring_no_body', 'description': '', 'parameters_json_schema': {
'properties': {
'foo': {'description': 'The foo thing.', 'title': 'Foo', 'type': 'integer'},
'bar': {'description': 'from fields', 'title': 'Bar', 'type': 'string'},
'foo': {'description': 'The foo thing.', 'title': 'Foo', 'type': 'integer'},
'bar': {'description': 'from fields', 'title': 'Bar', 'type': 'string'},
},
'required': ['foo', 'bar'],
'type': 'object',
'additionalProperties': False,
},
'required': ['foo', 'bar'],
'type': 'object',
'outer_typed_dict_key': None
}
)

Expand All @@ -228,10 +245,18 @@ def takes_just_model(model: Foo) -> str:
json_schema = json.loads(result.data)
assert json_schema == snapshot(
{
'title': 'Foo',
'properties': {'x': {'title': 'X', 'type': 'integer'}, 'y': {'title': 'Y', 'type': 'string'}},
'required': ['x', 'y'],
'type': 'object',
'name': 'takes_just_model',
'description': None,
'parameters_json_schema': {
'properties': {
'x': {'title': 'X', 'type': 'integer'},
'y': {'title': 'Y', 'type': 'string'},
},
'required': ['x', 'y'],
'title': 'Foo',
'type': 'object',
},
'outer_typed_dict_key': None
}
)

Expand All @@ -250,24 +275,29 @@ def takes_just_model(model: Foo, z: int) -> str:
json_schema = json.loads(result.data)
assert json_schema == snapshot(
{
'$defs': {
'Foo': {
'properties': {
'x': {'title': 'X', 'type': 'integer'},
'y': {'title': 'Y', 'type': 'string'},
},
'required': ['x', 'y'],
'title': 'Foo',
'type': 'object',
}
},
'additionalProperties': False,
'properties': {
'model': {'$ref': '#/$defs/Foo'},
'z': {'title': 'Z', 'type': 'integer'},
'name': 'takes_just_model',
'description': '',
'parameters_json_schema': {
'$defs': {
'Foo': {
'properties': {
'x': {'title': 'X', 'type': 'integer'},
'y': {'title': 'Y', 'type': 'string'},
},
'required': ['x', 'y'],
'title': 'Foo',
'type': 'object',
}
},
'properties': {
'model': {'$ref': '#/$defs/Foo'},
'z': {'title': 'Z', 'type': 'integer'},
},
'required': ['model', 'z'],
'type': 'object',
'additionalProperties': False,
},
'required': ['model', 'z'],
'type': 'object',
'outer_typed_dict_key': None
}
)

Expand Down
42 changes: 21 additions & 21 deletions uv.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading