354 lines
14 KiB
Python
354 lines
14 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
统一记忆入口 (memory_recall.py)
|
||
|
||
同时查询织忆、TencentDB、Soulful 三套记忆系统，合并去重后输出统一上下文。
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import threading
|
||
import urllib.error
|
||
import urllib.request
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
# ── 常量 ──────────────────────────────────────────────────────────────────────
|
||
|
||
ZHIYI_URL = "http://localhost:7821/api/v1/recall"
|
||
ZHIYI_HEADERS = {
|
||
"X-API-Key": "zhiyi-dev-key-2026",
|
||
"Content-Type": "application/json",
|
||
}
|
||
TENCENTDB_URL = "http://localhost:8420"
|
||
SOULFUL_DIR = Path.home() / ".hermes" / "soulful"
|
||
|
||
DEFAULT_TOP_K = 3
|
||
TIMEOUT_SEC = 5
|
||
|
||
# ── 工具函数 ──────────────────────────────────────────────────────────────────
|
||
|
||
def http_post(url: str, payload: dict, headers: dict | None = None, timeout: int = TIMEOUT_SEC) -> dict | None:
|
||
"""POST JSON，返回解析后的 dict；失败返回 None 并打印 warning。"""
|
||
try:
|
||
req = urllib.request.Request(
|
||
url,
|
||
data=json.dumps(payload).encode("utf-8"),
|
||
headers=headers or {},
|
||
method="POST",
|
||
)
|
||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||
return json.loads(resp.read().decode("utf-8"))
|
||
except Exception as e:
|
||
print(f" [warning] {url} -> {e}", flush=True)
|
||
return None
|
||
|
||
|
||
def read_json_file(path: Path) -> dict | list | None:
|
||
"""读取 JSON 文件，失败返回 None 并打印 warning。"""
|
||
try:
|
||
return json.loads(path.read_text(encoding="utf-8"))
|
||
except Exception as e:
|
||
print(f" [warning] read {path} -> {e}", flush=True)
|
||
return None
|
||
|
||
|
||
def read_jsonl(path: Path):
|
||
"""逐行读取 JSONL，yeild dict；失败打印 warning 并跳过。"""
|
||
try:
|
||
for line in path.read_text(encoding="utf-8").splitlines():
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
yield json.loads(line)
|
||
except json.JSONDecodeError:
|
||
pass
|
||
except Exception as e:
|
||
print(f" [warning] read {path} -> {e}", flush=True)
|
||
|
||
|
||
# ── 查询函数（供线程调用） ──────────────────────────────────────────────────────
|
||
|
||
def query_zhiyi(query: str, top_k: int) -> dict:
|
||
result = http_post(ZHIYI_URL, {"query": query, "top_k": top_k}, ZHIYI_HEADERS)
|
||
if result is None:
|
||
return {"error": True, "results": []}
|
||
return {"error": False, "results": result.get("results", [])}
|
||
|
||
|
||
def query_tencentdb_l1(query: str, top_k: int) -> dict:
|
||
result = http_post(
|
||
f"{TENCENTDB_URL}/search/memories",
|
||
{"query": query, "top_k": top_k},
|
||
)
|
||
if result is None:
|
||
return {"error": True, "results": []}
|
||
return {"error": False, "results": result.get("results", [])}
|
||
|
||
|
||
def query_tencentdb_l0(query: str, top_k: int) -> dict:
|
||
result = http_post(
|
||
f"{TENCENTDB_URL}/search/conversations",
|
||
{"query": query, "top_k": top_k},
|
||
)
|
||
if result is None:
|
||
return {"error": True, "results": []}
|
||
return {"error": False, "results": result.get("results", [])}
|
||
|
||
|
||
# ── Soulful 读取 ───────────────────────────────────────────────────────────────
|
||
|
||
def load_soulful() -> dict:
|
||
"""加载 Soulful 三件套，返回聚合 dict；各文件缺失不报错。"""
|
||
profile = read_json_file(SOULFUL_DIR / "user-profile.json") or {}
|
||
cares_data = read_json_file(SOULFUL_DIR / "cares-queue.json") or {}
|
||
heart_traces = list(read_jsonl(SOULFUL_DIR / "heart-traces.jsonl"))
|
||
|
||
return {
|
||
"error": False,
|
||
"profile": profile,
|
||
"cares": cares_data.get("cares", []) if isinstance(cares_data, dict) else [],
|
||
"heart_traces": heart_traces,
|
||
}
|
||
|
||
|
||
# ── 并行查询 ───────────────────────────────────────────────────────────────────
|
||
|
||
def parallel_query(query: str, top_k: int):
|
||
"""用三个线程同时查询三个后端 + Soulful 文件，返回 (zhiyi, l1, l0, soulful)。"""
|
||
|
||
def run(query: str, top_k: int, fn):
|
||
return fn(query, top_k)
|
||
|
||
threads = []
|
||
zhiyi_result = {}
|
||
l1_result = {}
|
||
l0_result = {}
|
||
soulful_result = {}
|
||
|
||
t1 = threading.Thread(target=lambda: zhiyi_result.update(run(query, top_k, query_zhiyi)))
|
||
t2 = threading.Thread(target=lambda: l1_result.update(run(query, top_k, query_tencentdb_l1)))
|
||
t3 = threading.Thread(target=lambda: l0_result.update(run(query, top_k, query_tencentdb_l0)))
|
||
t4 = threading.Thread(target=lambda: soulful_result.update(load_soulful()))
|
||
|
||
for t in (t1, t2, t3, t4):
|
||
t.start()
|
||
for t in (t1, t2, t3, t4):
|
||
t.join()
|
||
|
||
return zhiyi_result, l1_result, l0_result, soulful_result
|
||
|
||
|
||
# ── 格式化输出 ────────────────────────────────────────────────────────────────
|
||
|
||
import re
|
||
|
||
def _fmt_val(v):
|
||
"""格式化 profile 中的 value：dict/list 走 JSON.dumps，其余直接转 str。"""
|
||
if isinstance(v, (dict, list)):
|
||
return json.dumps(v, ensure_ascii=False)
|
||
return str(v)
|
||
|
||
def format_pretty(zhiyi, l1, l0, soulful):
|
||
"""pretty 模式：分区块显示。"""
|
||
lines = []
|
||
|
||
# ── 语义记忆（织忆） ──
|
||
if zhiyi.get("error"):
|
||
lines.append("## 语义记忆（织忆）")
|
||
lines.append("暂不可用")
|
||
else:
|
||
lines.append("## 语义记忆（织忆）")
|
||
for item in zhiyi.get("results", []):
|
||
if isinstance(item, dict):
|
||
content = item.get("content", "")
|
||
score = item.get("score", 0.0)
|
||
cat = item.get("category", "")
|
||
ts = item.get("timestamp", "")
|
||
else:
|
||
content = str(item)
|
||
score, cat, ts = 0.0, "", ""
|
||
cat_str = f" [{cat}]" if cat else ""
|
||
ts_str = f" ({ts})" if ts else ""
|
||
lines.append(f"[{content}]{cat_str}{ts_str} (score: {score:.2f})")
|
||
lines.append("")
|
||
|
||
# ── 人格记忆（TencentDB L1） ──
|
||
if l1.get("error"):
|
||
lines.append("## 人格记忆（TencentDB）")
|
||
lines.append("暂不可用\n")
|
||
else:
|
||
lines.append("## 人格记忆（TencentDB）")
|
||
results = l1.get("results", [])
|
||
# TencentDB 返回的 results 是 markdown 字符串，不是 JSON 数组
|
||
if isinstance(results, str):
|
||
# 整段输出，不拆分
|
||
lines.append(results)
|
||
else:
|
||
for item in results:
|
||
if isinstance(item, dict):
|
||
content = item.get("content", "")
|
||
score = item.get("score", 0.0)
|
||
ptype = item.get("type", "instruction")
|
||
scene = item.get("scene", "")
|
||
else:
|
||
content = str(item)
|
||
score, ptype, scene = 0.0, "instruction", ""
|
||
scene_str = f" [{scene}]" if scene else ""
|
||
lines.append(f"[type: {ptype}]{scene_str}")
|
||
lines.append(f"[{content}] (score: {score:.2f})")
|
||
lines.append("")
|
||
|
||
# ── 对话记忆（TencentDB L0） ──
|
||
if l0.get("error"):
|
||
lines.append("## 对话记忆（TencentDB L0）")
|
||
lines.append("暂不可用\n")
|
||
else:
|
||
lines.append("## 对话记忆（TencentDB L0）")
|
||
results = l0.get("results", [])
|
||
# TencentDB 返回的 results 是 markdown 字符串，不是 JSON 数组
|
||
if isinstance(results, str):
|
||
lines.append(results)
|
||
else:
|
||
for item in results:
|
||
if isinstance(item, dict):
|
||
msg = item.get("message", "")
|
||
score = item.get("score", 0.0)
|
||
ts = item.get("timestamp", "")
|
||
session = item.get("session", "")
|
||
else:
|
||
msg = str(item)
|
||
score, ts, session = 0.0, "", ""
|
||
session_str = f"[{session}]" if session else ""
|
||
ts_str = f" {ts}" if ts else ""
|
||
lines.append(f"{session_str}{ts_str}")
|
||
lines.append(f"[{msg}] (score: {score:.2f})")
|
||
lines.append("")
|
||
|
||
# ── 关系感知（Soulful） ──
|
||
lines.append("## 关系感知（Soulful）")
|
||
|
||
if soulful.get("error"):
|
||
lines.append("暂不可用\n")
|
||
else:
|
||
# 画像
|
||
lines.append("### 画像")
|
||
profile = soulful.get("profile", {})
|
||
if profile:
|
||
for key, value in profile.items():
|
||
lines.append(f"{key}: {_fmt_val(value)}")
|
||
else:
|
||
lines.append("(无)")
|
||
|
||
# 牵挂
|
||
lines.append("\n### 待关注（牵挂）")
|
||
cares = soulful.get("cares", [])
|
||
if cares:
|
||
for item in cares:
|
||
content = item.get("content", "")
|
||
due = item.get("follow_up_date") or item.get("due") or "null"
|
||
lines.append(f"- [{content}] (due: {due})")
|
||
else:
|
||
lines.append("(无)")
|
||
|
||
# 心迹
|
||
lines.append("\n### 心迹")
|
||
heart_traces = soulful.get("heart_traces", [])
|
||
if heart_traces:
|
||
for item in heart_traces:
|
||
ts = item.get("timestamp", "")
|
||
content = item.get("content", "")
|
||
lines.append(f"[{ts}] [{content}]")
|
||
else:
|
||
lines.append("(无)")
|
||
|
||
return "\n".join(lines)
|
||
|
||
|
||
def format_compact(zhiyi, l1, l0, soulful):
|
||
"""compact 模式：单行动态摘要。"""
|
||
parts = []
|
||
|
||
# 织忆
|
||
if not zhiyi.get("error"):
|
||
results = zhiyi.get("results", [])
|
||
if results:
|
||
top = results[0]
|
||
parts.append(f"织忆/{top.get('content', '')[:40]}(s={top.get('score',0):.2f})")
|
||
|
||
# L1 - TencentDB returns string for markdown, list for parsed
|
||
if not l1.get("error"):
|
||
results = l1.get("results", [])
|
||
if isinstance(results, str):
|
||
# Markdown string - extract content from first meaningful line
|
||
# L1 format: "- **[persona/instruction]** ..." then blank then " content"
|
||
# or " content" appears on 2nd or 3rd non-empty line
|
||
lines = [l for l in results.split("\n") if l.strip() and not l.startswith("Found") and not l.startswith("---")]
|
||
if len(lines) >= 2:
|
||
# Second line has the actual content text (after blank line)
|
||
content = lines[1].strip().lstrip("-* []").strip()
|
||
parts.append(f"人格:{content[:35]}...")
|
||
elif lines:
|
||
content = lines[0].strip().lstrip("-* []").strip()
|
||
parts.append(f"人格:{content[:35]}")
|
||
elif results:
|
||
top = results[0]
|
||
ptype = top.get("type", "?")
|
||
parts.append(f"人格/{ptype}:{top.get('content', '')[:30]}(s={top.get('score',0):.2f})")
|
||
|
||
# L0 - TencentDB returns string for markdown, list for parsed
|
||
if not l0.get("error"):
|
||
results = l0.get("results", [])
|
||
if isinstance(results, str):
|
||
# L0 format: "**[user]** Session: ... (score: X)\n\nMESSAGE CONTENT\n\n---..."
|
||
# Match header line + blank line + first non-empty content line
|
||
m = re.search(r'\*\*\[(user|assistant)\]\*\*.*?\n\n(.+?)(?:\n\n|---\Z)', results, re.DOTALL)
|
||
if m:
|
||
msg = m.group(2).strip()[:40]
|
||
parts.append(f"对话:{msg}")
|
||
elif "**[user]**" in results:
|
||
for line in results.split("\n"):
|
||
if "**[user]**" in line or "**[assistant]**" in line:
|
||
parts.append(f"对话:{line[:50]}")
|
||
break
|
||
elif results:
|
||
top = results[0]
|
||
msg = top.get("message", "")[:40]
|
||
parts.append(f"对话:{msg}(s={top.get('score',0):.2f})")
|
||
|
||
# Soulful 牵挂
|
||
if not soulful.get("error"):
|
||
cares = soulful.get("cares", [])
|
||
if cares:
|
||
parts.append(f"牵挂:{cares[0].get('content', '')[:30]}")
|
||
|
||
if not parts:
|
||
return "[memory] no results"
|
||
return "[memory] " + " | ".join(parts)
|
||
|
||
|
||
# ── 入口 ──────────────────────────────────────────────────────────────────────
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="统一记忆入口")
|
||
parser.add_argument("query", help="查询文本")
|
||
parser.add_argument("--top-k", type=int, default=DEFAULT_TOP_K, help=f"每系统返回条数 (default: {DEFAULT_TOP_K})")
|
||
parser.add_argument("--format", choices=["pretty", "compact"], default="pretty", help="输出格式 (default: pretty)")
|
||
args = parser.parse_args()
|
||
|
||
# 实际并行查询
|
||
zhiyi, l1, l0, soulful = parallel_query(args.query, args.top_k)
|
||
|
||
if args.format == "compact":
|
||
print(format_compact(zhiyi, l1, l0, soulful))
|
||
else:
|
||
print(format_pretty(zhiyi, l1, l0, soulful))
|
||
|
||
# 三个系统全挂 → 退出码 1
|
||
if zhiyi.get("error") and l1.get("error") and l0.get("error") and soulful.get("error"):
|
||
raise SystemExit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |