From 4bd79c9343cdee4d687891ede57a159932e7bcc5 Mon Sep 17 00:00:00 2001 From: Anirban Basu Date: Wed, 5 Nov 2025 07:21:48 +0900 Subject: [PATCH 1/7] feat: Added a mechanism to extract metadata from MCP tool call response. feat: Added a new MCP tool that attaches metadata to MCP TextContent. test: Added a test to call the aforementioned tool (failing as of now). --- pydantic_ai_slim/pydantic_ai/mcp.py | 21 ++++++++++++++++++--- tests/mcp_server.py | 28 ++++++++++++++++++++++++++++ tests/test_mcp.py | 22 ++++++++++++++++++++-- 3 files changed, 66 insertions(+), 5 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/mcp.py b/pydantic_ai_slim/pydantic_ai/mcp.py index 70364c0ad4..b37a9c0bed 100644 --- a/pydantic_ai_slim/pydantic_ai/mcp.py +++ b/pydantic_ai_slim/pydantic_ai/mcp.py @@ -255,10 +255,24 @@ async def direct_call_tool( # The MCP SDK wraps primitives and generic types like list in a `result` key, but we want to use the raw value returned by the tool function. # See https://github.com/modelcontextprotocol/python-sdk#structured-output if isinstance(structured, dict) and len(structured) == 1 and 'result' in structured: - return structured['result'] - return structured + return ( + messages.ToolReturn(return_value=structured['result'], metadata=result.meta) + if getattr(result, '_meta', None) + else structured['result'] + ) + return ( + messages.ToolReturn(return_value=structured, metadata=result.meta) + if getattr(result, '_meta', None) + else structured + ) mapped = [await self._map_tool_result_part(part) for part in result.content] + if getattr(result, '_meta', None): + return ( + messages.ToolReturn(return_value=mapped[0], metadata=result.meta) + if len(mapped) == 1 + else [messages.ToolReturn(return_value=mapped_item, metadata=result.meta) for mapped_item in mapped] + ) return mapped[0] if len(mapped) == 1 else mapped async def call_tool( @@ -875,9 +889,10 @@ def __eq__(self, value: object, /) -> bool: ToolResult = ( str | messages.BinaryContent + | messages.ToolReturn | dict[str, Any] | list[Any] - | Sequence[str | messages.BinaryContent | dict[str, Any] | list[Any]] + | Sequence[str | messages.BinaryContent | messages.ToolReturn | dict[str, Any] | list[Any]] ) """The result type of an MCP tool call.""" diff --git a/tests/mcp_server.py b/tests/mcp_server.py index 54b105ab29..6aa58ee165 100644 --- a/tests/mcp_server.py +++ b/tests/mcp_server.py @@ -68,6 +68,34 @@ async def get_image_resource_link() -> ResourceLink: ) +@mcp.tool(annotations=ToolAnnotations(title='Collatz Conjecture sequence generator')) +async def collatz_conjecture(n: int) -> TextContent: + """Generate the Collatz conjecture sequence for a given number. + This tool attaches response metadata. + + Args: + n: The starting number for the Collatz sequence. + Returns: + A list representing the Collatz sequence. + """ + if n <= 0: + raise ValueError('Startig number for the Collatz conjecture must be a positive integer.') + + sequence = [n] + while n != 1: + if n % 2 == 0: + n = n // 2 + else: + n = 3 * n + 1 + sequence.append(n) + response = TextContent(type='text', text=str(sequence)) + # attach metadata to the response + if response.meta is None: + response.meta = {} + response.meta['pydantic_ai'] = {'tool': 'collatz_conjecture', 'length': len(sequence)} + return response + + @mcp.resource('resource://kiwi.png', mime_type='image/png') async def kiwi_resource() -> bytes: return Path(__file__).parent.joinpath('assets/kiwi.png').read_bytes() diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 9342f49f10..b18c67125d 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -26,6 +26,7 @@ from pydantic_ai.agent import Agent from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior, UserError from pydantic_ai.mcp import MCPServerStreamableHTTP, load_mcp_servers +from pydantic_ai.messages import ToolReturn from pydantic_ai.models import Model from pydantic_ai.models.test import TestModel from pydantic_ai.tools import RunContext @@ -77,7 +78,7 @@ async def test_stdio_server(run_context: RunContext[int]): server = MCPServerStdio('python', ['-m', 'tests.mcp_server']) async with server: tools = [tool.tool_def for tool in (await server.get_tools(run_context)).values()] - assert len(tools) == snapshot(18) + assert len(tools) == snapshot(19) assert tools[0].name == 'celsius_to_fahrenheit' assert isinstance(tools[0].description, str) assert tools[0].description.startswith('Convert Celsius to Fahrenheit.') @@ -87,6 +88,23 @@ async def test_stdio_server(run_context: RunContext[int]): assert result == snapshot(32.0) +async def test_tool_response_metadata(run_context: RunContext[int]): + server = MCPServerStdio('python', ['-m', 'tests.mcp_server']) + async with server: + tools = [tool.tool_def for tool in (await server.get_tools(run_context)).values()] + assert len(tools) == snapshot(19) + assert tools[4].name == 'collatz_conjecture' + assert isinstance(tools[4].description, str) + assert tools[4].description.startswith('Generate the Collatz conjecture sequence for a given number.') + + # Test calling the Collatz conjecture generator tool + result = await server.direct_call_tool('collatz_conjecture', {'n': 7}) + assert isinstance(result, ToolReturn) + assert isinstance(result.return_value, str) + assert result.return_value == '[7, 22, 11, 34, 17, 52, 26, 13, 40, 20, 10, 5, 16, 8, 4, 2, 1]' + assert result.metadata == {'pydantic_ai': {'tool': 'collatz_conjecture', 'length': 17}} + + async def test_reentrant_context_manager(): server = MCPServerStdio('python', ['-m', 'tests.mcp_server']) async with server: @@ -138,7 +156,7 @@ async def test_stdio_server_with_cwd(run_context: RunContext[int]): server = MCPServerStdio('python', ['mcp_server.py'], cwd=test_dir) async with server: tools = await server.get_tools(run_context) - assert len(tools) == snapshot(18) + assert len(tools) == snapshot(19) async def test_process_tool_call(run_context: RunContext[int]) -> int: From 5c3f58f9e60ab2dbbdaaada05b2aaad0c998c119 Mon Sep 17 00:00:00 2001 From: Anirban Basu Date: Wed, 5 Nov 2025 08:47:32 +0900 Subject: [PATCH 2/7] fix: Attempted fix to check the _meta tag presence in the MCP response. --- pydantic_ai_slim/pydantic_ai/mcp.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/mcp.py b/pydantic_ai_slim/pydantic_ai/mcp.py index b37a9c0bed..7995adcd40 100644 --- a/pydantic_ai_slim/pydantic_ai/mcp.py +++ b/pydantic_ai_slim/pydantic_ai/mcp.py @@ -257,17 +257,17 @@ async def direct_call_tool( if isinstance(structured, dict) and len(structured) == 1 and 'result' in structured: return ( messages.ToolReturn(return_value=structured['result'], metadata=result.meta) - if getattr(result, '_meta', None) + if getattr(result, '_meta', None) is not None else structured['result'] ) return ( messages.ToolReturn(return_value=structured, metadata=result.meta) - if getattr(result, '_meta', None) + if getattr(result, '_meta', None) is not None else structured ) mapped = [await self._map_tool_result_part(part) for part in result.content] - if getattr(result, '_meta', None): + if getattr(result, '_meta', None) is not None: return ( messages.ToolReturn(return_value=mapped[0], metadata=result.meta) if len(mapped) == 1 From 0e5ae009b85002c5987599877952e974ac99fdee Mon Sep 17 00:00:00 2001 From: Anirban Basu Date: Wed, 5 Nov 2025 12:10:46 +0900 Subject: [PATCH 3/7] fix: Tests seem to work with the MCP metadata but are these exhaustive? --- pydantic_ai_slim/pydantic_ai/mcp.py | 15 +++++++++------ tests/mcp_server.py | 8 +++----- tests/test_mcp.py | 1 - 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/mcp.py b/pydantic_ai_slim/pydantic_ai/mcp.py index 7995adcd40..f539a2145f 100644 --- a/pydantic_ai_slim/pydantic_ai/mcp.py +++ b/pydantic_ai_slim/pydantic_ai/mcp.py @@ -254,20 +254,23 @@ async def direct_call_tool( ): # The MCP SDK wraps primitives and generic types like list in a `result` key, but we want to use the raw value returned by the tool function. # See https://github.com/modelcontextprotocol/python-sdk#structured-output - if isinstance(structured, dict) and len(structured) == 1 and 'result' in structured: + if isinstance(structured, dict) and ( + (len(structured) == 1 and 'result' in structured) + or (len(structured) == 2 and 'result' in structured and '_meta' in structured) + ): return ( - messages.ToolReturn(return_value=structured['result'], metadata=result.meta) - if getattr(result, '_meta', None) is not None + messages.ToolReturn(return_value=structured['result'], metadata=structured['_meta']) + if structured.get('_meta', None) is not None else structured['result'] ) return ( - messages.ToolReturn(return_value=structured, metadata=result.meta) - if getattr(result, '_meta', None) is not None + messages.ToolReturn(return_value=structured, metadata=structured['_meta']) + if structured.get('_meta', None) is not None else structured ) mapped = [await self._map_tool_result_part(part) for part in result.content] - if getattr(result, '_meta', None) is not None: + if result.meta: return ( messages.ToolReturn(return_value=mapped[0], metadata=result.meta) if len(mapped) == 1 diff --git a/tests/mcp_server.py b/tests/mcp_server.py index 6aa58ee165..1dd2af9a98 100644 --- a/tests/mcp_server.py +++ b/tests/mcp_server.py @@ -69,7 +69,7 @@ async def get_image_resource_link() -> ResourceLink: @mcp.tool(annotations=ToolAnnotations(title='Collatz Conjecture sequence generator')) -async def collatz_conjecture(n: int) -> TextContent: +async def collatz_conjecture(n: int) -> dict[str, Any]: """Generate the Collatz conjecture sequence for a given number. This tool attaches response metadata. @@ -88,11 +88,9 @@ async def collatz_conjecture(n: int) -> TextContent: else: n = 3 * n + 1 sequence.append(n) - response = TextContent(type='text', text=str(sequence)) + response: dict[str, Any] = {'result': str(sequence)} # attach metadata to the response - if response.meta is None: - response.meta = {} - response.meta['pydantic_ai'] = {'tool': 'collatz_conjecture', 'length': len(sequence)} + response['_meta'] = {'pydantic_ai': {'tool': 'collatz_conjecture', 'length': len(sequence)}} return response diff --git a/tests/test_mcp.py b/tests/test_mcp.py index b18c67125d..c4da0a3f95 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -100,7 +100,6 @@ async def test_tool_response_metadata(run_context: RunContext[int]): # Test calling the Collatz conjecture generator tool result = await server.direct_call_tool('collatz_conjecture', {'n': 7}) assert isinstance(result, ToolReturn) - assert isinstance(result.return_value, str) assert result.return_value == '[7, 22, 11, 34, 17, 52, 26, 13, 40, 20, 10, 5, 16, 8, 4, 2, 1]' assert result.metadata == {'pydantic_ai': {'tool': 'collatz_conjecture', 'length': 17}} From ad7359129e6ee7f54e2653c7cb0ae7f4b4c4e4ed Mon Sep 17 00:00:00 2001 From: Anirban Basu Date: Wed, 5 Nov 2025 16:57:21 +0900 Subject: [PATCH 4/7] chore: Improved metadata parsing for both structured content and TextContent. --- pydantic_ai_slim/pydantic_ai/mcp.py | 18 ++++++++--------- tests/mcp_server.py | 29 ++++++++++++++++++++-------- tests/test_mcp.py | 30 +++++++++++++++++++++-------- 3 files changed, 52 insertions(+), 25 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/mcp.py b/pydantic_ai_slim/pydantic_ai/mcp.py index f539a2145f..9660c400a2 100644 --- a/pydantic_ai_slim/pydantic_ai/mcp.py +++ b/pydantic_ai_slim/pydantic_ai/mcp.py @@ -270,12 +270,6 @@ async def direct_call_tool( ) mapped = [await self._map_tool_result_part(part) for part in result.content] - if result.meta: - return ( - messages.ToolReturn(return_value=mapped[0], metadata=result.meta) - if len(mapped) == 1 - else [messages.ToolReturn(return_value=mapped_item, metadata=result.meta) for mapped_item in mapped] - ) return mapped[0] if len(mapped) == 1 else mapped async def call_tool( @@ -391,17 +385,23 @@ async def _sampling_callback( async def _map_tool_result_part( self, part: mcp_types.ContentBlock - ) -> str | messages.BinaryContent | dict[str, Any] | list[Any]: + ) -> str | messages.ToolReturn | messages.BinaryContent | dict[str, Any] | list[Any]: # See https://github.com/jlowin/fastmcp/blob/main/docs/servers/tools.mdx#return-values + # Let's also check for metadata but it can be present in not just TextContent + metadata: dict[str, Any] | None = part.meta if isinstance(part, mcp_types.TextContent): text = part.text if text.startswith(('[', '{')): try: - return pydantic_core.from_json(text) + return ( + pydantic_core.from_json(text) + if metadata is None + else messages.ToolReturn(return_value=pydantic_core.from_json(text), metadata=metadata) + ) except ValueError: pass - return text + return text if metadata is None else messages.ToolReturn(return_value=text, metadata=metadata) elif isinstance(part, mcp_types.ImageContent): return messages.BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType) elif isinstance(part, mcp_types.AudioContent): diff --git a/tests/mcp_server.py b/tests/mcp_server.py index 1dd2af9a98..80ba0d806d 100644 --- a/tests/mcp_server.py +++ b/tests/mcp_server.py @@ -68,18 +68,20 @@ async def get_image_resource_link() -> ResourceLink: ) -@mcp.tool(annotations=ToolAnnotations(title='Collatz Conjecture sequence generator')) -async def collatz_conjecture(n: int) -> dict[str, Any]: +@mcp.tool(structured_output=False, annotations=ToolAnnotations(title='Collatz Conjecture sequence generator')) +async def get_collatz_conjecture(n: int) -> TextContent: """Generate the Collatz conjecture sequence for a given number. This tool attaches response metadata. Args: n: The starting number for the Collatz sequence. Returns: - A list representing the Collatz sequence. + A list representing the Collatz sequence with attached metadata. """ if n <= 0: - raise ValueError('Startig number for the Collatz conjecture must be a positive integer.') + raise ValueError('Starting number for the Collatz conjecture must be a positive integer.') + + input_param_n = n # store the original input value sequence = [n] while n != 1: @@ -88,10 +90,21 @@ async def collatz_conjecture(n: int) -> dict[str, Any]: else: n = 3 * n + 1 sequence.append(n) - response: dict[str, Any] = {'result': str(sequence)} - # attach metadata to the response - response['_meta'] = {'pydantic_ai': {'tool': 'collatz_conjecture', 'length': len(sequence)}} - return response + + return TextContent( + type='text', + text=str(sequence), + _meta={'pydantic_ai': {'tool': 'collatz_conjecture', 'n': input_param_n, 'length': len(sequence)}}, + ) + + +@mcp.tool() +async def get_structured_text_content_with_metadata() -> dict[str, Any]: + """Return structured dict with metadata.""" + return { + 'result': 'This is some text content.', + '_meta': {'pydantic_ai': {'source': 'get_structured_text_content_with_metadata'}}, + } @mcp.resource('resource://kiwi.png', mime_type='image/png') diff --git a/tests/test_mcp.py b/tests/test_mcp.py index c4da0a3f95..688b09c27d 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -78,7 +78,7 @@ async def test_stdio_server(run_context: RunContext[int]): server = MCPServerStdio('python', ['-m', 'tests.mcp_server']) async with server: tools = [tool.tool_def for tool in (await server.get_tools(run_context)).values()] - assert len(tools) == snapshot(19) + assert len(tools) == snapshot(20) assert tools[0].name == 'celsius_to_fahrenheit' assert isinstance(tools[0].description, str) assert tools[0].description.startswith('Convert Celsius to Fahrenheit.') @@ -92,16 +92,30 @@ async def test_tool_response_metadata(run_context: RunContext[int]): server = MCPServerStdio('python', ['-m', 'tests.mcp_server']) async with server: tools = [tool.tool_def for tool in (await server.get_tools(run_context)).values()] - assert len(tools) == snapshot(19) - assert tools[4].name == 'collatz_conjecture' + assert len(tools) == snapshot(20) + assert tools[4].name == 'get_collatz_conjecture' assert isinstance(tools[4].description, str) assert tools[4].description.startswith('Generate the Collatz conjecture sequence for a given number.') - # Test calling the Collatz conjecture generator tool - result = await server.direct_call_tool('collatz_conjecture', {'n': 7}) + result = await server.direct_call_tool('get_collatz_conjecture', {'n': 7}) assert isinstance(result, ToolReturn) - assert result.return_value == '[7, 22, 11, 34, 17, 52, 26, 13, 40, 20, 10, 5, 16, 8, 4, 2, 1]' - assert result.metadata == {'pydantic_ai': {'tool': 'collatz_conjecture', 'length': 17}} + assert result.return_value == snapshot([7, 22, 11, 34, 17, 52, 26, 13, 40, 20, 10, 5, 16, 8, 4, 2, 1]) + assert result.metadata == snapshot({'pydantic_ai': {'tool': 'collatz_conjecture', 'n': 7, 'length': 17}}) + + +async def test_tool_structured_response_metadata(run_context: RunContext[int]): + server = MCPServerStdio('python', ['-m', 'tests.mcp_server']) + async with server: + tools = [tool.tool_def for tool in (await server.get_tools(run_context)).values()] + assert len(tools) == snapshot(20) + assert tools[5].name == 'get_structured_text_content_with_metadata' + assert isinstance(tools[5].description, str) + assert tools[5].description.startswith('Return structured dict with metadata.') + + result = await server.direct_call_tool('get_structured_text_content_with_metadata', {}) + assert isinstance(result, ToolReturn) + assert result.return_value == 'This is some text content.' + assert result.metadata == snapshot({'pydantic_ai': {'source': 'get_structured_text_content_with_metadata'}}) async def test_reentrant_context_manager(): @@ -155,7 +169,7 @@ async def test_stdio_server_with_cwd(run_context: RunContext[int]): server = MCPServerStdio('python', ['mcp_server.py'], cwd=test_dir) async with server: tools = await server.get_tools(run_context) - assert len(tools) == snapshot(19) + assert len(tools) == snapshot(20) async def test_process_tool_call(run_context: RunContext[int]) -> int: From dfe81b6b1a7f6d345624da3ffec80c731407ffe1 Mon Sep 17 00:00:00 2001 From: Anirban Basu Date: Wed, 5 Nov 2025 22:14:17 +0900 Subject: [PATCH 5/7] chore: Added code to handle multi-modal content and metadata. todo: Exhaustive tests to improve coverage. --- pydantic_ai_slim/pydantic_ai/mcp.py | 64 +++++++++++++++++++++++++---- 1 file changed, 56 insertions(+), 8 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/mcp.py b/pydantic_ai_slim/pydantic_ai/mcp.py index 9660c400a2..c451101a2e 100644 --- a/pydantic_ai_slim/pydantic_ai/mcp.py +++ b/pydantic_ai_slim/pydantic_ai/mcp.py @@ -403,23 +403,71 @@ async def _map_tool_result_part( pass return text if metadata is None else messages.ToolReturn(return_value=text, metadata=metadata) elif isinstance(part, mcp_types.ImageContent): - return messages.BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType) + binary_response = messages.BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType) + return ( + binary_response + if metadata is None + else messages.ToolReturn( + return_value=f'See file {binary_response.identifier}', + content=[f'This is file {binary_response.identifier}:', binary_response], + metadata=metadata, + ) + ) elif isinstance(part, mcp_types.AudioContent): # NOTE: The FastMCP server doesn't support audio content. # See for more details. - return messages.BinaryContent( + binary_response = messages.BinaryContent( data=base64.b64decode(part.data), media_type=part.mimeType ) # pragma: no cover + return ( + binary_response + if metadata is None + else messages.ToolReturn( + return_value=f'See file {binary_response.identifier}', + content=[f'This is file {binary_response.identifier}:', binary_response], + metadata=metadata, + ) + ) elif isinstance(part, mcp_types.EmbeddedResource): resource = part.resource - return self._get_content(resource) - elif isinstance(part, mcp_types.ResourceLink): - resource_result: mcp_types.ReadResourceResult = await self._client.read_resource(part.uri) + response = self._get_content(resource) return ( - self._get_content(resource_result.contents[0]) - if len(resource_result.contents) == 1 - else [self._get_content(resource) for resource in resource_result.contents] + response + if metadata is None + else messages.ToolReturn( + return_value=response if isinstance(response, str) else f'See file {response.identifier}', + content=None if isinstance(response, str) else [f'This is file {response.identifier}:', response], + metadata=metadata, + ) ) + elif isinstance(part, mcp_types.ResourceLink): + resource_result: mcp_types.ReadResourceResult = await self._client.read_resource(part.uri) + if len(resource_result.contents) == 1: + response = self._get_content(resource_result.contents[0]) + return ( + response + if metadata is None + else messages.ToolReturn( + return_value=response if isinstance(response, str) else f'See file {response.identifier}', + content=None + if isinstance(response, str) + else [f'This is file {response.identifier}:', response], + metadata=metadata, + ) + ) + else: + responses = [self._get_content(resource) for resource in resource_result.contents] + return [ + response + if isinstance(response, str) + else messages.ToolReturn( + return_value=response if isinstance(response, str) else f'See file {response.identifier}', + content=None + if isinstance(response, str) + else [f'This is file {response.identifier}:', response], + ) + for response in responses + ] else: assert_never(part) From fd806f70621cedca53d9bf1b586490b59e7ccc43 Mon Sep 17 00:00:00 2001 From: Anirban Basu Date: Wed, 5 Nov 2025 22:32:20 +0900 Subject: [PATCH 6/7] chore: Added no cover pragma to some portions of the code. --- pydantic_ai_slim/pydantic_ai/mcp.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/mcp.py b/pydantic_ai_slim/pydantic_ai/mcp.py index c451101a2e..c10ea88b00 100644 --- a/pydantic_ai_slim/pydantic_ai/mcp.py +++ b/pydantic_ai_slim/pydantic_ai/mcp.py @@ -416,10 +416,8 @@ async def _map_tool_result_part( elif isinstance(part, mcp_types.AudioContent): # NOTE: The FastMCP server doesn't support audio content. # See for more details. - binary_response = messages.BinaryContent( - data=base64.b64decode(part.data), media_type=part.mimeType - ) # pragma: no cover - return ( + binary_response = messages.BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType) + return ( # pragma: no cover binary_response if metadata is None else messages.ToolReturn( @@ -456,8 +454,8 @@ async def _map_tool_result_part( ) ) else: - responses = [self._get_content(resource) for resource in resource_result.contents] - return [ + responses = [self._get_content(resource) for resource in resource_result.contents] # pragma: no cover + return [ # pragma: no cover response if isinstance(response, str) else messages.ToolReturn( From 39d47b5363ea42792967ac615a079372860c1027 Mon Sep 17 00:00:00 2001 From: Anirban Basu Date: Sun, 9 Nov 2025 18:14:52 +0900 Subject: [PATCH 7/7] experimental: Potential implementation of 3, 4 and 5 described in https://github.com/pydantic/pydantic-ai/pull/3339#issuecomment-3503048730 --- pydantic_ai_slim/pydantic_ai/mcp.py | 156 ++++++++++++++-------------- 1 file changed, 79 insertions(+), 77 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/mcp.py b/pydantic_ai_slim/pydantic_ai/mcp.py index 7569497f06..239d7a6364 100644 --- a/pydantic_ai_slim/pydantic_ai/mcp.py +++ b/pydantic_ai_slim/pydantic_ai/mcp.py @@ -254,23 +254,80 @@ async def direct_call_tool( ): # The MCP SDK wraps primitives and generic types like list in a `result` key, but we want to use the raw value returned by the tool function. # See https://github.com/modelcontextprotocol/python-sdk#structured-output - if isinstance(structured, dict) and ( - (len(structured) == 1 and 'result' in structured) - or (len(structured) == 2 and 'result' in structured and '_meta' in structured) - ): + if isinstance(structured, dict) and len(structured) == 1 and 'result' in structured: return ( - messages.ToolReturn(return_value=structured['result'], metadata=structured['_meta']) - if structured.get('_meta', None) is not None + messages.ToolReturn(return_value=structured['result'], metadata=result.meta) + if result.meta else structured['result'] ) + return messages.ToolReturn(return_value=structured, metadata=result.meta) if result.meta else structured + + mapped_part_metadata_tuple_list = [await self._map_tool_result_part(part) for part in result.content] + if ( + all(mapped_part_metadata_tuple[1] is None for mapped_part_metadata_tuple in mapped_part_metadata_tuple_list) + and result.meta is None + ): + # There is no metadata in the tool result or its parts, return just the mapped values return ( - messages.ToolReturn(return_value=structured, metadata=structured['_meta']) - if structured.get('_meta', None) is not None - else structured + mapped_part_metadata_tuple_list[0][0] + if len(mapped_part_metadata_tuple_list[0]) == 1 + else [mapped_part_metadata_tuple[0] for mapped_part_metadata_tuple in mapped_part_metadata_tuple_list] + ) + elif ( + all(mapped_part_metadata_tuple[1] is None for mapped_part_metadata_tuple in mapped_part_metadata_tuple_list) + and result.meta is not None + ): + # There is no metadata in the tool result parts, but there is metadata in the tool result + return messages.ToolReturn( + return_value=( + mapped_part_metadata_tuple_list[0][0] + if len(mapped_part_metadata_tuple_list[0]) == 1 + else [ + mapped_part_metadata_tuple[0] for mapped_part_metadata_tuple in mapped_part_metadata_tuple_list + ] + ), + metadata=result.meta, + ) + else: + # There is metadata in the tool result parts and there may be a metadata in the tool result, return a ToolReturn object + return_values: list[Any] = [] + user_contents: list[Any] = [] + return_metadata: dict[str, Any] = {} + for idx, (mapped_part, part_metadata) in enumerate(mapped_part_metadata_tuple_list): + if part_metadata is not None: + # Merge the metadata dictionaries, with part metadata taking precedence + if return_metadata.get('content', None) is None: + # Create an empty list if it doesn't exist yet + return_metadata['content'] = list[dict[str, Any]]() + return_metadata['content'].append({str(idx): part_metadata}) + if isinstance(mapped_part, messages.BinaryContent): + identifier = mapped_part.identifier + + return_values.append(f'See file {identifier}') + user_contents.append([f'This is file {identifier}:', mapped_part]) + else: + user_contents.append(mapped_part) + + if result.meta is not None and return_metadata.get('content', None) is not None: + # Merge the tool result metadata into the return metadata, with part metadata taking precedence + return_metadata['result'] = result.meta + elif result.meta is not None and return_metadata.get('content', None) is None: + return_metadata = result.meta + elif ( + result.meta is None + and return_metadata.get('content', None) is not None + and len(return_metadata['content']) == 1 + ): + # If there is only one content metadata, unwrap it + return_metadata = return_metadata['content'][0] + # TODO: What else should we cover here? + + # Finally, construct and return the ToolReturn object + return messages.ToolReturn( + return_value=return_values, + content=user_contents, + metadata=return_metadata, ) - - mapped = [await self._map_tool_result_part(part) for part in result.content] - return mapped[0] if len(mapped) == 1 else mapped async def call_tool( self, @@ -385,87 +442,32 @@ async def _sampling_callback( async def _map_tool_result_part( self, part: mcp_types.ContentBlock - ) -> str | messages.ToolReturn | messages.BinaryContent | dict[str, Any] | list[Any]: + ) -> tuple[str | messages.BinaryContent | dict[str, Any] | list[Any], dict[str, Any] | None]: # See https://github.com/jlowin/fastmcp/blob/main/docs/servers/tools.mdx#return-values - # Let's also check for metadata but it can be present in not just TextContent metadata: dict[str, Any] | None = part.meta if isinstance(part, mcp_types.TextContent): text = part.text if text.startswith(('[', '{')): try: - return ( - pydantic_core.from_json(text) - if metadata is None - else messages.ToolReturn(return_value=pydantic_core.from_json(text), metadata=metadata) - ) + return pydantic_core.from_json(text), metadata except ValueError: pass - return text if metadata is None else messages.ToolReturn(return_value=text, metadata=metadata) + return text, metadata elif isinstance(part, mcp_types.ImageContent): - binary_response = messages.BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType) - return ( - binary_response - if metadata is None - else messages.ToolReturn( - return_value=f'See file {binary_response.identifier}', - content=[f'This is file {binary_response.identifier}:', binary_response], - metadata=metadata, - ) - ) - elif isinstance(part, mcp_types.AudioContent): + return messages.BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType), metadata + elif isinstance(part, mcp_types.AudioContent): # pragma: no cover # NOTE: The FastMCP server doesn't support audio content. # See for more details. - binary_response = messages.BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType) - return ( # pragma: no cover - binary_response - if metadata is None - else messages.ToolReturn( - return_value=f'See file {binary_response.identifier}', - content=[f'This is file {binary_response.identifier}:', binary_response], - metadata=metadata, - ) - ) + return messages.BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType), metadata elif isinstance(part, mcp_types.EmbeddedResource): - resource = part.resource - response = self._get_content(resource) - return ( - response - if metadata is None - else messages.ToolReturn( - return_value=response if isinstance(response, str) else f'See file {response.identifier}', - content=None if isinstance(response, str) else [f'This is file {response.identifier}:', response], - metadata=metadata, - ) - ) + return self._get_content(part.resource), metadata elif isinstance(part, mcp_types.ResourceLink): resource_result: mcp_types.ReadResourceResult = await self._client.read_resource(part.uri) if len(resource_result.contents) == 1: - response = self._get_content(resource_result.contents[0]) - return ( - response - if metadata is None - else messages.ToolReturn( - return_value=response if isinstance(response, str) else f'See file {response.identifier}', - content=None - if isinstance(response, str) - else [f'This is file {response.identifier}:', response], - metadata=metadata, - ) - ) + return self._get_content(resource_result.contents[0]), metadata else: - responses = [self._get_content(resource) for resource in resource_result.contents] # pragma: no cover - return [ # pragma: no cover - response - if isinstance(response, str) - else messages.ToolReturn( - return_value=response if isinstance(response, str) else f'See file {response.identifier}', - content=None - if isinstance(response, str) - else [f'This is file {response.identifier}:', response], - ) - for response in responses - ] + return [self._get_content(resource) for resource in resource_result.contents], metadata else: assert_never(part) @@ -941,7 +943,7 @@ def __eq__(self, value: object, /) -> bool: | messages.ToolReturn | dict[str, Any] | list[Any] - | Sequence[str | messages.BinaryContent | messages.ToolReturn | dict[str, Any] | list[Any]] + | Sequence[str | messages.BinaryContent | dict[str, Any] | list[Any]] ) """The result type of an MCP tool call."""