mcp-use 1.3.11__py3-none-any.whl → 1.3.13__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of mcp-use might be problematic. Click here for more details.
- mcp_use/__init__.py +1 -1
- mcp_use/adapters/.deprecated +0 -0
- mcp_use/adapters/__init__.py +18 -7
- mcp_use/adapters/base.py +12 -185
- mcp_use/adapters/langchain_adapter.py +12 -264
- mcp_use/agents/adapters/__init__.py +10 -0
- mcp_use/agents/adapters/base.py +193 -0
- mcp_use/agents/adapters/langchain_adapter.py +228 -0
- mcp_use/agents/base.py +1 -1
- mcp_use/agents/managers/__init__.py +19 -0
- mcp_use/agents/managers/base.py +36 -0
- mcp_use/agents/managers/server_manager.py +131 -0
- mcp_use/agents/managers/tools/__init__.py +15 -0
- mcp_use/agents/managers/tools/base_tool.py +19 -0
- mcp_use/agents/managers/tools/connect_server.py +69 -0
- mcp_use/agents/managers/tools/disconnect_server.py +43 -0
- mcp_use/agents/managers/tools/get_active_server.py +29 -0
- mcp_use/agents/managers/tools/list_servers_tool.py +53 -0
- mcp_use/agents/managers/tools/search_tools.py +328 -0
- mcp_use/agents/mcpagent.py +88 -47
- mcp_use/agents/remote.py +168 -129
- mcp_use/auth/.deprecated +0 -0
- mcp_use/auth/__init__.py +19 -4
- mcp_use/auth/bearer.py +11 -12
- mcp_use/auth/oauth.py +11 -620
- mcp_use/auth/oauth_callback.py +16 -207
- mcp_use/client/__init__.py +1 -0
- mcp_use/client/auth/__init__.py +6 -0
- mcp_use/client/auth/bearer.py +23 -0
- mcp_use/client/auth/oauth.py +629 -0
- mcp_use/client/auth/oauth_callback.py +214 -0
- mcp_use/client/client.py +356 -0
- mcp_use/client/config.py +106 -0
- mcp_use/client/connectors/__init__.py +20 -0
- mcp_use/client/connectors/base.py +470 -0
- mcp_use/client/connectors/http.py +304 -0
- mcp_use/client/connectors/sandbox.py +332 -0
- mcp_use/client/connectors/stdio.py +109 -0
- mcp_use/client/connectors/utils.py +13 -0
- mcp_use/client/connectors/websocket.py +257 -0
- mcp_use/client/exceptions.py +31 -0
- mcp_use/client/middleware/__init__.py +50 -0
- mcp_use/client/middleware/logging.py +31 -0
- mcp_use/client/middleware/metrics.py +314 -0
- mcp_use/client/middleware/middleware.py +266 -0
- mcp_use/client/session.py +162 -0
- mcp_use/client/task_managers/__init__.py +20 -0
- mcp_use/client/task_managers/base.py +145 -0
- mcp_use/client/task_managers/sse.py +84 -0
- mcp_use/client/task_managers/stdio.py +69 -0
- mcp_use/client/task_managers/streamable_http.py +86 -0
- mcp_use/client/task_managers/websocket.py +68 -0
- mcp_use/client.py +12 -320
- mcp_use/config.py +20 -92
- mcp_use/connectors/.deprecated +0 -0
- mcp_use/connectors/__init__.py +46 -20
- mcp_use/connectors/base.py +12 -447
- mcp_use/connectors/http.py +13 -288
- mcp_use/connectors/sandbox.py +13 -297
- mcp_use/connectors/stdio.py +13 -96
- mcp_use/connectors/utils.py +15 -8
- mcp_use/connectors/websocket.py +13 -252
- mcp_use/exceptions.py +33 -18
- mcp_use/managers/.deprecated +0 -0
- mcp_use/managers/__init__.py +56 -17
- mcp_use/managers/base.py +13 -31
- mcp_use/managers/server_manager.py +13 -119
- mcp_use/managers/tools/__init__.py +45 -15
- mcp_use/managers/tools/base_tool.py +5 -16
- mcp_use/managers/tools/connect_server.py +5 -67
- mcp_use/managers/tools/disconnect_server.py +5 -41
- mcp_use/managers/tools/get_active_server.py +5 -26
- mcp_use/managers/tools/list_servers_tool.py +5 -51
- mcp_use/managers/tools/search_tools.py +17 -321
- mcp_use/middleware/.deprecated +0 -0
- mcp_use/middleware/__init__.py +89 -0
- mcp_use/middleware/logging.py +19 -0
- mcp_use/middleware/metrics.py +41 -0
- mcp_use/middleware/middleware.py +55 -0
- mcp_use/session.py +13 -149
- mcp_use/task_managers/.deprecated +0 -0
- mcp_use/task_managers/__init__.py +48 -20
- mcp_use/task_managers/base.py +13 -140
- mcp_use/task_managers/sse.py +13 -79
- mcp_use/task_managers/stdio.py +13 -64
- mcp_use/task_managers/streamable_http.py +15 -81
- mcp_use/task_managers/websocket.py +13 -63
- mcp_use/telemetry/events.py +58 -0
- mcp_use/telemetry/telemetry.py +71 -1
- mcp_use/types/.deprecated +0 -0
- mcp_use/types/sandbox.py +13 -18
- {mcp_use-1.3.11.dist-info → mcp_use-1.3.13.dist-info}/METADATA +66 -40
- mcp_use-1.3.13.dist-info/RECORD +109 -0
- mcp_use-1.3.11.dist-info/RECORD +0 -60
- mcp_use-1.3.11.dist-info/licenses/LICENSE +0 -21
- /mcp_use/{observability → agents/observability}/__init__.py +0 -0
- /mcp_use/{observability → agents/observability}/callbacks_manager.py +0 -0
- /mcp_use/{observability → agents/observability}/laminar.py +0 -0
- /mcp_use/{observability → agents/observability}/langfuse.py +0 -0
- {mcp_use-1.3.11.dist-info → mcp_use-1.3.13.dist-info}/WHEEL +0 -0
- {mcp_use-1.3.11.dist-info → mcp_use-1.3.13.dist-info}/entry_points.txt +0 -0
|
@@ -0,0 +1,69 @@
|
|
|
1
|
+
from typing import ClassVar
|
|
2
|
+
|
|
3
|
+
from pydantic import BaseModel, Field
|
|
4
|
+
|
|
5
|
+
from mcp_use.agents.managers.tools.base_tool import MCPServerTool
|
|
6
|
+
from mcp_use.errors.error_formatting import format_error
|
|
7
|
+
from mcp_use.logging import logger
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class ServerActionInput(BaseModel):
|
|
11
|
+
"""Base input for server-related actions"""
|
|
12
|
+
|
|
13
|
+
server_name: str = Field(description="The name of the MCP server")
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class ConnectServerTool(MCPServerTool):
|
|
17
|
+
"""Tool for connecting to a specific MCP server."""
|
|
18
|
+
|
|
19
|
+
name: ClassVar[str] = "connect_to_mcp_server"
|
|
20
|
+
description: ClassVar[str] = (
|
|
21
|
+
"Connect to a specific MCP (Model Context Protocol) server to use its "
|
|
22
|
+
"tools. Use this tool to connect to a specific server and use its tools."
|
|
23
|
+
)
|
|
24
|
+
args_schema: ClassVar[type[BaseModel]] = ServerActionInput
|
|
25
|
+
|
|
26
|
+
async def _arun(self, server_name: str) -> str:
|
|
27
|
+
"""Connect to a specific MCP server."""
|
|
28
|
+
# Check if server exists
|
|
29
|
+
servers = self.server_manager.client.get_server_names()
|
|
30
|
+
if server_name not in servers:
|
|
31
|
+
available = ", ".join(servers) if servers else "none"
|
|
32
|
+
return f"Server '{server_name}' not found. Available servers: {available}"
|
|
33
|
+
|
|
34
|
+
# If we're already connected to this server, just return
|
|
35
|
+
if self.server_manager.active_server == server_name:
|
|
36
|
+
return f"Already connected to MCP server '{server_name}'"
|
|
37
|
+
|
|
38
|
+
try:
|
|
39
|
+
# Create or get session for this server
|
|
40
|
+
try:
|
|
41
|
+
session = self.server_manager.client.get_session(server_name)
|
|
42
|
+
logger.debug(f"Using existing session for server '{server_name}'")
|
|
43
|
+
except ValueError:
|
|
44
|
+
logger.debug(f"Creating new session for server '{server_name}'")
|
|
45
|
+
session = await self.server_manager.client.create_session(server_name)
|
|
46
|
+
|
|
47
|
+
# Set as active server
|
|
48
|
+
self.server_manager.active_server = server_name
|
|
49
|
+
|
|
50
|
+
# Initialize server tools if not already initialized
|
|
51
|
+
if server_name not in self.server_manager._server_tools:
|
|
52
|
+
connector = session.connector
|
|
53
|
+
self.server_manager._server_tools[
|
|
54
|
+
server_name
|
|
55
|
+
] = await self.server_manager.adapter._create_tools_from_connectors([connector])
|
|
56
|
+
self.server_manager.initialized_servers[server_name] = True
|
|
57
|
+
|
|
58
|
+
server_tools = self.server_manager._server_tools.get(server_name, [])
|
|
59
|
+
num_tools = len(server_tools)
|
|
60
|
+
|
|
61
|
+
return f"Connected to MCP server '{server_name}'. {num_tools} tools are now available."
|
|
62
|
+
|
|
63
|
+
except Exception as e:
|
|
64
|
+
logger.error(f"Error connecting to server '{server_name}': {e}")
|
|
65
|
+
return format_error(e, server_name=server_name)
|
|
66
|
+
|
|
67
|
+
def _run(self, server_name: str) -> str:
|
|
68
|
+
"""Synchronous version that raises a NotImplementedError - use _arun instead."""
|
|
69
|
+
raise NotImplementedError("ConnectServerTool requires async execution. Use _arun instead.")
|
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
from typing import ClassVar
|
|
2
|
+
|
|
3
|
+
from pydantic import BaseModel
|
|
4
|
+
|
|
5
|
+
from mcp_use.agents.managers.tools.base_tool import MCPServerTool
|
|
6
|
+
from mcp_use.errors.error_formatting import format_error
|
|
7
|
+
from mcp_use.logging import logger
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class DisconnectServerInput(BaseModel):
|
|
11
|
+
"""Empty input for disconnecting from the current server"""
|
|
12
|
+
|
|
13
|
+
pass
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class DisconnectServerTool(MCPServerTool):
|
|
17
|
+
"""Tool for disconnecting from the currently active MCP server."""
|
|
18
|
+
|
|
19
|
+
name: ClassVar[str] = "disconnect_from_mcp_server"
|
|
20
|
+
description: ClassVar[str] = "Disconnect from the currently active MCP (Model Context Protocol) server"
|
|
21
|
+
args_schema: ClassVar[type[BaseModel]] = DisconnectServerInput
|
|
22
|
+
|
|
23
|
+
def _run(self, **kwargs) -> str:
|
|
24
|
+
"""Disconnect from the currently active MCP server."""
|
|
25
|
+
if not self.server_manager.active_server:
|
|
26
|
+
return "No MCP server is currently active, so there's nothing to disconnect from."
|
|
27
|
+
|
|
28
|
+
server_name = self.server_manager.active_server
|
|
29
|
+
try:
|
|
30
|
+
# Clear the active server
|
|
31
|
+
self.server_manager.active_server = None
|
|
32
|
+
|
|
33
|
+
# Note: We're not actually closing the session here, just 'deactivating'
|
|
34
|
+
# This way we keep the session cache without requiring reconnection if needed again
|
|
35
|
+
|
|
36
|
+
return f"Successfully disconnected from MCP server '{server_name}'."
|
|
37
|
+
except Exception as e:
|
|
38
|
+
logger.error(f"Error disconnecting from server '{server_name}': {e}")
|
|
39
|
+
return format_error(e, server_name=server_name)
|
|
40
|
+
|
|
41
|
+
async def _arun(self, **kwargs) -> str:
|
|
42
|
+
"""Async implementation of _run."""
|
|
43
|
+
return self._run(**kwargs)
|
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
from typing import ClassVar
|
|
2
|
+
|
|
3
|
+
from pydantic import BaseModel
|
|
4
|
+
|
|
5
|
+
from mcp_use.agents.managers.tools.base_tool import MCPServerTool
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
class CurrentServerInput(BaseModel):
|
|
9
|
+
"""Empty input for checking current server"""
|
|
10
|
+
|
|
11
|
+
pass
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class GetActiveServerTool(MCPServerTool):
|
|
15
|
+
"""Tool for getting the currently active MCP server."""
|
|
16
|
+
|
|
17
|
+
name: ClassVar[str] = "get_active_mcp_server"
|
|
18
|
+
description: ClassVar[str] = "Get the currently active MCP (Model Context Protocol) server"
|
|
19
|
+
args_schema: ClassVar[type[BaseModel]] = CurrentServerInput
|
|
20
|
+
|
|
21
|
+
def _run(self, **kwargs) -> str:
|
|
22
|
+
"""Get the currently active MCP server."""
|
|
23
|
+
if not self.server_manager.active_server:
|
|
24
|
+
return "No MCP server is currently active. Use connect_to_mcp_server to connect to a server."
|
|
25
|
+
return f"Currently active MCP server: {self.server_manager.active_server}"
|
|
26
|
+
|
|
27
|
+
async def _arun(self, **kwargs) -> str:
|
|
28
|
+
"""Async implementation of _run."""
|
|
29
|
+
return self._run(**kwargs)
|
|
@@ -0,0 +1,53 @@
|
|
|
1
|
+
from typing import ClassVar
|
|
2
|
+
|
|
3
|
+
from pydantic import BaseModel
|
|
4
|
+
|
|
5
|
+
from mcp_use.agents.managers.tools.base_tool import MCPServerTool
|
|
6
|
+
from mcp_use.errors.error_formatting import format_error
|
|
7
|
+
from mcp_use.logging import logger
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class listServersInput(BaseModel):
|
|
11
|
+
"""Empty input for listing available servers"""
|
|
12
|
+
|
|
13
|
+
pass
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class ListServersTool(MCPServerTool):
|
|
17
|
+
"""Tool for listing available MCP servers."""
|
|
18
|
+
|
|
19
|
+
name: ClassVar[str] = "list_mcp_servers"
|
|
20
|
+
description: ClassVar[str] = (
|
|
21
|
+
"Lists all available MCP (Model Context Protocol) servers that can be "
|
|
22
|
+
"connected to, along with the tools available on each server. "
|
|
23
|
+
"Use this tool to discover servers and see what functionalities they offer."
|
|
24
|
+
)
|
|
25
|
+
args_schema: ClassVar[type[BaseModel]] = listServersInput
|
|
26
|
+
|
|
27
|
+
def _run(self, **kwargs) -> str:
|
|
28
|
+
"""List all available MCP servers along with their available tools."""
|
|
29
|
+
servers = self.server_manager.client.get_server_names()
|
|
30
|
+
if not servers:
|
|
31
|
+
return "No MCP servers are currently defined."
|
|
32
|
+
|
|
33
|
+
result = "Available MCP servers:\n"
|
|
34
|
+
for i, server_name in enumerate(servers):
|
|
35
|
+
active_marker = " (ACTIVE)" if server_name == self.server_manager.active_server else ""
|
|
36
|
+
result += f"{i + 1}. {server_name}{active_marker}\n"
|
|
37
|
+
|
|
38
|
+
tools: list = []
|
|
39
|
+
try:
|
|
40
|
+
# Check cache first
|
|
41
|
+
if server_name in self.server_manager._server_tools:
|
|
42
|
+
tools = self.server_manager._server_tools[server_name]
|
|
43
|
+
tool_count = len(tools)
|
|
44
|
+
result += f" {tool_count} tools available for this server\n"
|
|
45
|
+
except Exception as e:
|
|
46
|
+
logger.error(f"Unexpected error listing tools for server '{server_name}': {e}")
|
|
47
|
+
return format_error(e, server=server_name, operation="list_tools")
|
|
48
|
+
|
|
49
|
+
return result
|
|
50
|
+
|
|
51
|
+
async def _arun(self, **kwargs) -> str:
|
|
52
|
+
"""Async implementation of _run - calls the synchronous version."""
|
|
53
|
+
return self._run(**kwargs)
|
|
@@ -0,0 +1,328 @@
|
|
|
1
|
+
import asyncio
|
|
2
|
+
import math
|
|
3
|
+
import time
|
|
4
|
+
from typing import ClassVar
|
|
5
|
+
|
|
6
|
+
from langchain_core.tools import BaseTool
|
|
7
|
+
from pydantic import BaseModel, Field
|
|
8
|
+
|
|
9
|
+
from mcp_use.agents.managers.tools.base_tool import MCPServerTool
|
|
10
|
+
from mcp_use.logging import logger
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class ToolSearchInput(BaseModel):
|
|
14
|
+
"""Input for searching for tools across MCP servers"""
|
|
15
|
+
|
|
16
|
+
query: str = Field(description="The search query to find relevant tools")
|
|
17
|
+
top_k: int = Field(
|
|
18
|
+
default=100,
|
|
19
|
+
description="The maximum number of tools to return (defaults to 100)",
|
|
20
|
+
)
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class SearchToolsTool(MCPServerTool):
|
|
24
|
+
"""Tool for searching for tools across all MCP servers using semantic search."""
|
|
25
|
+
|
|
26
|
+
name: ClassVar[str] = "search_mcp_tools"
|
|
27
|
+
description: ClassVar[str] = (
|
|
28
|
+
"Search for relevant tools across all MCP servers using semantic search. "
|
|
29
|
+
"Provide a description of the tool you think you might need to be able to perform "
|
|
30
|
+
"the task you are assigned. Do not be too specific, the search will give you many "
|
|
31
|
+
"options. It is important you search for the tool, not for the goal. "
|
|
32
|
+
"If your first search doesn't yield relevant results, try using different keywords "
|
|
33
|
+
"or more general terms."
|
|
34
|
+
)
|
|
35
|
+
args_schema: ClassVar[type[BaseModel]] = ToolSearchInput
|
|
36
|
+
|
|
37
|
+
def __init__(self, server_manager):
|
|
38
|
+
"""Initialize with server manager and create a search tool."""
|
|
39
|
+
super().__init__(server_manager)
|
|
40
|
+
self._search_tool = ToolSearchEngine(server_manager=server_manager)
|
|
41
|
+
|
|
42
|
+
async def _arun(self, query: str, top_k: int = 100) -> str:
|
|
43
|
+
"""Search for tools across all MCP servers using semantic search."""
|
|
44
|
+
# Make sure the index is ready, and if not, allow the search_tools method to handle it
|
|
45
|
+
# No need to manually check or build the index here as the search_tools method will do that
|
|
46
|
+
|
|
47
|
+
# Perform search using our search tool instance
|
|
48
|
+
results = await self._search_tool.search_tools(
|
|
49
|
+
query, top_k=top_k, active_server=self.server_manager.active_server
|
|
50
|
+
)
|
|
51
|
+
return results
|
|
52
|
+
|
|
53
|
+
def _run(self, query: str, top_k: int = 100) -> str:
|
|
54
|
+
"""Synchronous version that raises a NotImplementedError - use _arun instead."""
|
|
55
|
+
raise NotImplementedError("SearchToolsTool requires async execution. Use _arun instead.")
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
class ToolSearchEngine:
|
|
59
|
+
"""
|
|
60
|
+
Provides semantic search capabilities for MCP tools.
|
|
61
|
+
Uses vector similarity for semantic search with optional result caching.
|
|
62
|
+
"""
|
|
63
|
+
|
|
64
|
+
def __init__(self, server_manager=None, use_caching: bool = True):
|
|
65
|
+
"""
|
|
66
|
+
Initialize the tool search engine.
|
|
67
|
+
|
|
68
|
+
Args:
|
|
69
|
+
server_manager: The ServerManager instance to get tools from
|
|
70
|
+
use_caching: Whether to cache query results
|
|
71
|
+
"""
|
|
72
|
+
self.server_manager = server_manager
|
|
73
|
+
self.use_caching = use_caching
|
|
74
|
+
self.is_indexed = False
|
|
75
|
+
|
|
76
|
+
# Initialize model components (loaded on demand)
|
|
77
|
+
self.model = None
|
|
78
|
+
self.embedding_function = None
|
|
79
|
+
|
|
80
|
+
# Data storage
|
|
81
|
+
self.tool_embeddings = {} # Maps tool name to embedding vector
|
|
82
|
+
self.tools_by_name = {} # Maps tool name to tool instance
|
|
83
|
+
self.server_by_tool = {} # Maps tool name to server name
|
|
84
|
+
self.tool_texts = {} # Maps tool name to searchable text
|
|
85
|
+
self.query_cache = {} # Caches search results by query
|
|
86
|
+
|
|
87
|
+
def _load_model(self) -> bool:
|
|
88
|
+
"""Load the embedding model for semantic search if not already loaded."""
|
|
89
|
+
if self.model is not None:
|
|
90
|
+
return True
|
|
91
|
+
|
|
92
|
+
try:
|
|
93
|
+
from fastembed import TextEmbedding # optional dependency install with [search]
|
|
94
|
+
except ImportError as exc:
|
|
95
|
+
logger.error(
|
|
96
|
+
"The 'fastembed' library is not installed. "
|
|
97
|
+
"To use the search functionality, please install it by running: "
|
|
98
|
+
"pip install mcp-use[search]"
|
|
99
|
+
)
|
|
100
|
+
raise ImportError(
|
|
101
|
+
"The 'fastembed' library is not installed. "
|
|
102
|
+
"To use the server_manager functionality, please install it by running: "
|
|
103
|
+
"pip install mcp-use[search] "
|
|
104
|
+
"or disable the server_manager by setting use_server_manager=False in the MCPAgent constructor."
|
|
105
|
+
) from exc
|
|
106
|
+
|
|
107
|
+
try:
|
|
108
|
+
self.model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
|
|
109
|
+
self.embedding_function = lambda texts: list(self.model.embed(texts))
|
|
110
|
+
return True
|
|
111
|
+
except Exception as e:
|
|
112
|
+
logger.error(f"Failed to load the embedding model: {e}")
|
|
113
|
+
return False
|
|
114
|
+
|
|
115
|
+
async def start_indexing(self) -> None:
|
|
116
|
+
"""Index the tools from the server manager."""
|
|
117
|
+
if not self.server_manager:
|
|
118
|
+
return
|
|
119
|
+
|
|
120
|
+
# Get tools from server manager
|
|
121
|
+
server_tools = self.server_manager._server_tools
|
|
122
|
+
|
|
123
|
+
if not server_tools:
|
|
124
|
+
# Try to prefetch tools first
|
|
125
|
+
if hasattr(self.server_manager, "_prefetch_server_tools"):
|
|
126
|
+
await self.server_manager._prefetch_server_tools()
|
|
127
|
+
server_tools = self.server_manager._server_tools
|
|
128
|
+
|
|
129
|
+
if server_tools:
|
|
130
|
+
await self.index_tools(server_tools)
|
|
131
|
+
|
|
132
|
+
async def index_tools(self, server_tools: dict[str, list[BaseTool]]) -> None:
|
|
133
|
+
"""
|
|
134
|
+
Index all tools from all servers for search.
|
|
135
|
+
|
|
136
|
+
Args:
|
|
137
|
+
server_tools: dictionary mapping server names to their tools
|
|
138
|
+
"""
|
|
139
|
+
# Clear previous indexes
|
|
140
|
+
self.tool_embeddings = {}
|
|
141
|
+
self.tools_by_name = {}
|
|
142
|
+
self.server_by_tool = {}
|
|
143
|
+
self.tool_texts = {}
|
|
144
|
+
self.query_cache = {}
|
|
145
|
+
self.is_indexed = False
|
|
146
|
+
|
|
147
|
+
# Collect all tools and their descriptions
|
|
148
|
+
for server_name, tools in server_tools.items():
|
|
149
|
+
for tool in tools:
|
|
150
|
+
# Create text representation for search
|
|
151
|
+
tool_text = f"{tool.name}: {tool.description}"
|
|
152
|
+
|
|
153
|
+
# Store tool information
|
|
154
|
+
self.tools_by_name[tool.name] = tool
|
|
155
|
+
self.server_by_tool[tool.name] = server_name
|
|
156
|
+
self.tool_texts[tool.name] = tool_text.lower() # For case-insensitive search
|
|
157
|
+
|
|
158
|
+
if not self.tool_texts:
|
|
159
|
+
return
|
|
160
|
+
|
|
161
|
+
# Generate embeddings
|
|
162
|
+
if self._load_model():
|
|
163
|
+
tool_names = list(self.tool_texts.keys())
|
|
164
|
+
tool_texts = [self.tool_texts[name] for name in tool_names]
|
|
165
|
+
|
|
166
|
+
try:
|
|
167
|
+
embeddings = self.embedding_function(tool_texts)
|
|
168
|
+
for name, embedding in zip(tool_names, embeddings, strict=True):
|
|
169
|
+
self.tool_embeddings[name] = embedding
|
|
170
|
+
|
|
171
|
+
# Mark as indexed if we successfully embedded tools
|
|
172
|
+
self.is_indexed = len(self.tool_embeddings) > 0
|
|
173
|
+
except Exception:
|
|
174
|
+
return
|
|
175
|
+
|
|
176
|
+
def search(self, query: str, top_k: int = 5) -> list[tuple[BaseTool, str, float]]:
|
|
177
|
+
"""
|
|
178
|
+
Search for tools that match the query using semantic search.
|
|
179
|
+
|
|
180
|
+
Args:
|
|
181
|
+
query: The search query
|
|
182
|
+
top_k: Number of top results to return
|
|
183
|
+
|
|
184
|
+
Returns:
|
|
185
|
+
list of tuples containing (tool, server_name, score)
|
|
186
|
+
"""
|
|
187
|
+
if not self.is_indexed:
|
|
188
|
+
return []
|
|
189
|
+
|
|
190
|
+
# Check cache first
|
|
191
|
+
cache_key = f"semantic:{query}:{top_k}"
|
|
192
|
+
if self.use_caching and cache_key in self.query_cache:
|
|
193
|
+
return self.query_cache[cache_key]
|
|
194
|
+
|
|
195
|
+
# Ensure model and embeddings exist
|
|
196
|
+
if not self._load_model() or not self.tool_embeddings:
|
|
197
|
+
return []
|
|
198
|
+
|
|
199
|
+
# Generate embedding for the query
|
|
200
|
+
try:
|
|
201
|
+
query_embedding = self.embedding_function([query])[0]
|
|
202
|
+
except Exception:
|
|
203
|
+
return []
|
|
204
|
+
|
|
205
|
+
# Calculate similarity scores
|
|
206
|
+
scores = {}
|
|
207
|
+
for tool_name, embedding in self.tool_embeddings.items():
|
|
208
|
+
# Calculate cosine similarity using pure Python
|
|
209
|
+
similarity = self._cosine_similarity(query_embedding, embedding)
|
|
210
|
+
scores[tool_name] = float(similarity)
|
|
211
|
+
|
|
212
|
+
# Sort by score and get top_k results
|
|
213
|
+
sorted_results = sorted(scores.items(), key=lambda x: x[1], reverse=True)[:top_k]
|
|
214
|
+
|
|
215
|
+
# Format results
|
|
216
|
+
results = []
|
|
217
|
+
for tool_name, score in sorted_results:
|
|
218
|
+
tool = self.tools_by_name.get(tool_name)
|
|
219
|
+
server_name = self.server_by_tool.get(tool_name)
|
|
220
|
+
if tool and server_name:
|
|
221
|
+
results.append((tool, server_name, score))
|
|
222
|
+
|
|
223
|
+
# Cache results
|
|
224
|
+
if self.use_caching:
|
|
225
|
+
self.query_cache[cache_key] = results
|
|
226
|
+
|
|
227
|
+
return results
|
|
228
|
+
|
|
229
|
+
async def search_tools(self, query: str, top_k: int = 100, active_server: str = None) -> str:
|
|
230
|
+
"""
|
|
231
|
+
Search for tools across all MCP servers using semantic search.
|
|
232
|
+
|
|
233
|
+
Args:
|
|
234
|
+
query: The search query to find relevant tools
|
|
235
|
+
top_k: Number of top results to return
|
|
236
|
+
active_server: Name of the currently active server (for highlighting)
|
|
237
|
+
|
|
238
|
+
Returns:
|
|
239
|
+
String with formatted search results
|
|
240
|
+
"""
|
|
241
|
+
# Ensure the index is built or build it
|
|
242
|
+
if not self.is_indexed:
|
|
243
|
+
# Try to build the index
|
|
244
|
+
if self.server_manager and self.server_manager._server_tools:
|
|
245
|
+
await self.index_tools(self.server_manager._server_tools)
|
|
246
|
+
else:
|
|
247
|
+
# If we don't have server_manager or tools, try to index directly
|
|
248
|
+
await self.start_indexing()
|
|
249
|
+
|
|
250
|
+
# Wait for indexing to complete (maximum 10 seconds)
|
|
251
|
+
start_time = time.time()
|
|
252
|
+
timeout = 10 # seconds
|
|
253
|
+
while not self.is_indexed and (time.time() - start_time) < timeout:
|
|
254
|
+
await asyncio.sleep(0.5)
|
|
255
|
+
|
|
256
|
+
# If still not indexed, return a friendly message
|
|
257
|
+
if not self.is_indexed:
|
|
258
|
+
return (
|
|
259
|
+
"I'm still preparing the tool index. Please try your search again in a moment. "
|
|
260
|
+
"This usually takes just a few seconds to complete."
|
|
261
|
+
)
|
|
262
|
+
|
|
263
|
+
# If the server manager has an active server but it wasn't provided, use it
|
|
264
|
+
if active_server is None and self.server_manager and hasattr(self.server_manager, "active_server"):
|
|
265
|
+
active_server = self.server_manager.active_server
|
|
266
|
+
|
|
267
|
+
results = self.search(query, top_k=top_k)
|
|
268
|
+
if not results:
|
|
269
|
+
return (
|
|
270
|
+
"No relevant tools found. The search provided no results. "
|
|
271
|
+
"You can try searching again with different keywords. "
|
|
272
|
+
"Try using more general terms or focusing on the capability you need."
|
|
273
|
+
)
|
|
274
|
+
|
|
275
|
+
# If there's an active server, mark it in the results
|
|
276
|
+
if active_server:
|
|
277
|
+
# Create a new results list with marked active server
|
|
278
|
+
marked_results = []
|
|
279
|
+
for tool, server_name, score in results:
|
|
280
|
+
# If this is the active server, add "(ACTIVE)" marker
|
|
281
|
+
display_server = f"{server_name} (ACTIVE)" if server_name == active_server else server_name
|
|
282
|
+
marked_results.append((tool, display_server, score))
|
|
283
|
+
results = marked_results
|
|
284
|
+
|
|
285
|
+
# Format and return the results
|
|
286
|
+
return self._format_search_results(results)
|
|
287
|
+
|
|
288
|
+
def _format_search_results(self, results: list[tuple[BaseTool, str, float]]) -> str:
|
|
289
|
+
"""Format search results in a consistent format."""
|
|
290
|
+
formatted_output = "Search results\n\n"
|
|
291
|
+
|
|
292
|
+
for i, (tool, server_name, score) in enumerate(results):
|
|
293
|
+
# Format score as percentage
|
|
294
|
+
score_pct = f"{score * 100:.1f}%"
|
|
295
|
+
if i < 5:
|
|
296
|
+
logger.info(f"{i}: {tool.name} ({score_pct} match)")
|
|
297
|
+
formatted_output += f"[{i + 1}] Tool: {tool.name} ({score_pct} match)\n"
|
|
298
|
+
formatted_output += f" Server: {server_name}\n"
|
|
299
|
+
formatted_output += f" Description: {tool.description}\n\n"
|
|
300
|
+
|
|
301
|
+
# Add footer with information about how to use the results
|
|
302
|
+
formatted_output += "\nTo use a tool, connect to the appropriate server first, then invoke the tool."
|
|
303
|
+
|
|
304
|
+
return formatted_output
|
|
305
|
+
|
|
306
|
+
def _cosine_similarity(self, vec1: list[float], vec2: list[float]) -> float:
|
|
307
|
+
"""Calculate cosine similarity between two vectors.
|
|
308
|
+
|
|
309
|
+
Args:
|
|
310
|
+
vec1: First vector
|
|
311
|
+
vec2: Second vector
|
|
312
|
+
|
|
313
|
+
Returns:
|
|
314
|
+
Cosine similarity between the vectors
|
|
315
|
+
"""
|
|
316
|
+
# Calculate dot product
|
|
317
|
+
dot_product = sum(a * b for a, b in zip(vec1, vec2, strict=False))
|
|
318
|
+
|
|
319
|
+
# Calculate magnitudes
|
|
320
|
+
magnitude1 = math.sqrt(sum(a * a for a in vec1))
|
|
321
|
+
magnitude2 = math.sqrt(sum(b * b for b in vec2))
|
|
322
|
+
|
|
323
|
+
# Avoid division by zero
|
|
324
|
+
if magnitude1 == 0 or magnitude2 == 0:
|
|
325
|
+
return 0.0
|
|
326
|
+
|
|
327
|
+
# Calculate cosine similarity
|
|
328
|
+
return dot_product / (magnitude1 * magnitude2)
|