536 lines
20 KiB
Python
536 lines
20 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
统一记忆入口 (memory_recall.py)
|
||
|
||
同时查询织忆、TencentDB、Soulful 三套记忆系统，合并去重后输出统一上下文。
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import subprocess
|
||
import threading
|
||
import time
|
||
import urllib.error
|
||
import urllib.request
|
||
from pathlib import Path
|
||
|
||
# ── 常量 ──────────────────────────────────────────────────────────────────────
|
||
|
||
ZHIYI_URL = "http://localhost:7821/api/v1/recall"
|
||
ZHIYI_HEADERS = {
|
||
"X-API-Key": os.environ.get("ZHIYI_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 = 15
|
||
|
||
# ── 工具函数 ──────────────────────────────────────────────────────────────────
|
||
|
||
def http_post(url: str, payload: dict, headers: dict | None = None, timeout: int = TIMEOUT_SEC) -> dict | None:
|
||
"""POST JSON，返回解析后的 dict；失败自动重试最多3次，返回 None 并打印 warning。"""
|
||
delays = [1, 2, 4]
|
||
for attempt in range(4): # 首次 + 3次重试
|
||
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:
|
||
if attempt < 3:
|
||
print(f" [warning] {url} -> {e} (重试 {attempt+1}/3, 等待 {delays[attempt]}s)", flush=True)
|
||
time.sleep(delays[attempt])
|
||
else:
|
||
print(f" [warning] {url} -> {e} (已重试3次，放弃)", flush=True)
|
||
return None
|
||
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，yield 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)
|
||
|
||
|
||
# ── bge-embed keepalive ────────────────────────────────────────────────────────
|
||
|
||
BGE_HEALTH_URL = "http://localhost:8000/health"
|
||
_KEEPALIVE_RUNNING = True
|
||
|
||
|
||
def add_keepalive():
|
||
"""后台线程：每10秒ping bge-embed /health，防止服务因空闲超时断开。"""
|
||
global _KEEPALIVE_RUNNING
|
||
while _KEEPALIVE_RUNNING:
|
||
try:
|
||
req = urllib.request.Request(BGE_HEALTH_URL, method="GET")
|
||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||
resp.read() # 消费响应，无需解析
|
||
except Exception:
|
||
pass # bge-embed 暂时不可用也静默，ping 下次继续
|
||
time.sleep(10)
|
||
|
||
|
||
def start_keepalive():
|
||
"""启动 keepalive 守护线程（在 main 入口调用）。"""
|
||
t = threading.Thread(target=add_keepalive, daemon=True)
|
||
t.start()
|
||
|
||
|
||
def stop_keepalive():
|
||
"""停止 keepalive 线程。"""
|
||
global _KEEPALIVE_RUNNING
|
||
_KEEPALIVE_RUNNING = False
|
||
|
||
|
||
# ── 查询函数（供线程调用） ──────────────────────────────────────────────────────
|
||
|
||
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:
|
||
# 降级：织忆不可用时从 llm_context.json 读取缓存上下文
|
||
ctx = load_llm_context()
|
||
if ctx is not None:
|
||
print(f" [degrade] 织忆不可用，从 llm_context.json 读取缓存上下文", flush=True)
|
||
return {
|
||
"error": False,
|
||
"degraded": True,
|
||
"results": [],
|
||
"llm_context": ctx,
|
||
}
|
||
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,
|
||
}
|
||
|
||
|
||
# ── 统一读取入口（优先 llm_context.json）────────────────────────────────────────
|
||
|
||
LLM_CONTEXT_FILE = Path.home() / ".hermes" / "llm_context.json"
|
||
|
||
|
||
def load_llm_context() -> dict | None:
|
||
"""读取 llm_context.json（daemon 统一写入的合并上下文）。失败返回 None。"""
|
||
try:
|
||
if not LLM_CONTEXT_FILE.exists():
|
||
return None
|
||
return json.loads(LLM_CONTEXT_FILE.read_text(encoding="utf-8"))
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def query_llm_context() -> dict:
|
||
"""主查询：优先 llm_context.json，再补充织忆/L1/L0/soulful 详细数据。
|
||
|
||
返回格式同 parallel_query：(zhiyi, l1, l0, soulful, unified_context)
|
||
unified_context 包含 user_profile / active_scenes / short_term / cares / recent_moments
|
||
"""
|
||
ctx = load_llm_context()
|
||
soulful_data = load_soulful()
|
||
|
||
unified = {
|
||
"error": False,
|
||
"ctx": ctx, # llm_context v2 dict or None
|
||
"user_profile": (ctx or {}).get("user_profile", {}),
|
||
"active_scenes": (ctx or {}).get("active_scenes", []),
|
||
"short_term": (ctx or {}).get("short_term", {}),
|
||
"cares": soulful_data.get("cares", []), # 优先 soulful 原始
|
||
"heart_traces": soulful_data.get("heart_traces", []),
|
||
"distill_status": (ctx or {}).get("distill_status", "no_context"),
|
||
}
|
||
|
||
# 如果 ctx 有完整数据，织忆语义检索由主查询补充
|
||
# （此处不重复检索，llm_context.json 已含 daemon 蒸馏结果）
|
||
zhiyi = {"error": False, "results": []}
|
||
tddb_l1 = {"error": False, "results": []}
|
||
tddb_l0 = {"error": False, "results": []}
|
||
|
||
return zhiyi, tddb_l1, tddb_l0, soulful_data, unified
|
||
|
||
|
||
# ── 并行查询（fallback 模式）────────────────────────────────────────────────────
|
||
|
||
def parallel_query(query: str, top_k: int):
|
||
"""用三个线程同时查询三个后端 + Soulful 文件，返回 (zhiyi, l1, l0, soulful)。"""
|
||
|
||
def safe_update(result_dict, fn, *args, **kwargs):
|
||
"""安全执行查询，异常时标记 error"""
|
||
try:
|
||
result_dict.update(fn(*args, **kwargs))
|
||
except Exception as e:
|
||
result_dict["error"] = True
|
||
result_dict["results"] = []
|
||
|
||
threads = []
|
||
zhiyi_result = {}
|
||
l1_result = {}
|
||
l0_result = {}
|
||
soulful_result = {}
|
||
|
||
t1 = threading.Thread(target=lambda: safe_update(zhiyi_result, query_zhiyi, query, top_k))
|
||
t2 = threading.Thread(target=lambda: safe_update(l1_result, query_tencentdb_l1, query, top_k))
|
||
t3 = threading.Thread(target=lambda: safe_update(l0_result, query_tencentdb_l0, query, top_k))
|
||
t4 = threading.Thread(target=lambda: safe_update(soulful_result, 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 format_unified(ctx: dict):
|
||
"""格式化 llm_context.json v2 统一格式（方案A核心输出）。"""
|
||
lines = []
|
||
|
||
# user_profile
|
||
up = ctx.get("user_profile") or {}
|
||
if up:
|
||
lines.append("## 用户画像（L4-L5）")
|
||
if isinstance(up, dict):
|
||
for k, v in up.items():
|
||
if isinstance(v, dict):
|
||
lines.append(f" {k}:")
|
||
for k2, v2 in v.items():
|
||
lines.append(f" {k2}: {v2}")
|
||
elif isinstance(v, list):
|
||
lines.append(f" {k}: {', '.join(str(x) for x in v)}")
|
||
else:
|
||
lines.append(f" {k}: {v}")
|
||
else:
|
||
lines.append(f" {up}")
|
||
lines.append("")
|
||
|
||
# active_scenes
|
||
scenes = ctx.get("active_scenes") or []
|
||
if scenes:
|
||
lines.append(f"## 活跃场景（{len(scenes)} 个）")
|
||
for s in scenes[:5]:
|
||
name = s.get("name", s) if isinstance(s, dict) else str(s)
|
||
layer = s.get("layer", "?") if isinstance(s, dict) else "?"
|
||
lines.append(f" [{layer}] {name}")
|
||
lines.append("")
|
||
|
||
# short_term
|
||
st = ctx.get("short_term") or {}
|
||
if st:
|
||
lines.append("## 短期记忆（蒸馏计数）")
|
||
for k, v in st.items():
|
||
lines.append(f" {k}: {v}")
|
||
lines.append("")
|
||
|
||
# cares
|
||
cares = ctx.get("cares") or []
|
||
if cares:
|
||
lines.append(f"## 牵挂（{len(cares)} 条待关怀）")
|
||
for c in cares[:5]:
|
||
content = c.get("content", "")[:50]
|
||
due = c.get("due", c.get("follow_up_date", "?"))
|
||
lines.append(f" [{due}] {content}")
|
||
lines.append("")
|
||
|
||
# recent_moments
|
||
moments = ctx.get("recent_moments") or []
|
||
if moments:
|
||
lines.append(f"## 心迹（{len(moments)} 条）")
|
||
for m in moments[-3:]:
|
||
content = m.get("content", "")[:60]
|
||
ts = m.get("timestamp", "")[:10]
|
||
lines.append(f" [{ts}] {content}")
|
||
lines.append("")
|
||
|
||
# distill_status
|
||
ds = ctx.get("distill_status", "ok")
|
||
lines.append(f"蒸馏状态: {ds}")
|
||
|
||
return "\n".join(lines)
|
||
|
||
|
||
# ── 入口 ──────────────────────────────────────────────────────────────────────
|
||
|
||
def main():
|
||
start_keepalive() # 后台保持 bge-embed 连接活跃
|
||
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", "unified"], default="pretty", help="输出格式 (default: pretty)")
|
||
args = parser.parse_args()
|
||
|
||
if args.format == "unified":
|
||
# 新统一模式：读 llm_context.json（方案A）
|
||
_, _, _, _, unified = query_llm_context()
|
||
if unified.get("error"):
|
||
print("[memory] llm_context 不可用")
|
||
raise SystemExit(1)
|
||
ctx = unified.get("ctx")
|
||
if ctx is None:
|
||
print("[memory] llm_context.json 尚无数据（daemon 还未写入）")
|
||
raise SystemExit(1)
|
||
print(format_unified(ctx))
|
||
return
|
||
|
||
# legacy 模式：三系统并行查询
|
||
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() |