diff --git a/mcp/mcp_settings.json b/mcp/mcp_settings.json index 30b39af..9677e1e 100644 --- a/mcp/mcp_settings.json +++ b/mcp/mcp_settings.json @@ -14,6 +14,13 @@ "./mcp/multi_keyword_search_server.py", "{dataset_dir}" ] + }, + "rag_retrieve": { + "command": "python", + "args": [ + "./mcp/rag_retrieve_server.py", + "https://api-dev.gptbase.ai/v1/rag_retrieve/{unique_id}" + ] } } } diff --git a/mcp/rag_retrieve_server.py b/mcp/rag_retrieve_server.py new file mode 100644 index 0000000..e8ce699 --- /dev/null +++ b/mcp/rag_retrieve_server.py @@ -0,0 +1,198 @@ +#!/usr/bin/env python3 +""" +RAG检索MCP服务器 +调用本地RAG API进行文档检索 +""" + +import asyncio +import json +import sys +import os +from typing import Any, Dict, List + +try: + import requests +except ImportError: + print("Error: requests module is required. Please install it with: pip install requests") + sys.exit(1) + +from mcp_common import ( + create_error_response, + create_success_response, + create_initialize_response, + create_ping_response, + create_tools_list_response, + load_tools_from_json, + handle_mcp_streaming +) + + +def rag_retrieve(query: str, top_k: int = 50) -> Dict[str, Any]: + """调用RAG检索API""" + try: + url = "" + if len(sys.argv) > 1: + url = sys.argv[1] + + if not url: + return { + "content": [ + { + "type": "text", + "text": "Error: RAG API URL not provided. Please provide URL as command line argument." + } + ] + } + + headers = { + "content-type": "application/json" + } + data = { + "query": query, + "top_k": top_k + } + + # 发送POST请求 + response = requests.post(url, json=data, headers=headers, timeout=30) + + if response.status_code != 200: + return { + "content": [ + { + "type": "text", + "text": f"Error: RAG API returned status code {response.status_code}. Response: {response.text}" + } + ] + } + + # 解析响应 + try: + response_data = response.json() + except json.JSONDecodeError as e: + return { + "content": [ + { + "type": "text", + "text": f"Error: Failed to parse API response as JSON. Error: {str(e)}, Raw response: {response.text}" + } + ] + } + + # 提取markdown字段 + if "markdown" in response_data: + markdown_content = response_data["markdown"] + return { + "content": [ + { + "type": "text", + "text": markdown_content + } + ] + } + else: + return { + "content": [ + { + "type": "text", + "text": f"Error: 'markdown' field not found in API response. Response: {json.dumps(response_data, indent=2, ensure_ascii=False)}" + } + ] + } + + except requests.exceptions.RequestException as e: + return { + "content": [ + { + "type": "text", + "text": f"Error: Failed to connect to RAG API. {str(e)}" + } + ] + } + except Exception as e: + return { + "content": [ + { + "type": "text", + "text": f"Error: {str(e)}" + } + ] + } + + +async def handle_request(request: Dict[str, Any]) -> Dict[str, Any]: + """Handle MCP request""" + try: + method = request.get("method") + params = request.get("params", {}) + request_id = request.get("id") + + if method == "initialize": + return create_initialize_response(request_id, "rag-retrieve") + + elif method == "ping": + return create_ping_response(request_id) + + elif method == "tools/list": + # 从 JSON 文件加载工具定义 + tools = load_tools_from_json("rag_retrieve_tools.json") + if not tools: + # 如果 JSON 文件不存在，使用默认定义 + tools = [ + { + "name": "rag_retrieve", + "description": "调用RAG检索API，根据查询内容检索相关文档。返回包含相关内容的markdown格式结果。", + "inputSchema": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "检索查询内容" + }, + "top_k": { + "type": "integer", + "description": "返回结果的最大数量，默认50", + "default": 50 + } + }, + "required": ["query"] + } + } + ] + return create_tools_list_response(request_id, tools) + + elif method == "tools/call": + tool_name = params.get("name") + arguments = params.get("arguments", {}) + + if tool_name == "rag_retrieve": + query = arguments.get("query", "") + top_k = arguments.get("top_k", 50) + + if not query: + return create_error_response(request_id, -32602, "Missing required parameter: query") + + result = rag_retrieve(query, top_k) + + return { + "jsonrpc": "2.0", + "id": request_id, + "result": result + } + + else: + return create_error_response(request_id, -32601, f"Unknown tool: {tool_name}") + + else: + return create_error_response(request_id, -32601, f"Unknown method: {method}") + + except Exception as e: + return create_error_response(request.get("id"), -32603, f"Internal error: {str(e)}") + + +async def main(): + """Main entry point.""" + await handle_mcp_streaming(handle_request) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/mcp/rag_retrieve_tools.json b/mcp/rag_retrieve_tools.json new file mode 100644 index 0000000..3f09546 --- /dev/null +++ b/mcp/rag_retrieve_tools.json @@ -0,0 +1,21 @@ +[ + { + "name": "rag_retrieve", + "description": "调用RAG检索API，根据查询内容检索相关文档。返回包含相关内容的markdown格式结果。", + "inputSchema": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "检索查询内容" + }, + "top_k": { + "type": "integer", + "description": "返回结果的最大数量，默认50", + "default": 50 + } + }, + "required": ["query"] + } + } +] \ No newline at end of file diff --git a/utils/file_loaded_agent_manager.py b/utils/file_loaded_agent_manager.py index 7626dfd..52f377a 100644 --- a/utils/file_loaded_agent_manager.py +++ b/utils/file_loaded_agent_manager.py @@ -103,8 +103,8 @@ class FileLoadedAgentManager: import os # 实现参数优先级逻辑：传入参数 > 项目配置 > 默认配置 - final_system_prompt = load_system_prompt(project_dir, language, system_prompt, robot_type) - final_mcp_settings = load_mcp_settings(project_dir, mcp_settings) + final_system_prompt = load_system_prompt(project_dir, language, system_prompt, robot_type, unique_id) + final_mcp_settings = load_mcp_settings(project_dir, mcp_settings, unique_id) cache_key = self._get_cache_key(unique_id) diff --git a/utils/prompt_loader.py b/utils/prompt_loader.py index 10c279d..1565f72 100644 --- a/utils/prompt_loader.py +++ b/utils/prompt_loader.py @@ -7,9 +7,18 @@ import json from typing import List, Dict, Optional -def load_system_prompt(project_dir: str, language: str = None, system_prompt: str=None, robot_type: str = "agent") -> str: +def load_system_prompt(project_dir: str, language: str = None, system_prompt: str=None, robot_type: str = "agent", unique_id: str="") -> str: + # 获取语言显示名称 + language_display_map = { + 'zh': '中文', + 'en': 'English', + 'ja': '日本語', + 'jp': '日本語' + } + language_display = language_display_map.get(language, language if language else 'English') + if robot_type == "agent": - return system_prompt or "" + return system_prompt.replace("{language}", language_display).replace('{unique_id}', unique_id) or "" if robot_type == "catalog_agent": """ 优先使用项目目录的system_prompt，没有才使用默认的system_prompt_default.md @@ -52,18 +61,9 @@ def load_system_prompt(project_dir: str, language: str = None, system_prompt: st with open(readme_path, "r", encoding="utf-8") as f: readme = f.read().strip() - # 获取语言显示名称 - language_display_map = { - 'zh': '中文', - 'en': 'English', - 'ja': '日本語', - 'jp': '日本語' - } - language_display = language_display_map.get(language, language if language else 'English') - - return system_prompt_default.replace("{readme}", str(readme)).replace("{language}", language_display).replace("{extra_prompt}", system_prompt or "") or "" + return system_prompt_default.replace("{readme}", str(readme)).replace("{language}", language_display).replace("{extra_prompt}", system_prompt or "").replace('{unique_id}', unique_id) or "" else: - return system_prompt or "" + return system_prompt.replace("{language}", language_display).replace('{unique_id}', unique_id) or "" def get_available_prompt_languages() -> list: """ @@ -85,7 +85,7 @@ def get_available_prompt_languages() -> list: return available_languages -def replace_mcp_placeholders(mcp_settings: List[Dict], dataset_dir: str) -> List[Dict]: +def replace_mcp_placeholders(mcp_settings: List[Dict], dataset_dir: str, unique_id: str) -> List[Dict]: """ 替换 MCP 配置中的占位符 """ @@ -98,21 +98,21 @@ def replace_mcp_placeholders(mcp_settings: List[Dict], dataset_dir: str) -> List for key, value in obj.items(): if key == 'args' and isinstance(value, list): # 特别处理 args 列表 - obj[key] = [item.replace('{dataset_dir}', dataset_dir) if isinstance(item, str) else item + obj[key] = [item.replace('{dataset_dir}', dataset_dir).replace('{unique_id}', unique_id) if isinstance(item, str) else item for item in value] elif isinstance(value, (dict, list)): obj[key] = replace_placeholders_in_obj(value) elif isinstance(value, str): - obj[key] = value.replace('{dataset_dir}', dataset_dir) + obj[key] = value.replace('{dataset_dir}', dataset_dir).replace('{unique_id}', unique_id) elif isinstance(obj, list): return [replace_placeholders_in_obj(item) if isinstance(item, (dict, list)) else - item.replace('{dataset_dir}', dataset_dir) if isinstance(item, str) else item + item.replace('{dataset_dir}', dataset_dir).replace('{unique_id}', unique_id) if isinstance(item, str) else item for item in obj] return obj return replace_placeholders_in_obj(mcp_settings) -def load_mcp_settings(project_dir: str, mcp_settings: list=None) -> List[Dict]: +def load_mcp_settings(project_dir: str, mcp_settings: list=None, unique_id: str="") -> List[Dict]: """ 优先使用项目目录的mcp_settings.json，没有才使用默认的mcp/mcp_settings.json @@ -164,7 +164,7 @@ def load_mcp_settings(project_dir: str, mcp_settings: list=None) -> List[Dict]: # 计算 dataset_dir 用于替换 MCP 配置中的占位符 dataset_dir = os.path.join(project_dir, "dataset") # 替换 MCP 配置中的 {dataset_dir} 占位符 - mcp_settings = replace_mcp_placeholders(mcp_settings, dataset_dir) + mcp_settings = replace_mcp_placeholders(mcp_settings, dataset_dir, unique_id) return mcp_settings