From b11c36c7eca39ebad98229682854f084cff73b6c Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Fri, 21 Mar 2025 14:10:43 +0000 Subject: [PATCH 1/5] Add return docstring to the function description --- main.py | 26 +++++++++++++++++++++++++ pydantic_ai_slim/pydantic_ai/_griffe.py | 9 +++++++++ 2 files changed, 35 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000000..cc7e986fa4 --- /dev/null +++ b/main.py @@ -0,0 +1,26 @@ +from pydantic_ai import Agent +from pydantic_ai.models.openai import OpenAIModel +from pydantic_ai.providers.openai import OpenAIProvider + +provider = OpenAIProvider(base_url='http://127.0.0.1:11434/v1') +model = OpenAIModel('llama3.2', provider=provider) + +agent = Agent(model=model) + + +@agent.tool_plain +async def my_tool(a: int, b: int) -> int: + """Sum two numbers. + + Args: + a: First number. + b: Second number: + + Returns: + int: The sum between a and b. + """ + return a + b + + +result = agent.run_sync('Sum 10 + 15!') +print(result.data) diff --git a/pydantic_ai_slim/pydantic_ai/_griffe.py b/pydantic_ai_slim/pydantic_ai/_griffe.py index 29866f2934..def9329774 100644 --- a/pydantic_ai_slim/pydantic_ai/_griffe.py +++ b/pydantic_ai_slim/pydantic_ai/_griffe.py @@ -25,6 +25,7 @@ def doc_descriptions( Returns: A tuple of (main function description, parameter descriptions). """ + # TODO(Marcelo): Modify the docstring above. doc = func.__doc__ if doc is None: return '', {} @@ -45,6 +46,14 @@ def doc_descriptions( if main := next((p for p in sections if p.kind == DocstringSectionKind.text), None): main_desc = main.value + # TODO(Marcelo): How should we append the return docstring part to the description? + if return_ := next((p for p in sections if p.kind == DocstringSectionKind.returns), None): + return_statement = return_.value[0] + return_desc = return_statement.description + if return_type := return_statement.name: + return_desc = f'Returns: {return_desc} ({return_type})' + main_desc = '\n'.join((main_desc, return_desc)) + return main_desc, params From 6c19585ae1f071cf04ddbf7ac8aca408ce879c8e Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Sat, 22 Mar 2025 09:01:20 +0100 Subject: [PATCH 2/5] Is this fine? --- pydantic_ai_slim/pydantic_ai/_griffe.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_griffe.py b/pydantic_ai_slim/pydantic_ai/_griffe.py index def9329774..92d7426bf4 100644 --- a/pydantic_ai_slim/pydantic_ai/_griffe.py +++ b/pydantic_ai_slim/pydantic_ai/_griffe.py @@ -49,9 +49,9 @@ def doc_descriptions( # TODO(Marcelo): How should we append the return docstring part to the description? if return_ := next((p for p in sections if p.kind == DocstringSectionKind.returns), None): return_statement = return_.value[0] - return_desc = return_statement.description + return_desc = f'Returns {return_statement.description}' if return_type := return_statement.name: - return_desc = f'Returns: {return_desc} ({return_type})' + return_desc = f'{return_desc} ({return_type})' main_desc = '\n'.join((main_desc, return_desc)) return main_desc, params From e3145e3191b4a6fc379333be97485438848e4aab Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Sat, 22 Mar 2025 10:12:39 +0100 Subject: [PATCH 3/5] Add tests --- main.py | 26 ------ pydantic_ai_slim/pydantic_ai/_griffe.py | 30 ++++-- tests/test_tools.py | 116 ++++++++++++++++++++++++ 3 files changed, 139 insertions(+), 33 deletions(-) delete mode 100644 main.py diff --git a/main.py b/main.py deleted file mode 100644 index cc7e986fa4..0000000000 --- a/main.py +++ /dev/null @@ -1,26 +0,0 @@ -from pydantic_ai import Agent -from pydantic_ai.models.openai import OpenAIModel -from pydantic_ai.providers.openai import OpenAIProvider - -provider = OpenAIProvider(base_url='http://127.0.0.1:11434/v1') -model = OpenAIModel('llama3.2', provider=provider) - -agent = Agent(model=model) - - -@agent.tool_plain -async def my_tool(a: int, b: int) -> int: - """Sum two numbers. - - Args: - a: First number. - b: Second number: - - Returns: - int: The sum between a and b. - """ - return a + b - - -result = agent.run_sync('Sum 10 + 15!') -print(result.data) diff --git a/pydantic_ai_slim/pydantic_ai/_griffe.py b/pydantic_ai_slim/pydantic_ai/_griffe.py index 92d7426bf4..d834a3a3c8 100644 --- a/pydantic_ai_slim/pydantic_ai/_griffe.py +++ b/pydantic_ai_slim/pydantic_ai/_griffe.py @@ -22,10 +22,17 @@ def doc_descriptions( ) -> tuple[str, dict[str, str]]: """Extract the function description and parameter descriptions from a function's docstring. + The function parses the docstring using the specified format (or infers it if 'auto') + and extracts both the main description and parameter descriptions. If a returns section + is present in the docstring, the main description will be formatted as XML. + Returns: - A tuple of (main function description, parameter descriptions). + A tuple containing: + - str: Main description string, which may be either: + * Plain text if no returns section is present + * XML-formatted if returns section exists, including and tags + - dict[str, str]: Dictionary mapping parameter names to their descriptions """ - # TODO(Marcelo): Modify the docstring above. doc = func.__doc__ if doc is None: return '', {} @@ -46,13 +53,22 @@ def doc_descriptions( if main := next((p for p in sections if p.kind == DocstringSectionKind.text), None): main_desc = main.value - # TODO(Marcelo): How should we append the return docstring part to the description? if return_ := next((p for p in sections if p.kind == DocstringSectionKind.returns), None): return_statement = return_.value[0] - return_desc = f'Returns {return_statement.description}' - if return_type := return_statement.name: - return_desc = f'{return_desc} ({return_type})' - main_desc = '\n'.join((main_desc, return_desc)) + return_desc = return_statement.description + + if docstring_style == 'google': + return_type = return_statement.name + type_tag = f'{return_type}\n' if return_type else '' + else: + return_type = return_statement.annotation + type_tag = f'{return_type}\n' if return_type else '' + return_xml = f'\n{type_tag}{return_desc}\n' + + if main_desc: + main_desc = f'{main_desc}\n{return_xml}' + else: + main_desc = return_xml return main_desc, params diff --git a/tests/test_tools.py b/tests/test_tools.py index ebe42557c2..6d9883830d 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -187,6 +187,122 @@ def test_docstring_numpy(docstring_format: Literal['numpy', 'auto']): ) +def test_google_style_with_returns(): + agent = Agent(FunctionModel(get_json_schema)) + + def my_tool(x: int) -> str: + """A function that does something. + + Args: + x: The input value. + + Returns: + str: The result as a string. + """ + return str(x) + + agent.tool_plain(my_tool) + result = agent.run_sync('Hello') + json_schema = json.loads(result.data) + assert json_schema == snapshot( + { + 'name': 'my_tool', + 'description': """\ +A function that does something. + +str +The result as a string. +\ +""", + 'parameters_json_schema': { + 'additionalProperties': False, + 'properties': {'x': {'description': 'The input value.', 'type': 'integer'}}, + 'required': ['x'], + 'type': 'object', + }, + 'outer_typed_dict_key': None, + } + ) + + +def test_sphinx_style_with_returns(): + agent = Agent(FunctionModel(get_json_schema)) + + def my_tool(x: int) -> str: + """A sphinx function with returns. + + :param x: The input value. + :rtype: str + :return: The result as a string with type. + """ + return str(x) + + agent.tool_plain(docstring_format='sphinx')(my_tool) + result = agent.run_sync('Hello') + json_schema = json.loads(result.data) + assert json_schema == snapshot( + { + 'name': 'my_tool', + 'description': """\ +A sphinx function with returns. + +str +The result as a string with type. +\ +""", + 'parameters_json_schema': { + 'additionalProperties': False, + 'properties': {'x': {'description': 'The input value.', 'type': 'integer'}}, + 'required': ['x'], + 'type': 'object', + }, + 'outer_typed_dict_key': None, + } + ) + + +def test_numpy_style_with_returns(): + agent = Agent(FunctionModel(get_json_schema)) + + def my_tool(x: int) -> str: + """A numpy function with returns. + + Parameters + ---------- + x : int + The input value. + + Returns + ------- + str + The result as a string with type. + """ + return str(x) + + agent.tool_plain(docstring_format='numpy')(my_tool) + result = agent.run_sync('Hello') + json_schema = json.loads(result.data) + assert json_schema == snapshot( + { + 'name': 'my_tool', + 'description': """\ +A numpy function with returns. + +str +The result as a string with type. +\ +""", + 'parameters_json_schema': { + 'additionalProperties': False, + 'properties': {'x': {'description': 'The input value.', 'type': 'integer'}}, + 'required': ['x'], + 'type': 'object', + }, + 'outer_typed_dict_key': None, + } + ) + + def unknown_docstring(**kwargs: int) -> str: # pragma: no cover """Unknown style docstring.""" return str(kwargs) From 3be97bae548d308793661ef77a71db8e03378692 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Sat, 22 Mar 2025 10:17:22 +0100 Subject: [PATCH 4/5] Add tests --- tests/test_tools.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/tests/test_tools.py b/tests/test_tools.py index 6d9883830d..c975434d51 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -122,7 +122,6 @@ def sphinx_style_docstring(foo: int, /) -> str: # pragma: no cover """Sphinx style docstring. :param foo: The foo thing. - :return: The result. """ return str(foo) @@ -688,11 +687,7 @@ def ctx_tool(ctx: RunContext[int], x: int) -> int: async def tool_without_return_annotation_in_docstring() -> str: # pragma: no cover - """A tool that documents what it returns but doesn't have a return annotation in the docstring. - - Returns: - A value. - """ + """A tool that documents what it returns but doesn't have a return annotation in the docstring.""" return '' @@ -707,8 +702,7 @@ def test_suppress_griffe_logging(caplog: LogCaptureFixture): json_schema = json.loads(result.data) assert json_schema == snapshot( { - 'description': "A tool that documents what it returns but doesn't have a " - 'return annotation in the docstring.', + 'description': "A tool that documents what it returns but doesn't have a return annotation in the docstring.", 'name': 'tool_without_return_annotation_in_docstring', 'outer_typed_dict_key': None, 'parameters_json_schema': {'additionalProperties': False, 'properties': {}, 'type': 'object'}, From ea28a45a1606e502f9f8d34ab09d85a533ee10fb Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Sat, 22 Mar 2025 10:53:55 +0100 Subject: [PATCH 5/5] Add more tests --- pydantic_ai_slim/pydantic_ai/_griffe.py | 18 +++++++------ tests/test_tools.py | 36 ++++++++++++++++++++++--- 2 files changed, 43 insertions(+), 11 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_griffe.py b/pydantic_ai_slim/pydantic_ai/_griffe.py index d834a3a3c8..44d9092378 100644 --- a/pydantic_ai_slim/pydantic_ai/_griffe.py +++ b/pydantic_ai_slim/pydantic_ai/_griffe.py @@ -41,7 +41,14 @@ def doc_descriptions( parent = cast(GriffeObject, sig) docstring_style = _infer_docstring_style(doc) if docstring_format == 'auto' else docstring_format - docstring = Docstring(doc, lineno=1, parser=docstring_style, parent=parent) + docstring = Docstring( + doc, + lineno=1, + parser=docstring_style, + parent=parent, + # https://mkdocstrings.github.io/griffe/reference/docstrings/#google-options + parser_options={'returns_named_value': False, 'returns_multiple_items': False}, + ) with _disable_griffe_logging(): sections = docstring.parse() @@ -56,13 +63,8 @@ def doc_descriptions( if return_ := next((p for p in sections if p.kind == DocstringSectionKind.returns), None): return_statement = return_.value[0] return_desc = return_statement.description - - if docstring_style == 'google': - return_type = return_statement.name - type_tag = f'{return_type}\n' if return_type else '' - else: - return_type = return_statement.annotation - type_tag = f'{return_type}\n' if return_type else '' + return_type = return_statement.annotation + type_tag = f'{return_type}\n' if return_type else '' return_xml = f'\n{type_tag}{return_desc}\n' if main_desc: diff --git a/tests/test_tools.py b/tests/test_tools.py index c975434d51..cf7fbb99e4 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -189,7 +189,7 @@ def test_docstring_numpy(docstring_format: Literal['numpy', 'auto']): def test_google_style_with_returns(): agent = Agent(FunctionModel(get_json_schema)) - def my_tool(x: int) -> str: + def my_tool(x: int) -> str: # pragma: no cover """A function that does something. Args: @@ -227,7 +227,7 @@ def my_tool(x: int) -> str: def test_sphinx_style_with_returns(): agent = Agent(FunctionModel(get_json_schema)) - def my_tool(x: int) -> str: + def my_tool(x: int) -> str: # pragma: no cover """A sphinx function with returns. :param x: The input value. @@ -263,7 +263,7 @@ def my_tool(x: int) -> str: def test_numpy_style_with_returns(): agent = Agent(FunctionModel(get_json_schema)) - def my_tool(x: int) -> str: + def my_tool(x: int) -> str: # pragma: no cover """A numpy function with returns. Parameters @@ -302,6 +302,36 @@ def my_tool(x: int) -> str: ) +def only_returns_type() -> str: # pragma: no cover + """ + + Returns: + str: The result as a string. + """ + return 'foo' + + +def test_only_returns_type(): + agent = Agent(FunctionModel(get_json_schema)) + agent.tool_plain(only_returns_type) + + result = agent.run_sync('Hello') + json_schema = json.loads(result.data) + assert json_schema == snapshot( + { + 'name': 'only_returns_type', + 'description': """\ + +str +The result as a string. +\ +""", + 'parameters_json_schema': {'additionalProperties': False, 'properties': {}, 'type': 'object'}, + 'outer_typed_dict_key': None, + } + ) + + def unknown_docstring(**kwargs: int) -> str: # pragma: no cover """Unknown style docstring.""" return str(kwargs)