Move mcp conn to common

Signed-off-by: Jin Hai <haijin.chn@gmail.com>

agent/component/agent_with_tools.py

@@ -30,7 +30,7 @@ from api.db.services.mcp_server_service import MCPServerService
 from common.connection_utils import timeout
 from rag.prompts.generator import next_step, COMPLETE_TASK, analyze_task, \
     citation_prompt, reflect, rank_memories, kb_prompt, citation_plus, full_question, message_fit_in
-from rag.utils.mcp_tool_call_conn import MCPToolCallSession, mcp_tool_metadata_to_openai_tool
+from common.mcp_tool_call_conn import MCPToolCallSession, mcp_tool_metadata_to_openai_tool
 from agent.component.llm import LLMParam, LLM
 
 

agent/tools/base.py

@@ -21,9 +21,8 @@ from functools import partial
 from typing import TypedDict, List, Any
 from agent.component.base import ComponentParamBase, ComponentBase
 from common.misc_utils import hash_str2int
-from rag.llm.chat_model import ToolCallSession
 from rag.prompts.generator import kb_prompt
-from rag.utils.mcp_tool_call_conn import MCPToolCallSession
+from common.mcp_tool_call_conn import MCPToolCallSession, ToolCallSession
 from timeit import default_timer as timer
 
 

api/apps/mcp_server_app.py

@@ -25,7 +25,7 @@ from common.misc_utils import get_uuid
 from api.utils.api_utils import get_data_error_result, get_json_result, server_error_response, validate_request, \
     get_mcp_tools
 from api.utils.web_utils import get_float, safe_json_parse
-from rag.utils.mcp_tool_call_conn import MCPToolCallSession, close_multiple_mcp_toolcall_sessions
+from common.mcp_tool_call_conn import MCPToolCallSession, close_multiple_mcp_toolcall_sessions
 
 
 @manager.route("/list", methods=["POST"])  # noqa: F821

api/ragflow_server.py

@@ -41,7 +41,7 @@ from api.db.db_models import init_database_tables as init_web_db
 from api.db.init_data import init_web_data
 from common.versions import get_ragflow_version
 from common.config_utils import show_configs
-from rag.utils.mcp_tool_call_conn import shutdown_all_mcp_sessions
+from common.mcp_tool_call_conn import shutdown_all_mcp_sessions
 from rag.utils.redis_conn import RedisDistributedLock
 
 stop_event = threading.Event()

api/utils/api_utils.py

@@ -37,7 +37,7 @@ from peewee import OperationalError
 from common.constants import ActiveEnum
 from api.db.db_models import APIToken
 from api.utils.json_encode import CustomJSONEncoder
-from rag.utils.mcp_tool_call_conn import MCPToolCallSession, close_multiple_mcp_toolcall_sessions
+from common.mcp_tool_call_conn import MCPToolCallSession, close_multiple_mcp_toolcall_sessions
 from api.db.services.tenant_llm_service import LLMFactoriesService
 from common.connection_utils import timeout
 from common.constants import RetCode

common/mcp_tool_call_conn.py

@@ -21,7 +21,7 @@ import weakref
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures import TimeoutError as FuturesTimeoutError
 from string import Template
-from typing import Any, Literal
+from typing import Any, Literal, Protocol
 
 from typing_extensions import override
 
@@ -30,12 +30,15 @@ from mcp.client.session import ClientSession
 from mcp.client.sse import sse_client
 from mcp.client.streamable_http import streamablehttp_client
 from mcp.types import CallToolResult, ListToolsResult, TextContent, Tool
-from rag.llm.chat_model import ToolCallSession
 
 MCPTaskType = Literal["list_tools", "tool_call"]
 MCPTask = tuple[MCPTaskType, dict[str, Any], asyncio.Queue[Any]]
 
 
+class ToolCallSession(Protocol):
+    def tool_call(self, name: str, arguments: dict[str, Any]) -> str: ...
+
+
 class MCPToolCallSession(ToolCallSession):
     _ALL_INSTANCES: weakref.WeakSet["MCPToolCallSession"] = weakref.WeakSet()
 
@@ -106,7 +109,8 @@ class MCPToolCallSession(ToolCallSession):
                 await self._process_mcp_tasks(None, msg)
 
         else:
-            await self._process_mcp_tasks(None, f"Unsupported MCP server type: {self._mcp_server.server_type}, id: {self._mcp_server.id}")
+            await self._process_mcp_tasks(None,
+                                          f"Unsupported MCP server type: {self._mcp_server.server_type}, id: {self._mcp_server.id}")
 
     async def _process_mcp_tasks(self, client_session: ClientSession | None, error_message: str | None = None) -> None:
         while not self._close:
@@ -164,7 +168,8 @@ class MCPToolCallSession(ToolCallSession):
             raise
 
     async def _call_mcp_tool(self, name: str, arguments: dict[str, Any], timeout: float | int = 10) -> str:
-        result: CallToolResult = await self._call_mcp_server("tool_call", name=name, arguments=arguments, timeout=timeout)
+        result: CallToolResult = await self._call_mcp_server("tool_call", name=name, arguments=arguments,
+                                                             timeout=timeout)
 
         if result.isError:
             return f"MCP server error: {result.content}"
@@ -283,7 +288,8 @@ def close_multiple_mcp_toolcall_sessions(sessions: list[MCPToolCallSession]) ->
     except Exception:
         logging.exception("Exception during MCP session cleanup thread management")
 
-    logging.info(f"{len(sessions)} MCP sessions has been cleaned up. {len(list(MCPToolCallSession._ALL_INSTANCES))} in global context.")
+    logging.info(
+        f"{len(sessions)} MCP sessions has been cleaned up. {len(list(MCPToolCallSession._ALL_INSTANCES))} in global context.")
 
 
 def shutdown_all_mcp_sessions():
@@ -298,7 +304,7 @@ def shutdown_all_mcp_sessions():
     logging.info("All MCPToolCallSession instances have been closed.")
 
 
-def mcp_tool_metadata_to_openai_tool(mcp_tool: Tool|dict) -> dict[str, Any]:
+def mcp_tool_metadata_to_openai_tool(mcp_tool: Tool | dict) -> dict[str, Any]:
     if isinstance(mcp_tool, dict):
         return {
             "type": "function",

rag/llm/chat_model.py

@@ -22,7 +22,6 @@ import re
 import time
 from abc import ABC
 from copy import deepcopy
-from typing import Any, Protocol
 from urllib.parse import urljoin
 
 import json_repair
@@ -65,10 +64,6 @@ LENGTH_NOTIFICATION_CN = "······\n由于大模型的上下文窗口大小
 LENGTH_NOTIFICATION_EN = "...\nThe answer is truncated by your chosen LLM due to its limitation on context length."
 
 
-class ToolCallSession(Protocol):
-    def tool_call(self, name: str, arguments: dict[str, Any]) -> str: ...
-
-
 class Base(ABC):
     def __init__(self, key, model_name, base_url, **kwargs):
         timeout = int(os.environ.get("LM_TIMEOUT_SECONDS", 600))