From 7b538d496734c5d4e21c882896985a67103be4f3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9C=B1=E6=BD=AE?= Date: Tue, 7 Oct 2025 14:35:07 +0800 Subject: [PATCH] add extra_prompt --- .gitignore | 1 + __pycache__/fastapi_app.cpython-312.pyc | Bin 13108 -> 13087 bytes __pycache__/gbase_agent.cpython-312.pyc | Bin 7991 -> 6870 bytes agent_prompt.txt | 8 +- fastapi_app.py | 4 +- gbase_agent.py | 59 ++--- llm_config.json | 65 ----- mcp/directory_tree_wrapper_server.py | 61 ----- mcp/mcp_wrapper.py | 308 ------------------------ mcp/ripgrep_wrapper_server.py | 61 ----- 10 files changed, 30 insertions(+), 537 deletions(-) delete mode 100644 llm_config.json delete mode 100644 mcp/directory_tree_wrapper_server.py delete mode 100644 mcp/mcp_wrapper.py delete mode 100644 mcp/ripgrep_wrapper_server.py diff --git a/.gitignore b/.gitignore index 156a1fd..dafdaaa 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ .DS_Store projects/* +workspace diff --git a/__pycache__/fastapi_app.cpython-312.pyc b/__pycache__/fastapi_app.cpython-312.pyc index d34da561400a0ea163462cd19d31e14588baf667..8479c3f1b8cd85a5a0dd74f913ca95bd34491209 100644 GIT binary patch delta 2085 zcmZvddrVVT9LMjuebZiS3vGe6a-jl!z{*1uopW@u2z9c_f;yus4fg_8Z1ME+m@A^& z@WI@M9+MG=Q7@&?AI5%2NHX549J}$eap2ixn+%Ncyerq8H zkQHmdljpTId4s#HK=tH;=_Y!Od|Wbgo2WjnR!f;Q0}Bd2n$Bpc9@fuT(uQt<;;zmw#c;2+6J;i>_MlJX=VFWpP_y= z?czvF=2~1S<)=NVAZ^h_Q~kJN))$O|NiYjpqwE*i=~x2o$(CGXm)%CT<(ZTT6e%KK z<+uV+6!MeMj6d05!Lp534xUJ81S=F-icrTs7Q4j_@meZ%}imtsyFjI2!Kc<09@6F)v~q!^?TkLRo2#B`k)w;}(c%I>Pa=Bw@S*L_aV@SkQ zSmyYc`h>o54K40Jn8&-BX)D8-cE3E$Hjzz*9o3K2j4pS?azb;OJIQ_CduCv4;0N1O zp!r&$`KEULE&q}!|ElZ$Rl~hgHs5udZ^FFaFr7iNTs~`*?P8}H#<2q!X9w8()EP+D zx(1Opju=zTpAQ`5=S4>bf{W4KKz@BSbEyc`yOo!U177M}HuCj8_2qIOwXfKXRCL9M z>dVzvO4VRjt}@fagc1@H*{HrsJz;^wggq1DJ{7PWL8MA%wJvB;Osou;gGR-qQAM3e ziz1lAOj>hH!3xEs2L(&DlU@zj6^h_uGE)>pn@ONJpgKx7XePO@IE0L!!u}lxoB*68 zGbLW3D_-678I#b=<8JkY!gkmBfUvDDoa5_?2?tY?R9vrj6ZzQ{?;d1|dw% zl$!34812DdL%bTW0svJm|9N4!B`IZBjNT0>&e=-xO!S_P0=}f5`@Q#^ zd;0FV=iYas{quI?q`{EGv8VRjtY=*8>TsgUIz4NnCa^f1IWf}3^`o;Hoz4 zN|L}k`F8l%VJ}ar##D2|gcEZ_A*ndRiTb3H(SSV0i^ijh4wB^HUcL=}%eN9Iu%H2T z)SV6YtVQ1gihYGDu1(ReK27?0_7d~h>*NVVSueQ+))(YEOrlw|h`C1<=dEMPIa`tp z+2m3fDhQAT_#be&z)buwRbZ{Oz5NzJs{hIrPB1)?MC+^>^svKG-MI+HqFp6C+YN0l8!e(9;La3}X9 zJ(CV=3D!jBN$SR)Z0d8q^Pcx*i}O4E1effC_ne-ZU)vNUle%r(B5s82Q!nE7%`4&L z+HB1E9NpBP!`_B;yZeQt5PgL$#sznrM)E5B> zh6Lu@>zyn8-N1W&7JbZH0tl`K&PXC6%ajhH+YrKqz=NJvCi;4)h+Tt}hc^2H0y=bA@WebjuQdB#>YWheWYBq8_JcS(unlUrAYs8S4U z#&OhVp*84Lf5BuBhJ!=oEzndQG-nDIS44e?GcZ~e8faxQ&38z$6x&KK;yB(+6mR@Y zeb0TG@O-3Iy>WIburTM!AG1y|lbKYcweeV3+RD~K@l`4QJH?xazKQt%2QRIy>xqkf z-I8)C{T}4%5A<46GU9*8vF4#2goO=`HbR0Yo9?qu8woy8athC^0_E&FXkYl%e@@TL AOaK4? diff --git a/__pycache__/gbase_agent.cpython-312.pyc b/__pycache__/gbase_agent.cpython-312.pyc index 4d01082467949cd56b34331b7bd06dc4938651c5..53f128139a3613a33c8d3a51d88df5c069cdd88b 100644 GIT binary patch delta 1880 zcmZ`(T})g>6uxtRc9&&$*|G~Pu)r?AdrM(UDL+!I(YCY-!89$kLY1t{0NXCRo4I!z zX|`e2#s||>>lh!5F|9GC225)(@ujgQ`e5u!6OyL6eQiw>edym)6BEz53lxY>GT%Kj z=bJOXb}76_ zigV*J7CU(a?PFU=b+A<>tzT+bREr)Wtx>lyBP={4(dgHNHqxf0=0$A-e@tA{YGUK< zto`Ki&C75^8|_R3kZBVz#>FwYA@SeL7Mtk;X7l1|AR^7n{}Z=$o!zpLJulD8`u649 z@*E+$*spdIpSYwplF;(6YLC3V90KzILJ+{xk|vuoEX~NKlBVw9)##}}s1%LPeo`eS zoI6R|@_FYN5#8>ijj%itFwkVu-6HH9{%<2)4OU8fw@Do>!n7UVwshue9xoXQi~~mfFL+ zU`G7qX>mKjP@q34h^zk49>v6g-ST9u;AhpmObL~#@}+WCGgoURNqKFZ;V_+o`05lT zuZpN|Z1`Dt)-O+jC8>&j(8xsF3cfv5@+$r~h+*CkR+Tw%AoLN)5hOL?;7Ts!(8 zvU&u5Nqn>Iu)h;%yceMlq5sqJgA0RT>8m<~aQ3+l8pks8eF*yz4j>Ex7+9bQol(M&n89Ufv6u9;d>AU0H#!Nt2!v& z3oJ1H!X}R{7gK!^W_5?9C1zOK#B~@Wt2{35Rz%5=_@`oT1pn4jO$OIy0?aR<8yNt&2JY4scT7xF&cgEF zD*qr3eDC<{;94>6w7k3cTJY2x@N1(-#`bxg8XfSHvzUN zj>kAVY37nV0nYD{umrG{Mq%GWhvY9W9J=H7-Sbz7-PNCyE^)K^A{h~H)qHrsu9piI z*xQ4WSX7SZ(2xOE%UaxI1~#EE`8#^??yqekVzuoM^=NG~c}dLIwvvdrR2w{jx>yRN zW+}K7hdB?xgBzbbxP1Q4E0^zI{}@!olVUliW;4ktX7C7%_zVKx9-ah%M+wiXE{~3a zIf7=V@Q3S2Ky-%hcjH;L6sSY-_S*Ltv*UOViw6=5YVU80#Hkpr8;RMn?nIA$(CmV` z2H}awrPyBnJ~;3dG^WLkI<{RQk;fjH^gi;{kg{vaLkT{=820DZ@rfrmd^`Z~Xh0mV JzpG-ze*tcGlB56t delta 2683 zcmaJ?ZERat89wLU>+5glJ9aX+wbL|?+az9F(xlN=5wwBFzyRp4< zuk)oFB|-%@jgm}9+k&cz7HLvSqHR@xAq_-9;t%j6!P8n@`hy?sD1T75PMXvoi1)pY zlTjg#bsnF4-jDO1_v?Jp{oyhHpZ$Ihfe~A~J@&^p{7v-d*J3`J-mO$!AR~K}YQTsR zS8CL-;)PkOhBLBKcfm0dRqBD(qU^^#r2%H6O$!2TuW~?Xxpk#kU3rPyX}!7yt3w_db2=hqpia z)oe{yLenOa>3lAe9qk&|^Evx^L#!4=9=?=Sg&s|6+S3&YhQ1qwskZ@u(Qf}nAbR=m z$L`q5!CO^zi!-myEZ4klekl@9`we&OukP3b-36={8#=@_N4JwNc=mbG-6G8U^^||$ zX*p5J(>cl5sschxTB)ZwhbfriBxUV}6NK!cP*b=>6v04M_Ate9FIG)gMeGuJN2;v9 zM@`Wbnt*XmRzy>tBrGftQy%wYb@huzU`xw(dYW=Yc1|`h@9bg9iewVdEGQ2ASLAm@ zQ!*XA#XEdr)_F3os9Gd8dngvk<&8)tU(6|yY%cPU$Gal^lc}thI%J&d=bOHL^Und>w_xhm03f7+U8TUDg@G?*5{Q)i z(S@gABWZ2@^*LgHBFy9YL2nv=>IhbezJAv>t!ZQdv;gi`naRV!iodmG_ike9kH6v*! z3L8SgvjDmb5V42d-B=G&5b{X92!NrK%dzCY4>SzPq#?R+@}{$Xd0^AI_jV|}{K8Mq z{^abX@0LOxE2lO?iC6s?sK+XQDa)BEWRsAKmd~~^iTzDwyrt^R z)0U%b_G#c_&jZAT@@2n)YuR$x9g$(8vTYvAvr{kjh5zfWoK`nF&Vay=dAPcV9$kOF zx|6z+s83>j{I!}WeU`7*9E;d}S29$54uJYDOdWj#Zt;Xmp55GCJ0{}=%8hpY_1bH6 z{w#d6mjEn~oQ4|9WHoixe;%rtsd^!w)78GX({A1H_T>sigFT8OX@nHQc@&mGR?R7K zhaD@XB0)jW5JmxXY;-%}LnjuJMyQx&XydPuV+~;2*i!wwP!%%`Z_p9$ZT!J;d^{L7 zR{KD3D<%bd1whB_i)_7lw%Sq9%3!mERD7qlwSLlg0JI9hR{A851>5GMK(wSWRV&QO zot+rFxM+LN5{+UOdJSkoXBM$s?*4A+?v)>Xx^m^tZx+G#2Fg0qV|jSnbcUK(5n&Qx z3Ls7yR!3YI8iM5rE)$kBnN;%WWYTh_*eL%?u!`36TfsZMSSFSTv18c%%B5r9*{&pP zS>=q!Ot4L|{l*yoA~e`;ySf)C7~k^i?|KZbQ`FRBYzg)#OufLzn$)ieH1fGe1n>#; p2(H4P`I>3f2jV>ff7f;UZ{x-9+{e$)`@*#9vt#_rroTz3`#-iVmS6w? diff --git a/agent_prompt.txt b/agent_prompt.txt index e7826b3..9572aaa 100644 --- a/agent_prompt.txt +++ b/agent_prompt.txt @@ -186,12 +186,8 @@ query_pattern = simple_field_match(conditions[0]) # 先匹配主要条件 - **渐进式扩展**：逐步放宽查询条件以发现更多相关数据 - **交叉验证**：使用多种方法验证搜索结果的完整性 -## 重要说明 -1. 查询的设备类型为第一优先级，比如笔记本和台式机。 -2. 针对"CPU处理器"和"GPU显卡"的查询，因为命名方式多样性，查询优先级最低。 -3. 如果确实无法找到完全匹配的数据，根据用户要求，可接受性能更高(更低)的CPU处理器和GPU显卡是作为代替。 --- -**执行提醒**：始终使用完整的文件路径参数调用工具，确保数据访问的准确性和安全性。在查询执行过程中，动态调整策略以适应不同的数据特征和查询需求。 +**重要说明**：所有文件路径中的 `[当前数据目录]` 将通过系统消息动态提供，请根据实际的数据目录路径进行操作。始终使用完整的文件路径参数调用工具，确保数据访问的准确性和安全性。在查询执行过程中，动态调整策略以适应不同的数据特征和查询需求。 + -**重要说明**：所有文件路径中的 `[当前数据目录]` 将通过系统消息动态提供，请根据实际的数据目录路径进行操作。 diff --git a/fastapi_app.py b/fastapi_app.py index 4281e08..691e6e1 100644 --- a/fastapi_app.py +++ b/fastapi_app.py @@ -62,6 +62,7 @@ class ChatRequest(BaseModel): extra: Optional[Dict] = None stream: Optional[bool] = False file_url: Optional[str] = None + extra_prompt: Optional[str] = None class ChatResponse(BaseModel): @@ -174,12 +175,13 @@ async def chat_completions(request: ChatRequest): # 动态设置请求的模型，支持从接口传入api_key、model_server和extra参数 update_agent_llm(agent, request.model, request.api_key, request.model_server) + extra_prompt = request.extra_prompt if request.extra_prompt else "" # 构建包含项目信息的消息上下文 messages = [ # 项目信息系统消息 { "role": "user", - "content": f"当前项目来自ZIP URL: {zip_url}，项目目录: {project_dir}。所有文件路径中的 '[当前数据目录]' 请替换为: {project_dir}" + "content": f"当前项目来自ZIP URL: {zip_url}，项目目录: {project_dir}。所有文件路径中的 '[当前数据目录]' 请替换为: {project_dir}\n"+ extra_prompt }, # 用户消息批量转换 *[{"role": msg.role, "content": msg.content} for msg in request.messages] diff --git a/gbase_agent.py b/gbase_agent.py index ebab503..53ecc9d 100644 --- a/gbase_agent.py +++ b/gbase_agent.py @@ -57,11 +57,6 @@ def read_mcp_settings_with_project_restriction(project_data_dir: str): return mcp_settings_json -def read_system_prompt(): - with open("./agent_prompt.txt", "r", encoding="utf-8") as f: - return f.read().strip() - - def read_system_prompt(): """读取通用的无状态系统prompt""" with open("./agent_prompt.txt", "r", encoding="utf-8") as f: @@ -73,32 +68,28 @@ def init_agent_service(): return init_agent_service_universal("qwen3-next") -def read_llm_config(): - """读取LLM配置文件""" - with open("./llm_config.json", "r") as f: - return json.load(f) def init_agent_service_with_project(project_id: str, project_data_dir: str, model_name: str = "qwen3-next"): """支持项目目录的agent初始化函数 - 保持向后兼容""" - llm_cfg = read_llm_config() - # 读取通用的系统prompt（无状态） system = read_system_prompt() # 读取MCP工具配置 tools = read_mcp_settings_with_project_restriction(project_data_dir) - # 使用指定的模型配置 - if model_name not in llm_cfg: - raise ValueError(f"Model '{model_name}' not found in llm_config.json. Available models: {list(llm_cfg.keys())}") + # 创建默认的LLM配置（可以通过update_agent_llm动态更新） + llm_config = { + "model": model_name, + "model_server": "https://openrouter.ai/api/v1", # 默认服务器 + "api_key": "default-key" # 默认密钥，实际使用时需要通过API传入 + } - llm_instance = llm_cfg[model_name] - if "llm_class" in llm_instance: - llm_instance = llm_instance.get("llm_class", TextChatAtOAI)(llm_instance) + # 创建LLM实例 + llm_instance = TextChatAtOAI(llm_config) bot = Assistant( - llm=llm_instance, # 使用指定的模型实例 + llm=llm_instance, # 使用默认LLM初始化，可通过update_agent_llm动态更新 name=f"数据库助手-{project_id}", description=f"项目 {project_id} 数据库查询", system_message=system, @@ -110,26 +101,24 @@ def init_agent_service_with_project(project_id: str, project_data_dir: str, mode def init_agent_service_universal(): """创建无状态的通用助手实例（使用默认LLM，可动态切换）""" - llm_cfg = read_llm_config() - # 读取通用的系统prompt（无状态） system = read_system_prompt() # 读取基础的MCP工具配置（不包含项目限制） tools = read_mcp_settings() - # 使用默认模型创建助手实例 - default_model = "qwen3-next" # 默认模型 - if default_model not in llm_cfg: - # 如果默认模型不存在，使用第一个可用模型 - default_model = list(llm_cfg.keys())[0] + # 创建默认的LLM配置（可以通过update_agent_llm动态更新） + llm_config = { + "model": "qwen3-next", # 默认模型 + "model_server": "https://openrouter.ai/api/v1", # 默认服务器 + "api_key": "default-key" # 默认密钥，实际使用时需要通过API传入 + } - llm_instance = llm_cfg[default_model] - if "llm_class" in llm_instance: - llm_instance = llm_instance.get("llm_class", TextChatAtOAI)(llm_instance) + # 创建LLM实例 + llm_instance = TextChatAtOAI(llm_config) bot = Assistant( - llm=llm_instance, # 使用默认LLM初始化 + llm=llm_instance, # 使用默认LLM初始化，可通过update_agent_llm动态更新 name="通用数据检索助手", description="无状态通用数据检索助手", system_message=system, @@ -163,8 +152,8 @@ def update_agent_llm(agent, model_name: str, api_key: str = None, model_server: def test(query="数据库里有几张表"): - # Define the agent - bot = init_agent_service() + # Define the agent - 使用通用初始化 + bot = init_agent_service_universal() # Chat messages = [] @@ -182,8 +171,8 @@ def test(query="数据库里有几张表"): def app_tui(): - # Define the agent - bot = init_agent_service() + # Define the agent - 使用通用初始化 + bot = init_agent_service_universal() # Chat messages = [] @@ -209,8 +198,8 @@ def app_tui(): def app_gui(): - # Define the agent - bot = init_agent_service() + # Define the agent - 使用通用初始化 + bot = init_agent_service_universal() chatbot_config = { "prompt.suggestions": [ "数据库里有几张表", diff --git a/llm_config.json b/llm_config.json deleted file mode 100644 index c5df2e2..0000000 --- a/llm_config.json +++ /dev/null @@ -1,65 +0,0 @@ -{ - "llama-33": { - "model": "gbase-llama-33", - "model_server": "http://llmapi:9009/v1", - "api_key": "any" - }, - "gpt-oss-120b": { - "model": "openai/gpt-oss-120b", - "model_server": "https://openrouter.ai/api/v1", - "api_key": "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", - "generate_cfg": { - "use_raw_api": true, - "fncall_prompt_type": "nous" - } - }, - "claude-3.7": { - "model": "claude-3-7-sonnet-20250219", - "model_server": "https://one.felo.me/v1", - "api_key": "sk-9gtHriq7C3jAvepq5dA0092a5cC24a54Aa83FbC99cB88b21-2", - "generate_cfg": { - "use_raw_api": true - } - }, - "gpt-4o": { - "model": "gpt-4o", - "model_server": "https://one-dev.felo.me/v1", - "api_key": "sk-hsKClH0Z695EkK5fDdB2Ec2fE13f4fC1B627BdBb8e554b5b-4", - "generate_cfg": { - "use_raw_api": true, - "fncall_prompt_type": "nous" - } - }, - "glm-45": { - "model_server": "https://open.bigmodel.cn/api/paas/v4", - "api_key": "0c9cbaca9d2bbf864990f1e1decdf340.dXRMsZCHTUbPQ0rm", - "model": "glm-4.5", - "generate_cfg": { - "use_raw_api": true - } - }, - "qwen3-next": { - "model": "qwen/qwen3-next-80b-a3b-instruct", - "model_server": "https://openrouter.ai/api/v1", - "api_key": "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" - }, - "deepresearch": { - "model": "alibaba/tongyi-deepresearch-30b-a3b", - "model_server": "https://openrouter.ai/api/v1", - "api_key": "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" - }, - "qwen3-coder": { - "model": "Qwen/Qwen3-Coder-30B-A3B-Instruct", - "model_server": "https://api-inference.modelscope.cn/v1", - "api_key": "ms-92027446-2787-4fd6-af01-f002459ec556" - }, - "openrouter-gpt4o": { - "model": "openai/gpt-4o", - "model_server": "https://openrouter.ai/api/v1", - "api_key": "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", - "generate_cfg": { - "use_raw_api": true, - "fncall_prompt_type": "nous" - } - } -} \ No newline at end of file diff --git a/mcp/directory_tree_wrapper_server.py b/mcp/directory_tree_wrapper_server.py deleted file mode 100644 index f533988..0000000 --- a/mcp/directory_tree_wrapper_server.py +++ /dev/null @@ -1,61 +0,0 @@ -#!/usr/bin/env python3 -""" -目录树 MCP包装器服务器 -提供安全的目录结构查看功能，限制在项目目录内 -""" - -import asyncio -import json -import sys -from mcp_wrapper import handle_wrapped_request - - -async def main(): - """主入口点""" - try: - while True: - # 从stdin读取 - line = await asyncio.get_event_loop().run_in_executor(None, sys.stdin.readline) - if not line: - break - - line = line.strip() - if not line: - continue - - try: - request = json.loads(line) - response = await handle_wrapped_request(request, "directory-tree-wrapper") - - # 写入stdout - sys.stdout.write(json.dumps(response) + "\n") - sys.stdout.flush() - - except json.JSONDecodeError: - error_response = { - "jsonrpc": "2.0", - "error": { - "code": -32700, - "message": "Parse error" - } - } - sys.stdout.write(json.dumps(error_response) + "\n") - sys.stdout.flush() - - except Exception as e: - error_response = { - "jsonrpc": "2.0", - "error": { - "code": -32603, - "message": f"Internal error: {str(e)}" - } - } - sys.stdout.write(json.dumps(error_response) + "\n") - sys.stdout.flush() - - except KeyboardInterrupt: - pass - - -if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file diff --git a/mcp/mcp_wrapper.py b/mcp/mcp_wrapper.py deleted file mode 100644 index 879b641..0000000 --- a/mcp/mcp_wrapper.py +++ /dev/null @@ -1,308 +0,0 @@ -#!/usr/bin/env python3 -""" -通用MCP工具包装器 -为所有MCP工具提供目录访问控制和安全限制 -""" - -import os -import sys -import json -import asyncio -import subprocess -from typing import Dict, Any, Optional -from pathlib import Path - - -class MCPSecurityWrapper: - """MCP安全包装器""" - - def __init__(self, allowed_directory: str): - self.allowed_directory = os.path.abspath(allowed_directory) - self.project_id = os.getenv("PROJECT_ID", "default") - - def validate_path(self, path: str) -> str: - """验证路径是否在允许的目录内""" - if not os.path.isabs(path): - path = os.path.abspath(path) - - # 规范化路径 - path = os.path.normpath(path) - - # 检查路径遍历 - if ".." in path.split(os.sep): - raise ValueError(f"路径遍历攻击被阻止: {path}") - - # 检查是否在允许的目录内 - if not path.startswith(self.allowed_directory): - raise ValueError(f"访问被拒绝: {path} 不在允许的目录 {self.allowed_directory} 内") - - return path - - def safe_execute_command(self, command: list, cwd: Optional[str] = None) -> Dict[str, Any]: - """安全执行命令，限制工作目录""" - if cwd is None: - cwd = self.allowed_directory - else: - cwd = self.validate_path(cwd) - - # 设置环境变量限制 - env = os.environ.copy() - env["PWD"] = cwd - env["PROJECT_DATA_DIR"] = self.allowed_directory - env["PROJECT_ID"] = self.project_id - - try: - result = subprocess.run( - command, - cwd=cwd, - env=env, - capture_output=True, - text=True, - timeout=30 # 30秒超时 - ) - - return { - "success": result.returncode == 0, - "stdout": result.stdout, - "stderr": result.stderr, - "returncode": result.returncode - } - except subprocess.TimeoutExpired: - return { - "success": False, - "stdout": "", - "stderr": "命令执行超时", - "returncode": -1 - } - except Exception as e: - return { - "success": False, - "stdout": "", - "stderr": str(e), - "returncode": -1 - } - - -async def handle_wrapped_request(request: Dict[str, Any], tool_name: str) -> Dict[str, Any]: - """处理包装后的MCP请求""" - try: - method = request.get("method") - params = request.get("params", {}) - request_id = request.get("id") - - allowed_dir = get_allowed_directory() - wrapper = MCPSecurityWrapper(allowed_dir) - - if method == "initialize": - return { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "protocolVersion": "2024-11-05", - "capabilities": { - "tools": {} - }, - "serverInfo": { - "name": f"{tool_name}-wrapper", - "version": "1.0.0" - } - } - } - - elif method == "ping": - return { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "pong": True - } - } - - elif method == "tools/list": - return { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "tools": get_tool_definitions(tool_name) - } - } - - elif method == "tools/call": - return await execute_tool_call(wrapper, tool_name, params, request_id) - - else: - return { - "jsonrpc": "2.0", - "id": request_id, - "error": { - "code": -32601, - "message": f"Unknown method: {method}" - } - } - - except Exception as e: - return { - "jsonrpc": "2.0", - "id": request.get("id"), - "error": { - "code": -32603, - "message": f"Internal error: {str(e)}" - } - } - - -def get_allowed_directory(): - """获取允许访问的目录""" - project_dir = os.getenv("PROJECT_DATA_DIR", "./data") - return os.path.abspath(project_dir) - - -def get_tool_definitions(tool_name: str) -> list: - """根据工具名称返回工具定义""" - if tool_name == "ripgrep-wrapper": - return [ - { - "name": "ripgrep_search", - "description": "在项目目录内搜索文本", - "inputSchema": { - "type": "object", - "properties": { - "pattern": {"type": "string", "description": "搜索模式"}, - "path": {"type": "string", "description": "搜索路径（相对于项目目录）"}, - "maxResults": {"type": "integer", "default": 100} - }, - "required": ["pattern"] - } - } - ] - elif tool_name == "directory-tree-wrapper": - return [ - { - "name": "get_directory_tree", - "description": "获取项目目录结构", - "inputSchema": { - "type": "object", - "properties": { - "path": {"type": "string", "description": "目录路径（相对于项目目录）"}, - "max_depth": {"type": "integer", "default": 3} - } - } - } - ] - else: - return [] - - -async def execute_tool_call(wrapper: MCPSecurityWrapper, tool_name: str, params: Dict[str, Any], request_id: Any) -> Dict[str, Any]: - """执行工具调用""" - try: - if tool_name == "ripgrep-wrapper": - return await execute_ripgrep_search(wrapper, params, request_id) - elif tool_name == "directory-tree-wrapper": - return await execute_directory_tree(wrapper, params, request_id) - else: - return { - "jsonrpc": "2.0", - "id": request_id, - "error": { - "code": -32601, - "message": f"Unknown tool: {tool_name}" - } - } - except Exception as e: - return { - "jsonrpc": "2.0", - "id": request_id, - "error": { - "code": -32603, - "message": str(e) - } - } - - -async def execute_ripgrep_search(wrapper: MCPSecurityWrapper, params: Dict[str, Any], request_id: Any) -> Dict[str, Any]: - """执行ripgrep搜索""" - pattern = params.get("pattern", "") - path = params.get("path", ".") - max_results = params.get("maxResults", 100) - - # 验证和构建搜索路径 - search_path = os.path.join(wrapper.allowed_directory, path) - search_path = wrapper.validate_path(search_path) - - # 构建ripgrep命令 - command = [ - "rg", - "--json", - "--max-count", str(max_results), - pattern, - search_path - ] - - result = wrapper.safe_execute_command(command) - - if result["success"]: - return { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "content": [ - { - "type": "text", - "text": result["stdout"] - } - ] - } - } - else: - return { - "jsonrpc": "2.0", - "id": request_id, - "error": { - "code": -32603, - "message": f"搜索失败: {result['stderr']}" - } - } - - -async def execute_directory_tree(wrapper: MCPSecurityWrapper, params: Dict[str, Any], request_id: Any) -> Dict[str, Any]: - """执行目录树获取""" - path = params.get("path", ".") - max_depth = params.get("max_depth", 3) - - # 验证和构建目录路径 - dir_path = os.path.join(wrapper.allowed_directory, path) - dir_path = wrapper.validate_path(dir_path) - - # 构建目录树命令 - command = [ - "find", - dir_path, - "-type", "d", - "-maxdepth", str(max_depth) - ] - - result = wrapper.safe_execute_command(command) - - if result["success"]: - return { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "content": [ - { - "type": "text", - "text": result["stdout"] - } - ] - } - } - else: - return { - "jsonrpc": "2.0", - "id": request_id, - "error": { - "code": -32603, - "message": f"获取目录树失败: {result['stderr']}" - } - } \ No newline at end of file diff --git a/mcp/ripgrep_wrapper_server.py b/mcp/ripgrep_wrapper_server.py deleted file mode 100644 index ce6ac84..0000000 --- a/mcp/ripgrep_wrapper_server.py +++ /dev/null @@ -1,61 +0,0 @@ -#!/usr/bin/env python3 -""" -Ripgrep MCP包装器服务器 -提供安全的文本搜索功能，限制在项目目录内 -""" - -import asyncio -import json -import sys -from mcp_wrapper import handle_wrapped_request - - -async def main(): - """主入口点""" - try: - while True: - # 从stdin读取 - line = await asyncio.get_event_loop().run_in_executor(None, sys.stdin.readline) - if not line: - break - - line = line.strip() - if not line: - continue - - try: - request = json.loads(line) - response = await handle_wrapped_request(request, "ripgrep-wrapper") - - # 写入stdout - sys.stdout.write(json.dumps(response) + "\n") - sys.stdout.flush() - - except json.JSONDecodeError: - error_response = { - "jsonrpc": "2.0", - "error": { - "code": -32700, - "message": "Parse error" - } - } - sys.stdout.write(json.dumps(error_response) + "\n") - sys.stdout.flush() - - except Exception as e: - error_response = { - "jsonrpc": "2.0", - "error": { - "code": -32603, - "message": f"Internal error: {str(e)}" - } - } - sys.stdout.write(json.dumps(error_response) + "\n") - sys.stdout.flush() - - except KeyboardInterrupt: - pass - - -if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file