diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index abdb0aa5d0..cc6dbd076c 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -18,7 +18,7 @@ from datetime import timedelta from functools import partial from inspect import isawaitable -from typing import TYPE_CHECKING, Any, Literal, TypedDict, cast +from typing import TYPE_CHECKING, Any, Literal, TypeAlias, TypedDict, cast from opentelemetry import propagate from opentelemetry import trace as otel_trace @@ -60,6 +60,18 @@ logger = logging.getLogger(__name__) +if TYPE_CHECKING: + _MCPSamplingContentBlock: TypeAlias = ( + types.TextContent + | types.ImageContent + | types.AudioContent + | types.EmbeddedResource + | types.ResourceLink + | types.ToolUseContent + ) +else: + _MCPSamplingContentBlock = Any + class MCPSpecificApproval(TypedDict, total=False): """Represents the specific approval mode for an MCP tool. @@ -779,14 +791,21 @@ def _parse_content_from_mcp( def _prepare_content_for_mcp( self, content: Content, - ) -> ( - types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink | None - ): + ) -> _MCPSamplingContentBlock | None: """Prepare an Agent Framework content type for MCP.""" from mcp import types if content.type == "text": return types.TextContent(type="text", text=content.text) # type: ignore[attr-defined] + if content.type == "function_call": + if not content.call_id or not content.name: + return None + return types.ToolUseContent( + type="tool_use", + id=content.call_id, + name=content.name, + input=content.parse_arguments() or {}, + ) if content.type == "data": if content.media_type and content.media_type.startswith("image/"): return types.ImageContent(type="image", data=content.uri, mimeType=content.media_type) # type: ignore[attr-defined] @@ -821,13 +840,9 @@ def _prepare_content_for_mcp( def _prepare_message_for_mcp( self, content: Message, - ) -> list[ - types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink - ]: + ) -> list[_MCPSamplingContentBlock]: """Prepare a Message for MCP format.""" - messages: list[ - types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink - ] = [] + messages: list[_MCPSamplingContentBlock] = [] for item in content.contents: mcp_content = self._prepare_content_for_mcp(item) if mcp_content: @@ -1433,7 +1448,7 @@ async def sampling_callback( self, context: RequestContext[ClientSession, Any], params: types.CreateMessageRequestParams, - ) -> types.CreateMessageResult | types.ErrorData: + ) -> types.CreateMessageResult | types.CreateMessageResultWithTools | types.ErrorData: """Callback function for sampling. This function is called when the MCP server sends a ``sampling/createMessage`` @@ -1460,8 +1475,8 @@ async def sampling_callback( params: The message creation request parameters. Returns: - Either a CreateMessageResult with the generated message or ErrorData if the request - is denied, rate limited, or generation fails. + Either a CreateMessageResult/CreateMessageResultWithTools with the generated message or ErrorData if the + request is denied, rate limited, or generation fails. """ from mcp import types @@ -1544,7 +1559,21 @@ async def sampling_callback( code=types.INTERNAL_ERROR, message="Failed to get chat message content.", ) - mcp_contents = self._prepare_message_for_mcp(response.messages[0]) + mcp_contents = [ + mcp_content for message in response.messages for mcp_content in self._prepare_message_for_mcp(message) + ] + tool_use_contents: list[types.SamplingMessageContentBlock] = [] + for content in mcp_contents: + if isinstance(content, types.ToolUseContent): + tool_use_contents.append(content) + if tool_use_contents: + return types.CreateMessageResultWithTools( + role="assistant", + content=tool_use_contents, + model=response.model or "unknown", + stopReason="toolUse", + ) + # grab the first content that is of type TextContent or ImageContent mcp_content = next( (content for content in mcp_contents if isinstance(content, (types.TextContent, types.ImageContent))), diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 4b61cd2451..62f5684c48 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -2954,6 +2954,65 @@ async def test_mcp_tool_sampling_callback_no_response_and_successful_message_cre assert success.content.text == "Hello" +async def test_mcp_tool_sampling_callback_returns_tool_use_results(): + """Test sampling callback returns structured MCP tool-use content from function calls.""" + tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve) + + mock_chat_client = AsyncMock() + mock_chat_client.get_response.return_value = Mock( + messages=[ + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="call-1", + name="Answer", + arguments={"answer": "hello"}, + ), + ], + ), + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="call-2", + name="Citations", + arguments={"sources": ["docs"]}, + ), + ], + ), + ], + model="test-model", + ) + tool.client = mock_chat_client + + params = Mock() + params.messages = [types.PromptMessage(role="user", content=types.TextContent(type="text", text="Answer"))] + params.temperature = None + params.maxTokens = 100 + params.stopSequences = None + params.systemPrompt = None + params.tools = [ + types.Tool(name="Answer", description="Return an answer", inputSchema={"type": "object"}), + types.Tool(name="Citations", description="Return source citations", inputSchema={"type": "object"}), + ] + params.toolChoice = None + + result = await tool.sampling_callback(Mock(), params) + + assert isinstance(result, types.CreateMessageResultWithTools) + assert result.role == "assistant" + assert result.model == "test-model" + assert result.stopReason == "toolUse" + assert isinstance(result.content, list) + tool_use_contents = [content for content in result.content if isinstance(content, types.ToolUseContent)] + assert tool_use_contents == result.content + assert [(content.id, content.name, content.input) for content in tool_use_contents] == [ + ("call-1", "Answer", {"answer": "hello"}), + ("call-2", "Citations", {"sources": ["docs"]}), + ] + + async def test_mcp_tool_logging_callback_logs_at_requested_level() -> None: tool = MCPStdioTool(name="test_tool", command="python")