mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-11 09:08:31 +00:00
Fix: Add connect_timeout parameter to get_mcp_tools method
- Add Optional import to crew_base.py for type annotation - Modify get_mcp_tools to accept connect_timeout parameter with 30s default - Pass connect_timeout to MCPServerAdapter constructor - Update existing test to verify default timeout is passed - Add test for custom timeout values - Add test for backward compatibility Fixes issue #3463 where @CrewBase users couldn't configure MCP server connection timeouts Co-Authored-By: João <joao@crewai.com>
This commit is contained in:
@@ -1,7 +1,7 @@
|
|||||||
import inspect
|
import inspect
|
||||||
import logging
|
import logging
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any, Callable, Dict, TypeVar, cast, List
|
from typing import Any, Callable, Dict, TypeVar, cast, List, Optional
|
||||||
from crewai.tools import BaseTool
|
from crewai.tools import BaseTool
|
||||||
|
|
||||||
import yaml
|
import yaml
|
||||||
@@ -86,7 +86,7 @@ def CrewBase(cls: T) -> T:
|
|||||||
import types
|
import types
|
||||||
return types.MethodType(_close_mcp_server, self)
|
return types.MethodType(_close_mcp_server, self)
|
||||||
|
|
||||||
def get_mcp_tools(self, *tool_names: list[str]) -> List[BaseTool]:
|
def get_mcp_tools(self, *tool_names: list[str], connect_timeout: Optional[int] = 30) -> List[BaseTool]:
|
||||||
if not self.mcp_server_params:
|
if not self.mcp_server_params:
|
||||||
return []
|
return []
|
||||||
|
|
||||||
@@ -94,7 +94,7 @@ def CrewBase(cls: T) -> T:
|
|||||||
|
|
||||||
adapter = getattr(self, '_mcp_server_adapter', None)
|
adapter = getattr(self, '_mcp_server_adapter', None)
|
||||||
if not adapter:
|
if not adapter:
|
||||||
self._mcp_server_adapter = MCPServerAdapter(self.mcp_server_params)
|
self._mcp_server_adapter = MCPServerAdapter(self.mcp_server_params, connect_timeout=connect_timeout)
|
||||||
|
|
||||||
return self._mcp_server_adapter.tools.filter_by_names(tool_names or None)
|
return self._mcp_server_adapter.tools.filter_by_names(tool_names or None)
|
||||||
|
|
||||||
|
|||||||
@@ -284,4 +284,53 @@ def test_internal_crew_with_mcp():
|
|||||||
assert crew.reporting_analyst().tools == [simple_tool, another_simple_tool]
|
assert crew.reporting_analyst().tools == [simple_tool, another_simple_tool]
|
||||||
assert crew.researcher().tools == [simple_tool]
|
assert crew.researcher().tools == [simple_tool]
|
||||||
|
|
||||||
adapter_mock.assert_called_once_with({"host": "localhost", "port": 8000})
|
adapter_mock.assert_called_once_with({"host": "localhost", "port": 8000}, connect_timeout=30)
|
||||||
|
|
||||||
|
def test_internal_crew_with_mcp_custom_timeout():
|
||||||
|
with patch("embedchain.client.Client.setup"):
|
||||||
|
from crewai_tools import MCPServerAdapter
|
||||||
|
from crewai_tools.adapters.mcp_adapter import ToolCollection
|
||||||
|
|
||||||
|
@CrewBase
|
||||||
|
class CrewWithCustomTimeout(InternalCrew):
|
||||||
|
mcp_server_params = {"host": "localhost", "port": 8000}
|
||||||
|
|
||||||
|
@agent
|
||||||
|
def test_agent(self):
|
||||||
|
return Agent(
|
||||||
|
config=self.agents_config["researcher"],
|
||||||
|
tools=self.get_mcp_tools("simple_tool", connect_timeout=60),
|
||||||
|
)
|
||||||
|
|
||||||
|
mock = Mock(spec=MCPServerAdapter)
|
||||||
|
mock.tools = ToolCollection([simple_tool])
|
||||||
|
with patch("crewai_tools.MCPServerAdapter", return_value=mock) as adapter_mock:
|
||||||
|
crew = CrewWithCustomTimeout()
|
||||||
|
assert crew.test_agent().tools == [simple_tool]
|
||||||
|
|
||||||
|
adapter_mock.assert_called_once_with({"host": "localhost", "port": 8000}, connect_timeout=60)
|
||||||
|
|
||||||
|
|
||||||
|
def test_internal_crew_with_mcp_backward_compatibility():
|
||||||
|
with patch("embedchain.client.Client.setup"):
|
||||||
|
from crewai_tools import MCPServerAdapter
|
||||||
|
from crewai_tools.adapters.mcp_adapter import ToolCollection
|
||||||
|
|
||||||
|
@CrewBase
|
||||||
|
class CrewWithoutTimeout(InternalCrew):
|
||||||
|
mcp_server_params = {"host": "localhost", "port": 8000}
|
||||||
|
|
||||||
|
@agent
|
||||||
|
def test_agent(self):
|
||||||
|
return Agent(
|
||||||
|
config=self.agents_config["researcher"],
|
||||||
|
tools=self.get_mcp_tools("simple_tool"),
|
||||||
|
)
|
||||||
|
|
||||||
|
mock = Mock(spec=MCPServerAdapter)
|
||||||
|
mock.tools = ToolCollection([simple_tool])
|
||||||
|
with patch("crewai_tools.MCPServerAdapter", return_value=mock) as adapter_mock:
|
||||||
|
crew = CrewWithoutTimeout()
|
||||||
|
assert crew.test_agent().tools == [simple_tool]
|
||||||
|
|
||||||
|
adapter_mock.assert_called_once_with({"host": "localhost", "port": 8000}, connect_timeout=30)
|
||||||
|
|||||||
Reference in New Issue
Block a user