Yuzhiran Dev
301 lines
12 KiB
Python
301 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
MCP Search Server — provides web search via opencode infrastructure.
|
|
|
|
Two search methods (automatic fallback):
|
|
1. npx opencode run (rate-limited but returns real web results)
|
|
2. opencode-go API + model training data (no rate limit, less fresh)
|
|
|
|
Usage:
|
|
python3 scripts/mcp_search_server.py # MCP server (stdio)
|
|
python3 scripts/mcp_search_server.py --query Q # one-shot search
|
|
python3 scripts/mcp_search_server.py --url U # one-shot webfetch
|
|
"""
|
|
import json, os, subprocess, sys, time
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
|
CACHE_FILE = PROJECT_ROOT / "automation" / "data" / "mcp_search_cache.json"
|
|
SESSION_FILE = PROJECT_ROOT / "automation" / "data" / "mcp_session.txt"
|
|
CACHE_TTL = 3600
|
|
SESSION_TITLE = "opencode搜索"
|
|
|
|
API_BASE = "https://opencode.ai/zen/go/v1"
|
|
API_KEY = os.environ.get("OPENCODE_API_KEY", "")
|
|
if not API_KEY:
|
|
try:
|
|
from dotenv import load_dotenv
|
|
env_path = PROJECT_ROOT / "platform" / "backend" / ".env"
|
|
load_dotenv(env_path)
|
|
API_KEY = os.environ.get("OPENCODE_API_KEY", "")
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# ── session (reuse same session for all MCP searches) ─────────────
|
|
def _load_session() -> Optional[str]:
|
|
if SESSION_FILE.exists():
|
|
try:
|
|
return SESSION_FILE.read_text().strip() or None
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
def _save_session_from_output(stdout: str):
|
|
for line in stdout.strip().split("\n"):
|
|
try:
|
|
ev = json.loads(line)
|
|
sid = ev.get("sessionID") or ev.get("part", {}).get("sessionID")
|
|
if sid:
|
|
SESSION_FILE.parent.mkdir(parents=True, exist_ok=True)
|
|
SESSION_FILE.write_text(sid)
|
|
return
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# ── cache ─────────────────────────────────────────────────────────
|
|
def _check_cache(query: str) -> Optional[List[Dict]]:
|
|
if not CACHE_FILE.exists():
|
|
return None
|
|
try:
|
|
data = json.loads(CACHE_FILE.read_text())
|
|
entry = data.get(query)
|
|
if entry and time.time() - entry.get("ts", 0) < CACHE_TTL:
|
|
return entry.get("results")
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
def _write_cache(query: str, results: List[Dict]):
|
|
CACHE_FILE.parent.mkdir(parents=True, exist_ok=True)
|
|
data = {}
|
|
if CACHE_FILE.exists():
|
|
try:
|
|
data = json.loads(CACHE_FILE.read_text())
|
|
except Exception:
|
|
pass
|
|
data[query] = {"ts": time.time(), "results": results}
|
|
keys = sorted(data.keys(), key=lambda k: data[k].get("ts", 0), reverse=True)[:200]
|
|
CACHE_FILE.write_text(json.dumps({k: data[k] for k in keys}, ensure_ascii=False))
|
|
|
|
|
|
# ── method 1: npx opencode run ────────────────────────────────────
|
|
def _search_via_opencode_cli(query: str, max_results: int) -> Optional[List[Dict]]:
|
|
"""Use npx opencode run to execute websearch tool (short timeout)."""
|
|
sid = _load_session()
|
|
args = ["npx", "opencode", "run", f"websearch {query}", "--format", "json", "--title", SESSION_TITLE]
|
|
if sid:
|
|
args.extend(["--session", sid, "--continue"])
|
|
try:
|
|
r = subprocess.run(
|
|
args, capture_output=True, text=True, timeout=15,
|
|
env={**os.environ, "OPENCODE_DISABLE_AUTOUPDATE": "1"}
|
|
)
|
|
except subprocess.TimeoutExpired:
|
|
return None
|
|
except Exception:
|
|
return None
|
|
if r.returncode != 0:
|
|
return None
|
|
# Save session ID for reuse
|
|
_save_session_from_output(r.stdout)
|
|
for line in r.stdout.strip().split("\n"):
|
|
try:
|
|
ev = json.loads(line)
|
|
if ev.get("type") == "tool_use":
|
|
part = ev.get("part", {})
|
|
state = part.get("state", {})
|
|
if part.get("tool") == "websearch" and state.get("status") == "completed":
|
|
data = json.loads(state["output"])
|
|
results = []
|
|
for item in (data.get("results") or [])[:max_results]:
|
|
url = (item.get("url") or "").strip()
|
|
title = (item.get("title") or "").strip()
|
|
excerpts = item.get("excerpts") or []
|
|
content = (excerpts[0] if excerpts else "")[:500]
|
|
if url and title:
|
|
results.append({"title": title, "url": url, "content": content, "source": "opencode_cli"})
|
|
return results
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
# ── method 2: opencode-go API + training data ─────────────────────
|
|
def _search_via_api(query: str, max_results: int) -> Optional[List[Dict]]:
|
|
"""Use opencode-go API to answer query from training data (no rate limit)."""
|
|
if not API_KEY:
|
|
return None
|
|
import requests
|
|
try:
|
|
resp = requests.post(
|
|
f"{API_BASE}/chat/completions",
|
|
headers={"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"},
|
|
json={
|
|
"model": "deepseek-v4-flash",
|
|
"messages": [{"role": "user", "content": (
|
|
f"你现在是一个网络搜索工具。用户查询: {query[:100]}\n\n"
|
|
f"请根据你的训练数据，提供{max_results}条最相关的网页结果，包含标题、URL和摘要。"
|
|
f"以JSON格式输出: [{{\"title\":\"...\",\"url\":\"...\",\"content\":\"...\"}}]"
|
|
f"仅输出JSON数组，不要其他文字。如果URL不确定，用合理占位。"
|
|
)}],
|
|
"temperature": 0.3,
|
|
"max_tokens": 2000,
|
|
},
|
|
timeout=30
|
|
)
|
|
data = resp.json()
|
|
content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
|
|
# Extract JSON array
|
|
import re as _re
|
|
m = _re.search(r'\[.*?\]', content, _re.DOTALL)
|
|
if m:
|
|
items = json.loads(m.group())
|
|
if isinstance(items, list):
|
|
for item in items:
|
|
item["source"] = "opencode_api"
|
|
return items[:max_results]
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
# ── search ─────────────────────────────────────────────────────────
|
|
def web_search(query: str, max_results: int = 8) -> List[Dict]:
|
|
max_results = min(max_results, 10)
|
|
cached = _check_cache(query)
|
|
if cached:
|
|
return cached[:max_results]
|
|
|
|
results = _search_via_opencode_cli(query, max_results)
|
|
if results:
|
|
_write_cache(query, results)
|
|
return results
|
|
|
|
results = _search_via_api(query, max_results)
|
|
if results:
|
|
_write_cache(query, results)
|
|
return results
|
|
|
|
return []
|
|
|
|
|
|
def webfetch(url: str) -> Optional[str]:
|
|
sid = _load_session()
|
|
args = ["npx", "opencode", "run", f"webfetch {url}", "--format", "json", "--title", SESSION_TITLE]
|
|
if sid:
|
|
args.extend(["--session", sid, "--continue"])
|
|
try:
|
|
r = subprocess.run(
|
|
args, capture_output=True, text=True, timeout=60,
|
|
env={**os.environ, "OPENCODE_DISABLE_AUTOUPDATE": "1"}
|
|
)
|
|
if r.returncode == 0:
|
|
_save_session_from_output(r.stdout)
|
|
for line in r.stdout.strip().split("\n"):
|
|
try:
|
|
ev = json.loads(line)
|
|
if ev.get("type") == "tool_use":
|
|
p = ev.get("part", {})
|
|
s = p.get("state", {})
|
|
if p.get("tool") == "webfetch" and s.get("status") == "completed":
|
|
return s.get("output", "")[:10000]
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
# ── MCP protocol (JSON-RPC 2.0 over stdio) ────────────────────────
|
|
def _read_msg() -> Optional[Dict]:
|
|
line = sys.stdin.readline()
|
|
if not line:
|
|
return None
|
|
try:
|
|
return json.loads(line)
|
|
except json.JSONDecodeError:
|
|
return None
|
|
|
|
def _send_msg(msg: Dict):
|
|
sys.stdout.write(json.dumps(msg, ensure_ascii=False) + "\n")
|
|
sys.stdout.flush()
|
|
|
|
def _send_error(req_id: Any, code: int, message: str):
|
|
_send_msg({"jsonrpc": "2.0", "id": req_id, "error": {"code": code, "message": message}})
|
|
|
|
def _send_result(req_id: Any, result: Any):
|
|
_send_msg({"jsonrpc": "2.0", "id": req_id, "result": result})
|
|
|
|
|
|
def serve():
|
|
sys.stdin.reconfigure(encoding="utf-8")
|
|
sys.stdout.reconfigure(encoding="utf-8")
|
|
while True:
|
|
msg = _read_msg()
|
|
if msg is None:
|
|
break
|
|
req_id = msg.get("id")
|
|
method = msg.get("method", "")
|
|
params = msg.get("params", {})
|
|
if method == "initialize":
|
|
_send_result(req_id, {
|
|
"protocolVersion": "2024-11-05",
|
|
"capabilities": {"tools": {"listChanged": False}},
|
|
"serverInfo": {"name": "opencode-search-mcp", "version": "1.0.0"}
|
|
})
|
|
elif method == "notifications/initialized":
|
|
pass
|
|
elif method == "tools/list":
|
|
_send_result(req_id, {"tools": [
|
|
{"name": "web_search", "description": "Search the web. Returns up to 10 results with title, url, content.", "inputSchema": {
|
|
"type": "object", "properties": {
|
|
"query": {"type": "string", "description": "Search query"},
|
|
"max_results": {"type": "number", "description": "Max results (1-10)", "default": 8}
|
|
}, "required": ["query"]
|
|
}},
|
|
{"name": "webfetch", "description": "Fetch and extract content from a URL.", "inputSchema": {
|
|
"type": "object", "properties": {"url": {"type": "string", "description": "URL to fetch"}},
|
|
"required": ["url"]
|
|
}}
|
|
]})
|
|
elif method == "tools/call":
|
|
name = params.get("name", "")
|
|
args = params.get("arguments", {})
|
|
try:
|
|
if name == "web_search":
|
|
results = web_search(args.get("query", ""), int(args.get("max_results", 8)))
|
|
_send_result(req_id, {"content": [{"type": "text", "text": json.dumps(results, ensure_ascii=False)}]})
|
|
elif name == "webfetch":
|
|
content = webfetch(args.get("url", ""))
|
|
_send_result(req_id, {"content": [{"type": "text", "text": content or "Failed to fetch URL"}]})
|
|
else:
|
|
_send_error(req_id, -32601, f"Unknown tool: {name}")
|
|
except Exception as e:
|
|
_send_error(req_id, -32603, str(e))
|
|
elif method == "shutdown":
|
|
_send_result(req_id, {})
|
|
break
|
|
else:
|
|
_send_error(req_id, -32601, f"Unknown method: {method}")
|
|
|
|
|
|
def main():
|
|
if "--query" in sys.argv:
|
|
idx = sys.argv.index("--query")
|
|
q = sys.argv[idx + 1] if idx + 1 < len(sys.argv) else ""
|
|
print(json.dumps(web_search(q), ensure_ascii=False, indent=2))
|
|
return
|
|
if "--url" in sys.argv:
|
|
idx = sys.argv.index("--url")
|
|
u = sys.argv[idx + 1] if idx + 1 < len(sys.argv) else ""
|
|
print(webfetch(u) or "Failed")
|
|
return
|
|
serve()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|