diff --git a/src/backend/base/langflow/components/tools/mcp_sse.py b/src/backend/base/langflow/components/tools/mcp_sse.py index 82e6c34e5..af9d857e2 100644 --- a/src/backend/base/langflow/components/tools/mcp_sse.py +++ b/src/backend/base/langflow/components/tools/mcp_sse.py @@ -1,4 +1,5 @@ # from langflow.field_typing import Data +import asyncio from contextlib import AsyncExitStack import httpx @@ -10,7 +11,6 @@ from langflow.components.tools.mcp_stdio import create_input_schema_from_json_sc from langflow.custom import Component from langflow.field_typing import Tool from langflow.io import MessageTextInput, Output -from langflow.utils.async_helpers import timeout_context # Define constant for status code HTTP_TEMPORARY_REDIRECT = 307 @@ -38,20 +38,32 @@ class MCPSseClient: if headers is None: headers = {} url = await self.pre_check_redirect(url) - - async with timeout_context(timeout_seconds): - sse_transport = await self.exit_stack.enter_async_context( - sse_client(url, headers, timeout_seconds, sse_read_timeout_seconds) + try: + await asyncio.wait_for( + self._connect_with_timeout(url, headers, timeout_seconds, sse_read_timeout_seconds), + timeout=timeout_seconds, ) - self.sse, self.write = sse_transport - self.session = await self.exit_stack.enter_async_context(ClientSession(self.sse, self.write)) - - await self.session.initialize() - # List available tools + if self.session is None: + msg = "Session not initialized" + raise ValueError(msg) response = await self.session.list_tools() + except asyncio.TimeoutError as err: + error_message = f"Connection to {url} timed out after {timeout_seconds} seconds" + raise TimeoutError(error_message) from err + else: # Only executed if no TimeoutError return response.tools + async def _connect_with_timeout( + self, url: str, headers: dict[str, str] | None, timeout_seconds: int, sse_read_timeout_seconds: int + ): + sse_transport = await self.exit_stack.enter_async_context( + sse_client(url, headers, timeout_seconds, sse_read_timeout_seconds) + ) + self.sse, self.write = sse_transport + self.session = await self.exit_stack.enter_async_context(ClientSession(self.sse, self.write)) + await self.session.initialize() + class MCPSse(Component): client = MCPSseClient() @@ -89,6 +101,7 @@ class MCPSse(Component): Tool( name=tool.name, # maybe format this description=tool.description, + args_schema=args_schema, coroutine=create_tool_coroutine(tool.name, args_schema, self.client.session), func=create_tool_func(tool.name, self.client.session), ) diff --git a/src/backend/tests/unit/api/v1/test_api_schemas.py b/src/backend/tests/unit/api/v1/test_api_schemas.py index f58bcf45d..2a73afe22 100644 --- a/src/backend/tests/unit/api/v1/test_api_schemas.py +++ b/src/backend/tests/unit/api/v1/test_api_schemas.py @@ -6,7 +6,6 @@ from langflow.api.v1.schemas import ResultDataResponse, VertexBuildResponse from langflow.schema.schema import OutputValue from langflow.serialization import serialize from langflow.services.tracing.schema import Log - from pydantic import BaseModel # Use a smaller test size for hypothesis diff --git a/src/backend/tests/unit/base/tools/test_component_toolkit.py b/src/backend/tests/unit/base/tools/test_component_toolkit.py index c9d627973..bf8df7a18 100644 --- a/src/backend/tests/unit/base/tools/test_component_toolkit.py +++ b/src/backend/tests/unit/base/tools/test_component_toolkit.py @@ -8,7 +8,6 @@ from langflow.components.outputs.chat import ChatOutput from langflow.components.tools.calculator import CalculatorToolComponent from langflow.graph import Graph from langflow.schema.data import Data - from pydantic import BaseModel diff --git a/src/backend/tests/unit/base/tools/test_toolmodemixin.py b/src/backend/tests/unit/base/tools/test_toolmodemixin.py index 059464113..27518360c 100644 --- a/src/backend/tests/unit/base/tools/test_toolmodemixin.py +++ b/src/backend/tests/unit/base/tools/test_toolmodemixin.py @@ -21,7 +21,6 @@ from langflow.io import ( TableInput, ) from langflow.schema import Data - from pydantic import BaseModel diff --git a/src/backend/tests/unit/components/helpers/test_structured_output_component.py b/src/backend/tests/unit/components/helpers/test_structured_output_component.py index afe7adf59..8e1d1c32d 100644 --- a/src/backend/tests/unit/components/helpers/test_structured_output_component.py +++ b/src/backend/tests/unit/components/helpers/test_structured_output_component.py @@ -5,8 +5,8 @@ import pytest from langflow.components.helpers.structured_output import StructuredOutputComponent from langflow.helpers.base_model import build_model_from_schema from langflow.inputs.inputs import TableInput - from pydantic import BaseModel + from tests.base import ComponentTestBaseWithoutClient from tests.unit.mock_language_model import MockLanguageModel diff --git a/src/backend/tests/unit/graph/graph/state/test_state_model.py b/src/backend/tests/unit/graph/graph/state/test_state_model.py index 4278637b9..869366b8e 100644 --- a/src/backend/tests/unit/graph/graph/state/test_state_model.py +++ b/src/backend/tests/unit/graph/graph/state/test_state_model.py @@ -5,7 +5,6 @@ from langflow.graph import Graph from langflow.graph.graph.constants import Finish from langflow.graph.state.model import create_state_model from langflow.template.field.base import UNDEFINED - from pydantic import Field diff --git a/src/backend/tests/unit/helpers/test_base_model_from_schema.py b/src/backend/tests/unit/helpers/test_base_model_from_schema.py index 3bc2bcd2d..d07a4908e 100644 --- a/src/backend/tests/unit/helpers/test_base_model_from_schema.py +++ b/src/backend/tests/unit/helpers/test_base_model_from_schema.py @@ -4,9 +4,8 @@ from typing import Any import pytest from langflow.helpers.base_model import build_model_from_schema -from pydantic_core import PydanticUndefined - from pydantic import BaseModel +from pydantic_core import PydanticUndefined class TestBuildModelFromSchema: diff --git a/src/backend/tests/unit/inputs/test_inputs.py b/src/backend/tests/unit/inputs/test_inputs.py index eacb2a40b..2a946a2a9 100644 --- a/src/backend/tests/unit/inputs/test_inputs.py +++ b/src/backend/tests/unit/inputs/test_inputs.py @@ -23,7 +23,6 @@ from langflow.inputs.inputs import ( ) from langflow.inputs.utils import instantiate_input from langflow.schema.message import Message - from pydantic import ValidationError diff --git a/src/backend/tests/unit/mock_language_model.py b/src/backend/tests/unit/mock_language_model.py index 8c06e112b..186d127cf 100644 --- a/src/backend/tests/unit/mock_language_model.py +++ b/src/backend/tests/unit/mock_language_model.py @@ -1,9 +1,8 @@ from unittest.mock import MagicMock from langchain_core.language_models import BaseLanguageModel -from typing_extensions import override - from pydantic import BaseModel, Field +from typing_extensions import override class MockLanguageModel(BaseLanguageModel, BaseModel): diff --git a/src/backend/tests/unit/serialization/test_serialization.py b/src/backend/tests/unit/serialization/test_serialization.py index dff548472..e8110906c 100644 --- a/src/backend/tests/unit/serialization/test_serialization.py +++ b/src/backend/tests/unit/serialization/test_serialization.py @@ -9,7 +9,6 @@ from hypothesis import strategies as st from langchain_core.documents import Document from langflow.serialization.constants import MAX_ITEMS_LENGTH, MAX_TEXT_LENGTH from langflow.serialization.serialization import serialize, serialize_or_str - from pydantic import BaseModel as PydanticBaseModel from pydantic.v1 import BaseModel as PydanticV1BaseModel diff --git a/src/backend/tests/unit/test_schema.py b/src/backend/tests/unit/test_schema.py index 86aaf4500..97fd930e3 100644 --- a/src/backend/tests/unit/test_schema.py +++ b/src/backend/tests/unit/test_schema.py @@ -7,7 +7,6 @@ from langflow.schema.data import Data from langflow.template import Input, Output from langflow.template.field.base import UNDEFINED from langflow.type_extraction.type_extraction import post_process_type - from pydantic import ValidationError diff --git a/src/backend/tests/unit/test_template.py b/src/backend/tests/unit/test_template.py index 6a4ad938c..6b2127178 100644 --- a/src/backend/tests/unit/test_template.py +++ b/src/backend/tests/unit/test_template.py @@ -2,7 +2,6 @@ import importlib import pytest from langflow.utils.util import build_template_from_function, get_base_classes, get_default_factory - from pydantic import BaseModel