diff --git a/src/backend/base/langflow/components/tools/mcp_component.py b/src/backend/base/langflow/components/tools/mcp_component.py index 06ecfd46b..fbe920b08 100644 --- a/src/backend/base/langflow/components/tools/mcp_component.py +++ b/src/backend/base/langflow/components/tools/mcp_component.py @@ -17,7 +17,7 @@ from langflow.inputs.inputs import InputTypes from langflow.io import MessageTextInput, MultilineInput, Output, TabInput from langflow.io.schema import flatten_schema, schema_to_langflow_inputs from langflow.logging import logger -from langflow.schema import Message +from langflow.schema import DataFrame def maybe_unflatten_dict(flat: dict[str, Any]) -> dict[str, Any]: @@ -409,7 +409,7 @@ class MCPToolsComponent(Component): logger.exception(msg) raise ValueError(msg) from e - async def build_output(self) -> Message: + async def build_output(self) -> DataFrame: """Build output with improved error handling and validation.""" try: await self.update_tools() @@ -426,8 +426,12 @@ class MCPToolsComponent(Component): output = await exec_tool.coroutine(**unflattened_kwargs) - return Message(text=output.content[len(output.content) - 1].text) - return Message(text="You must select a tool", error=True) + tool_content = [] + for item in output.content: + item_dict = item.model_dump() + tool_content.append(item_dict) + return DataFrame(data=tool_content) + return DataFrame(data=[{"error": "You must select a tool"}]) except Exception as e: msg = f"Error in build_output: {e!s}" logger.exception(msg) diff --git a/src/backend/tests/unit/components/tools/test_mcp_component.py b/src/backend/tests/unit/components/tools/test_mcp_component.py index 2feed2b6a..6c22ee1bc 100644 --- a/src/backend/tests/unit/components/tools/test_mcp_component.py +++ b/src/backend/tests/unit/components/tools/test_mcp_component.py @@ -108,7 +108,10 @@ class TestMCPToolsComponent(ComponentTestBaseWithoutClient): # Mock the coroutine response mock_response = AsyncMock() - mock_response.content = [MagicMock(text="Test response")] + mock_content_item = MagicMock() + mock_content_item.text = "Test response" + mock_content_item.model_dump.return_value = {"text": "Test response"} + mock_response.content = [mock_content_item] mock_create_coroutine.return_value = AsyncMock(return_value=mock_response) # Create a mock tool and add it to the cache @@ -126,7 +129,8 @@ class TestMCPToolsComponent(ComponentTestBaseWithoutClient): mock_get_inputs.return_value = {"test_tool": [mock_input]} output = await component.build_output() - assert output.text == "Test response" + # Use iloc to access the first row's 'text' column value + assert output.iloc[0]["text"] == "Test response" # Verify the mocks were called correctly mock_get_inputs.assert_called_once_with(component.tools) mock_structured_tool.coroutine.assert_called_once_with(test_param="test value")