83 lines
2.9 KiB
Python
Executable File
83 lines
2.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
cbm-query-wrapper.py — P3: CBM 查询钩子
|
|
封装 CBM CLI 调用，查询结果自动写入织忆。
|
|
用法:
|
|
python3 cbm-query-wrapper.py search_graph --project X --name-pattern Y
|
|
python3 cbm-query-wrapper.py get_architecture --project X
|
|
"""
|
|
import json, os, subprocess, sys, urllib.request, urllib.error
|
|
|
|
CBM = os.path.expanduser("~/.local/bin/codebase-memory-mcp")
|
|
ZHIYI_URL = "http://localhost:7821/api/v1/commit"
|
|
ZHIYI_KEY = os.environ.get("ZHIYI_API_KEY", "zhiyi-dev-key-2026")
|
|
|
|
def zhiyi_write(content, category="distilled"):
|
|
"""写入织忆"""
|
|
payload = json.dumps({
|
|
"content": content,
|
|
"category": category,
|
|
"agent_id": "a06"
|
|
}).encode()
|
|
req = urllib.request.Request(ZHIYI_URL, data=payload,
|
|
headers={"X-API-Key": ZHIYI_KEY, "Content-Type": "application/json"})
|
|
try:
|
|
urllib.request.urlopen(req, timeout=5)
|
|
return True
|
|
except urllib.error.URLError as e:
|
|
print(f"[cbm-wrapper] ⚠ 织忆写入失败: {e}", file=sys.stderr)
|
|
return False
|
|
|
|
def run_cbm(args):
|
|
"""执行 CBM CLI 并返回 JSON 结果"""
|
|
cmd = [CBM, "cli"] + args
|
|
try:
|
|
result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
|
|
return result.stdout, result.returncode
|
|
except FileNotFoundError:
|
|
print(f"[cbm-wrapper] ❌ CBM binary not found: {CBM}", file=sys.stderr)
|
|
return "", 1
|
|
except subprocess.TimeoutExpired:
|
|
print("[cbm-wrapper] ⚠ CBM timed out (60s)", file=sys.stderr)
|
|
return "", 1
|
|
|
|
def main():
|
|
if len(sys.argv) < 2:
|
|
print("用法: cbm-query-wrapper.py <cbm_cli_args...>")
|
|
sys.exit(1)
|
|
|
|
tool = sys.argv[1]
|
|
args = sys.argv[1:]
|
|
|
|
# 执行 CBM 查询
|
|
stdout, rc = run_cbm(args)
|
|
print(stdout, end="")
|
|
|
|
if rc != 0:
|
|
sys.exit(rc)
|
|
|
|
# 对有价值的结果写入织忆
|
|
try:
|
|
data = json.loads(stdout)
|
|
project = data.get("project", data.get("name", "unknown"))
|
|
nodes = data.get("nodes", data.get("total_nodes", 0))
|
|
edges = data.get("edges") or data.get("total_edges") or 0
|
|
|
|
if tool in ("index_repository",) and int(nodes) > 0:
|
|
content = f"CBM索引完成: {project} ({nodes}节点/{edges}边)"
|
|
zhiyi_write(content)
|
|
elif tool in ("get_architecture", "list_projects"):
|
|
content = f"CBM架构查询: {project}"
|
|
zhiyi_write(content)
|
|
elif tool in ("search_graph", "trace_path"):
|
|
results = data.get("results") or data.get("nodes") or []
|
|
if results:
|
|
names = [r.get("name","?") for r in results[:5]]
|
|
content = f"CBM查询 '{tool}': 找到 {len(results)} 个结果 ({', '.join(names)})"
|
|
zhiyi_write(content)
|
|
except (json.JSONDecodeError, KeyError, TypeError, ValueError):
|
|
pass # 非 JSON 输出或不匹配模式就不写入
|
|
|
|
if __name__ == "__main__":
|
|
main()
|