762 lines
34 KiB
Python
762 lines
34 KiB
Python
"""hermes-zhiyi — Bridge plugin: Hermes Agent → ZhiYi MemoryWeave REST API.
|
||
|
||
Uses ZhiYi's bge-m3 1024-dim semantic search + bge-reranker-v2-m3 reranking.
|
||
Writes go to ZhiYi via /commit; reads come from /recall.
|
||
|
||
Usage:
|
||
1. Set memory.provider: zhiyi in config.yaml
|
||
2. Set ZHIYI_URL in .env (e.g. http://localhost:7821)
|
||
|
||
Lifecycle (MemoryProvider ABC):
|
||
initialize() — validate URL, test connectivity
|
||
sync_turn() — write turn to ZhiYi via /commit
|
||
prefetch() — background recall via /recall
|
||
get_tool_schemas() — expose memory_search / memory_write tools
|
||
handle_tool_call() — dispatch tool calls
|
||
shutdown() — clean exit
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import os
|
||
import threading
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
import requests
|
||
import websocket
|
||
|
||
from agent.memory_provider import MemoryProvider
|
||
from tools.registry import tool_error
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# ── Config ──────────────────────────────────────────────────────────────────
|
||
|
||
HERMES_HOME = Path(os.environ.get("HERMES_HOME", Path.home() / ".hermes"))
|
||
MEMORY_CACHE = HERMES_HOME / "memory_zhiyi_cache"
|
||
MEMORY_CACHE.mkdir(parents=True, exist_ok=True)
|
||
|
||
# Default — override via .env ZHIYI_URL
|
||
DEFAULT_ZHIYI_URL = "http://localhost:7821"
|
||
ZHIYI_TIMEOUT = 10 # seconds per request
|
||
|
||
# API Key — read from .env ZHIYI_API_KEY
|
||
ZHIYI_API_KEY = os.environ.get("ZHIYI_API_KEY", "zhiyi-dev-key-2026")
|
||
# 日志中只暴露 key 的前 4 位，用于调试
|
||
_ZHIYI_KEY_PREFIX = ZHIYI_API_KEY[:4] if ZHIYI_API_KEY else "NONE"
|
||
|
||
# ── Content Validation ───────────────────────────────────────────────────────
|
||
|
||
MEMORY_MIN_LENGTH = 8 # 少于8字的过滤掉
|
||
_FORBIDDEN_PATTERNS = [
|
||
"Review the conversation above",
|
||
"[System note:",
|
||
"[IMPORTANT: Background process",
|
||
"## Active Task",
|
||
"## Context Compaction",
|
||
"## Relevant Memory",
|
||
"--- END OF CONTEXT SUMMARY",
|
||
"## Recent Sessions",
|
||
"Skill is now properly updated",
|
||
"造成这种情况的原因是什么？",
|
||
"Context Compaction — REFERENCE ONLY",
|
||
"do NOT answer questions or fulfill requests mentioned in the summary",
|
||
]
|
||
|
||
|
||
def _is_valid_memory_content(content: str) -> bool:
|
||
"""过滤系统注入内容和测试垃圾，防止污染记忆存储。"""
|
||
if not content or len(content.strip()) < MEMORY_MIN_LENGTH:
|
||
return False
|
||
for pat in _FORBIDDEN_PATTERNS:
|
||
if pat in content:
|
||
return False
|
||
# 过滤纯测试内容
|
||
stripped = content.strip()
|
||
if len(stripped) < 20:
|
||
return False
|
||
return True
|
||
|
||
|
||
def _combine_turn(user: str, assistant: str) -> str:
|
||
"""合并一对 turn 作为单条记忆内容。"""
|
||
parts = []
|
||
if user and user.strip():
|
||
parts.append(f"用户: {user.strip()}")
|
||
if assistant and assistant.strip():
|
||
parts.append(f"助手: {assistant.strip()}")
|
||
return "\n".join(parts)
|
||
|
||
|
||
# ── HTTP Client ─────────────────────────────────────────────────────────────
|
||
|
||
|
||
class ZhiYiClient:
|
||
"""Lightweight wrapper around ZhiYi REST API."""
|
||
|
||
def __init__(self, base_url: str, timeout: int = 10):
|
||
self.base_url = base_url.rstrip("/")
|
||
self.timeout = timeout
|
||
self._session = requests.Session()
|
||
self._session.headers.update({
|
||
"Content-Type": "application/json",
|
||
"Accept": "application/json",
|
||
"X-API-Key": ZHIYI_API_KEY,
|
||
})
|
||
|
||
def _url(self, path: str) -> str:
|
||
return f"{self.base_url}{path}"
|
||
|
||
def health(self) -> bool:
|
||
"""检查 ZhiYi 服务是否可达。"""
|
||
try:
|
||
r = self._session.get(self._url("/health"), timeout=3)
|
||
return r.status_code == 200
|
||
except Exception:
|
||
return False
|
||
|
||
def commit(self, content: str, category: str = "episodes", metadata: Optional[Dict] = None,
|
||
agent_id: str = "hermes-a06", namespace: Optional[str] = None) -> Optional[str]:
|
||
"""写入一条记忆到 ZhiYi。返回 commit_id 或 episode_id 或 None。失败自动重试一次。"""
|
||
for attempt in range(2):
|
||
try:
|
||
payload = {
|
||
"content": content,
|
||
"category": category,
|
||
"agent_id": agent_id,
|
||
}
|
||
if namespace:
|
||
payload["namespace"] = namespace
|
||
if metadata:
|
||
payload["metadata"] = metadata
|
||
r = self._session.post(self._url("/api/v1/commit"), json=payload, timeout=self.timeout)
|
||
if r.status_code in (200, 201):
|
||
data = r.json()
|
||
cid = (
|
||
data.get("commit_id")
|
||
or data.get("episode_id")
|
||
or data.get("distilled_id")
|
||
or data.get("id")
|
||
or (str(data) if data.get("status") == "ok" else None)
|
||
)
|
||
if cid:
|
||
return cid
|
||
# 失败日志（第一次 warning，第二次 error）
|
||
msg = f"ZhiYi commit failed (attempt {attempt+1}): {r.status_code} {r.text[:200]}"
|
||
if attempt == 0:
|
||
logger.warning("%s — retrying...", msg)
|
||
else:
|
||
logger.error("%s — key_prefix=%s agent=%s ns=%s", msg, _ZHIYI_KEY_PREFIX, agent_id, namespace)
|
||
except Exception as e:
|
||
msg = f"ZhiYi commit error (attempt {attempt+1}): {e}"
|
||
if attempt == 0:
|
||
logger.warning("%s — retrying...", msg)
|
||
else:
|
||
logger.error(msg)
|
||
return None
|
||
|
||
def recall(self, query: str, top_k: int = 5,
|
||
agent_id: str = "hermes-a06") -> List[Dict[str, Any]]:
|
||
"""语义搜索 ZhiYi 记忆。返回 [{content, score, category}, ...]。"""
|
||
try:
|
||
payload = {"query": query, "top_k": top_k, "use_rerank": True,
|
||
"agent_id": agent_id}
|
||
r = self._session.post(self._url("/api/v1/recall"), json=payload, timeout=self.timeout)
|
||
if r.status_code == 200:
|
||
data = r.json()
|
||
return data.get("results", [])
|
||
logger.warning("ZhiYi recall failed: %s %s", r.status_code, r.text[:200])
|
||
return []
|
||
except Exception as e:
|
||
logger.warning("ZhiYi recall error: %s", e)
|
||
return []
|
||
|
||
def stats(self) -> Dict[str, Any]:
|
||
"""获取 ZhiYi 统计信息。"""
|
||
try:
|
||
r = self._session.get(self._url("/api/v1/stats"), timeout=self.timeout)
|
||
if r.status_code == 200:
|
||
return r.json()
|
||
return {}
|
||
except Exception:
|
||
return {}
|
||
|
||
def mark_useful(self, memory_id: str) -> bool:
|
||
"""标记一条记忆为有用。"""
|
||
try:
|
||
r = self._session.post(self._url("/api/v1/feedback/useful"),
|
||
json={"memory_id": memory_id, "agent_id": "hermes-a06"},
|
||
timeout=self.timeout)
|
||
return r.status_code == 200
|
||
except Exception:
|
||
return False
|
||
|
||
def mark_not_useful(self, memory_id: str, reason: str = "") -> bool:
|
||
"""标记一条记忆为无用。"""
|
||
try:
|
||
r = self._session.post(self._url("/api/v1/feedback/not-useful"),
|
||
json={"memory_id": memory_id, "agent_id": "hermes-a06", "reason": reason},
|
||
timeout=self.timeout)
|
||
return r.status_code == 200
|
||
except Exception:
|
||
return False
|
||
|
||
def metrics(self) -> Dict[str, Any]:
|
||
"""获取自优化 7 项核心指标。"""
|
||
try:
|
||
r = self._session.get(self._url("/api/v1/metrics/self"), timeout=self.timeout)
|
||
if r.status_code == 200:
|
||
return r.json()
|
||
return {}
|
||
except Exception:
|
||
return {}
|
||
|
||
def search_notes(self, entity: str, max_hops: int = 2, max_notes: int = 5) -> List[Dict[str, Any]]:
|
||
"""图谱导航 + Obsidian 笔记关联搜索。"""
|
||
try:
|
||
params = {"entity": entity, "max_hops": max_hops, "max_notes": max_notes}
|
||
r = self._session.get(self._url("/api/v1/graph/notes"), params=params, timeout=self.timeout)
|
||
if r.status_code == 200:
|
||
data = r.json()
|
||
return data.get("notes", [])
|
||
return []
|
||
except Exception as e:
|
||
logger.warning("ZhiYi search_notes error: %s", e)
|
||
return []
|
||
|
||
def graph_navigate(self, entity: str, max_hops: int = 2,
|
||
namespace: Optional[str] = None) -> Dict[str, Any]:
|
||
"""图谱导航：查询实体的 N 跳关系网络。返回 {entity, count, paths}。"""
|
||
try:
|
||
payload = {"entity": entity, "max_hops": max_hops}
|
||
if namespace:
|
||
payload["namespace"] = namespace
|
||
r = self._session.post(self._url("/api/v1/graph/navigate"), json=payload, timeout=self.timeout)
|
||
if r.status_code == 200:
|
||
return r.json()
|
||
logger.warning("ZhiYi graph_navigate failed: %s %s", r.status_code, r.text[:200])
|
||
return {}
|
||
except Exception as e:
|
||
logger.warning("ZhiYi graph_navigate error: %s", e)
|
||
return {}
|
||
|
||
def graph_stats(self) -> Dict[str, Any]:
|
||
"""获取图谱统计：节点数、边数、密度。"""
|
||
try:
|
||
r = self._session.get(self._url("/api/v1/graph/stats"), timeout=self.timeout)
|
||
if r.status_code == 200:
|
||
return r.json()
|
||
return {}
|
||
except Exception as e:
|
||
logger.warning("ZhiYi graph_stats error: %s", e)
|
||
return {}
|
||
|
||
|
||
# ── Memory Provider ──────────────────────────────────────────────────────────
|
||
|
||
|
||
class HermesZhiYiMemoryProvider(MemoryProvider):
|
||
"""ZhiYi MemoryWeave bridge for Hermes Agent.
|
||
|
||
Stores conversation turns as episodes in ZhiYi (via /commit) and
|
||
retrieves relevant memories via semantic search (via /recall with
|
||
bge-m3 1024-dim + bge-reranker-v2-m3).
|
||
"""
|
||
|
||
def __init__(self):
|
||
self._client: Optional[ZhiYiClient] = None
|
||
self._session_id: str = ""
|
||
self._platform: str = "cli"
|
||
self._turn_counter: int = 0
|
||
self._write_queue: List[Dict] = []
|
||
self._queue_lock = threading.Lock()
|
||
self._prefetch_cache: str = "" # last prefetch result
|
||
self._prefetch_lock = threading.RLock()
|
||
self._started: bool = False
|
||
self._ws_thread: Optional[threading.Thread] = None
|
||
self._ws_running: bool = False
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "zhiyi"
|
||
|
||
# ── Availability ──────────────────────────────────────────────────────────
|
||
|
||
def is_available(self) -> bool:
|
||
"""检查 ZHIYI_URL 配置和服务连通性。"""
|
||
url = os.environ.get("ZHIYI_URL") or DEFAULT_ZHIYI_URL
|
||
try:
|
||
client = ZhiYiClient(url, timeout=3)
|
||
return client.health()
|
||
except Exception:
|
||
return False
|
||
|
||
# ── Lifecycle ─────────────────────────────────────────────────────────────
|
||
|
||
def initialize(self, session_id: str, **kwargs) -> None:
|
||
"""初始化 ZhiYi 客户端，注册当前 session。"""
|
||
self._session_id = session_id
|
||
self._platform = kwargs.get("platform", "cli")
|
||
url = os.environ.get("ZHIYI_URL") or DEFAULT_ZHIYI_URL
|
||
self._client = ZhiYiClient(url, timeout=ZHIYI_TIMEOUT)
|
||
self._started = True
|
||
self._turn_counter = 0
|
||
logger.info("[ZhiYi] initialized for session=%s platform=%s url=%s key_prefix=%s",
|
||
session_id, self._platform, url, _ZHIYI_KEY_PREFIX)
|
||
# 自检：验证 API Key 是否有效
|
||
try:
|
||
ok = self._client.health()
|
||
if ok:
|
||
logger.info("[ZhiYi] health check OK — key_prefix=%s", _ZHIYI_KEY_PREFIX)
|
||
else:
|
||
logger.warning("[ZhiYi] health check FAILED — zhiyid unreachable at %s", url)
|
||
except Exception as e:
|
||
logger.warning("[ZhiYi] health check error: %s", e)
|
||
# 启动 WebSocket 事件监听
|
||
self._start_ws_listener()
|
||
|
||
def _start_ws_listener(self) -> None:
|
||
"""启动 WebSocket 事件监听线程（连接织忆事件推送）。"""
|
||
if self._ws_running:
|
||
return
|
||
self._ws_running = True
|
||
self._ws_thread = threading.Thread(target=self._ws_listen, daemon=True, name="zhiyi-ws")
|
||
self._ws_thread.start()
|
||
logger.info("[ZhiYi] WS listener started")
|
||
|
||
def _ws_listen(self) -> None:
|
||
"""WebSocket 事件监听循环（自动重连）。"""
|
||
agent_id = os.environ.get("ZHIYI_AGENT_ID", "hermes-a06")
|
||
ws_url = (os.environ.get("ZHIYI_URL") or DEFAULT_ZHIYI_URL).replace("http://", "ws://").replace("https://", "wss://")
|
||
ws_url = f"{ws_url}/api/v1/ws/{agent_id}"
|
||
|
||
def on_message(ws_app, message):
|
||
try:
|
||
data = json.loads(message)
|
||
event_type = data.get("event", data.get("type", "unknown"))
|
||
logger.info("[ZhiYi] WS event: %s", event_type)
|
||
# prefetch.push 事件更新预取缓存
|
||
if event_type == "prefetch.push":
|
||
memories = data.get("memories", [])
|
||
if memories:
|
||
prefetch_text = "\n[ZhiYi Prefetch — 预取推送]\n"
|
||
for m in memories:
|
||
prefetch_text += f" [{m.get('score', 0):.2f}] {m.get('content', '')[:200]}\n"
|
||
with self._prefetch_lock:
|
||
self._prefetch_cache = prefetch_text
|
||
except json.JSONDecodeError:
|
||
pass
|
||
|
||
def on_error(ws_app, error):
|
||
logger.warning("[ZhiYi] WS error: %s", error)
|
||
|
||
def on_close(ws_app, close_status_code, close_msg):
|
||
logger.info("[ZhiYi] WS closed: %s %s", close_status_code, close_msg)
|
||
|
||
def on_open(ws_app):
|
||
logger.info("[ZhiYi] WS connected to %s", ws_url)
|
||
|
||
while self._ws_running:
|
||
try:
|
||
ws_app = websocket.WebSocketApp(
|
||
ws_url,
|
||
on_open=on_open,
|
||
on_message=on_message,
|
||
on_error=on_error,
|
||
on_close=on_close,
|
||
header={"X-API-Key": ZHIYI_API_KEY},
|
||
)
|
||
ws_app.run_forever(ping_interval=30, ping_timeout=10)
|
||
except Exception as e:
|
||
logger.warning("[ZhiYi] WS exception: %s", e)
|
||
# 重连延迟
|
||
for _ in range(30):
|
||
if not self._ws_running:
|
||
break
|
||
time.sleep(1)
|
||
|
||
def shutdown(self) -> None:
|
||
"""flush pending writes before exit."""
|
||
self._flush_queue()
|
||
self._ws_running = False
|
||
self._started = False
|
||
logger.info("[ZhiYi] shutdown")
|
||
|
||
# ── Write path ───────────────────────────────────────────────────────────
|
||
|
||
def sync_turn(self, user_content: str, assistant_content: str, *, session_id: str = "") -> None:
|
||
"""每轮对话结束后调用：将 turn 写入队列，异步 commit 到 ZhiYi。"""
|
||
if not self._started or not self._client:
|
||
return
|
||
|
||
combined = _combine_turn(user_content, assistant_content)
|
||
if not _is_valid_memory_content(combined):
|
||
return
|
||
|
||
self._turn_counter += 1
|
||
entry = {
|
||
"content": combined,
|
||
"category": "episodes",
|
||
"metadata": {
|
||
"session_id": session_id or self._session_id,
|
||
"turn": self._turn_counter,
|
||
"platform": self._platform,
|
||
},
|
||
}
|
||
|
||
with self._queue_lock:
|
||
self._write_queue.append(entry)
|
||
|
||
# 队列超过一定量就 flush，避免积压
|
||
if len(self._write_queue) >= 3:
|
||
self._flush_queue()
|
||
|
||
def _flush_queue(self) -> None:
|
||
"""将队列中的记忆批量写入 ZhiYi。"""
|
||
with self._queue_lock:
|
||
if not self._write_queue:
|
||
return
|
||
entries = self._write_queue[:]
|
||
self._write_queue.clear()
|
||
|
||
# 从 metadata 中读取 agent_id（sync_turn 写入时携带），否则用默认值
|
||
for entry in entries:
|
||
meta = entry.get("metadata", {}) or {}
|
||
self._client.commit(
|
||
content=entry["content"],
|
||
category=entry.get("category", "episodes"),
|
||
metadata=meta,
|
||
agent_id=meta.get("agent_id", "hermes-a06"),
|
||
namespace=meta.get("namespace"),
|
||
)
|
||
|
||
# ── Read path ────────────────────────────────────────────────────────────
|
||
|
||
def prefetch(self, query: str, *, session_id: str = "") -> str:
|
||
"""每次 API 调用前触发：执行语义搜索 + 图谱导航 Obsidian 笔记，返回最相关记忆。"""
|
||
if not self._client or not query or len(query.strip()) < 2:
|
||
return ""
|
||
|
||
blocks = ["[ZhiYi Memory — relevant past context]"]
|
||
|
||
# 语义搜索
|
||
results = self._client.recall(query.strip(), top_k=3,
|
||
agent_id=os.environ.get("ZHIYI_AGENT_ID", "hermes-a06"))
|
||
for r in results:
|
||
score = r.get("score", 0)
|
||
content = r.get("content", "") or r.get("text", "")
|
||
cat = r.get("category", "")
|
||
if content:
|
||
blocks.append(f" [{score:.2f}][{cat}] {content[:500]}")
|
||
|
||
# 图谱导航 + Obsidian 笔记（从 query 提取关键词作为实体）
|
||
notes = self._client.search_notes(query.strip(), max_hops=2, max_notes=3)
|
||
if notes:
|
||
blocks.append("\n[ZhiYi Graph — related Obsidian notes]")
|
||
for n in notes:
|
||
title = n.get("title", "无标题")
|
||
path = n.get("path", "")
|
||
snippet = n.get("snippet", "")[:200]
|
||
score = n.get("score", 0)
|
||
entities = ", ".join(n.get("entities", []))
|
||
blocks.append(f" [{score:.0f}] {title} ({path})")
|
||
blocks.append(f" \"{snippet}\"")
|
||
if entities:
|
||
blocks.append(f" via entities: {entities}")
|
||
|
||
text = "\n".join(blocks)
|
||
|
||
# 合并 WebSocket prefetch.push 推送事件
|
||
with self._prefetch_lock:
|
||
if self._prefetch_cache and "[ZhiYi Prefetch" in self._prefetch_cache:
|
||
text += "\n" + self._prefetch_cache
|
||
|
||
with self._prefetch_lock:
|
||
self._prefetch_cache = text
|
||
return text
|
||
|
||
def queue_prefetch(self, query: str, *, session_id: str = "") -> None:
|
||
"""空实现 — prefetch 已是同步的，不需要额外的异步队列。"""
|
||
pass
|
||
|
||
# ── Tool interface ────────────────────────────────────────────────────────
|
||
|
||
def get_tool_schemas(self) -> List[Dict[str, Any]]:
|
||
return [
|
||
{
|
||
"name": "memory_search",
|
||
"description": "Search ZhiYi semantic memory for relevant past information. Use for: user preferences, earlier decisions, established facts, project context.",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"query": {
|
||
"type": "string",
|
||
"description": "Semantic search query describing what to find. Be specific — the query is encoded with bge-m3 1024-dim and matched against all memories.",
|
||
},
|
||
"top_k": {
|
||
"type": "integer",
|
||
"description": "Maximum number of results to return (default: 5, max: 20).",
|
||
"default": 5,
|
||
},
|
||
},
|
||
"required": ["query"],
|
||
},
|
||
},
|
||
{
|
||
"name": "memory_write",
|
||
"description": "Explicitly save an important fact or preference to ZhiYi memory. Use when: user tells you something important, you discover a key fact, user confirms a decision.",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"content": {
|
||
"type": "string",
|
||
"description": "The memory content to save. Should be a complete, meaningful statement. Be specific — include context.",
|
||
},
|
||
"category": {
|
||
"type": "string",
|
||
"description": "Category for this memory: 'distilled' (important facts), 'episodes' (conversation turns). Default: 'distilled'.",
|
||
"default": "distilled",
|
||
},
|
||
},
|
||
"required": ["content"],
|
||
},
|
||
},
|
||
{
|
||
"name": "memory_stats",
|
||
"description": "Get ZhiYi memory statistics: total documents, vector dimensions, index status.",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {},
|
||
},
|
||
},
|
||
{
|
||
"name": "memory_feedback",
|
||
"description": "Mark a memory as useful or not-useful. Feedback activates ZhiYi's self-optimization — useful memories get higher priority, not-useful memories trigger quality monitoring and potential deprecation.",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"memory_id": {
|
||
"type": "string",
|
||
"description": "The memory ID to provide feedback on (from memory_search results).",
|
||
},
|
||
"useful": {
|
||
"type": "boolean",
|
||
"description": "True = this memory was helpful, False = this memory was not helpful.",
|
||
},
|
||
"reason": {
|
||
"type": "string",
|
||
"description": "Why this memory was not useful (optional, only used when useful=false).",
|
||
},
|
||
},
|
||
"required": ["memory_id", "useful"],
|
||
},
|
||
},
|
||
{
|
||
"name": "memory_metrics",
|
||
"description": "Get ZhiYi self-optimization metrics: recall hit rate, usefulness rate, gap closure rate, distillation loss, and more. Use this to check memory system health.",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {},
|
||
},
|
||
},
|
||
{
|
||
"name": "memory_graph_navigate",
|
||
"description": "Navigate the ZhiYi knowledge graph to understand relationships around an entity. Shows N-hop relationship network — what concepts, people, or projects are connected, and how. Use when: analyzing project structure, understanding who/what is related to a topic, exploring context dependencies.",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"entity": {
|
||
"type": "string",
|
||
"description": "The central entity to navigate from (e.g. '牧尘', '织忆', '小唯', 'openclaw'). Chinese and English both work.",
|
||
},
|
||
"max_hops": {
|
||
"type": "integer",
|
||
"description": "Maximum path depth (default: 2). 1-hop = direct neighbors only; 2-hop = friends-of-friends. Larger values return more paths but are slower.",
|
||
"default": 2,
|
||
},
|
||
},
|
||
"required": ["entity"],
|
||
},
|
||
},
|
||
{
|
||
"name": "memory_graph_stats",
|
||
"description": "Get ZhiYi knowledge graph statistics: total nodes, edges, and graph density. Use to check graph size and health.",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {},
|
||
},
|
||
},
|
||
]
|
||
|
||
def handle_tool_call(self, tool_name: str, args: Dict[str, Any], **kwargs) -> str:
|
||
if tool_name == "memory_search":
|
||
return self._tool_memory_search(args.get("query", ""), args.get("top_k", 5))
|
||
elif tool_name == "memory_write":
|
||
return self._tool_memory_write(args.get("content", ""), args.get("category", "distilled"))
|
||
elif tool_name == "memory_feedback":
|
||
return self._tool_memory_feedback(
|
||
args.get("memory_id", ""), args.get("useful", True), args.get("reason", ""))
|
||
elif tool_name == "memory_metrics":
|
||
return self._tool_memory_metrics()
|
||
elif tool_name == "memory_stats":
|
||
return self._tool_memory_stats()
|
||
elif tool_name == "memory_graph_navigate":
|
||
return self._tool_memory_graph_navigate(
|
||
args.get("entity", ""), args.get("max_hops", 2))
|
||
elif tool_name == "memory_graph_stats":
|
||
return self._tool_memory_graph_stats()
|
||
return tool_error(tool_name, "Unknown tool")
|
||
|
||
# ── Tool handlers ─────────────────────────────────────────────────────────
|
||
|
||
def _tool_memory_search(self, query: str, top_k: int) -> str:
|
||
if not self._client:
|
||
return json.dumps({"success": False, "error": "ZhiYi client not initialized"})
|
||
if top_k > 20:
|
||
top_k = 20
|
||
results = self._client.recall(query, top_k=top_k)
|
||
if not results:
|
||
return json.dumps({"success": True, "results": [], "message": "No relevant memories found"})
|
||
formatted = []
|
||
for r in results:
|
||
content = r.get("content", "") or r.get("text", "")
|
||
formatted.append({
|
||
"content": content[:1000],
|
||
"score": round(r.get("score", 0), 4),
|
||
"category": r.get("category", ""),
|
||
"quality_score": round(r.get("quality_score", 0), 4),
|
||
"v_value": round(r.get("v_value", 0), 4),
|
||
})
|
||
return json.dumps({"success": True, "count": len(formatted), "results": formatted})
|
||
|
||
def _tool_memory_write(self, content: str, category: str) -> str:
|
||
if not self._client:
|
||
return json.dumps({"success": False, "error": "ZhiYi client not initialized"})
|
||
if not _is_valid_memory_content(content):
|
||
return json.dumps({"success": False, "error": "Content too short or invalid"})
|
||
commit_id = self._client.commit(
|
||
content=content, category=category,
|
||
agent_id=os.environ.get("ZHIYI_AGENT_ID", "hermes-a06"),
|
||
)
|
||
if commit_id:
|
||
return json.dumps({"success": True, "commit_id": commit_id})
|
||
return json.dumps({"success": False, "error": "Commit failed — check ZhiYi server logs"})
|
||
|
||
def _tool_memory_stats(self) -> str:
|
||
if not self._client:
|
||
return json.dumps({"success": False, "error": "ZhiYi client not initialized"})
|
||
stats = self._client.stats()
|
||
if not stats:
|
||
return json.dumps({"success": False, "error": "Could not reach ZhiYi /stats endpoint"})
|
||
return json.dumps({"success": True, "stats": stats})
|
||
|
||
def _tool_memory_feedback(self, memory_id: str, useful: bool, reason: str = "") -> str:
|
||
if not self._client:
|
||
return json.dumps({"success": False, "error": "ZhiYi client not initialized"})
|
||
if not memory_id:
|
||
return json.dumps({"success": False, "error": "memory_id required"})
|
||
if useful:
|
||
ok = self._client.mark_useful(memory_id)
|
||
else:
|
||
ok = self._client.mark_not_useful(memory_id, reason)
|
||
if ok:
|
||
return json.dumps({"success": True, "memory_id": memory_id, "useful": useful,
|
||
"message": "Feedback recorded — VProp + Dashboard updated"})
|
||
return json.dumps({"success": False, "error": "ZhiYi feedback API failed"})
|
||
|
||
def _tool_memory_metrics(self) -> str:
|
||
if not self._client:
|
||
return json.dumps({"success": False, "error": "ZhiYi client not initialized"})
|
||
metrics = self._client.metrics()
|
||
if not metrics:
|
||
return json.dumps({"success": False, "error": "Could not reach ZhiYi /metrics/self"})
|
||
return json.dumps({"success": True, "metrics": metrics})
|
||
|
||
def _tool_memory_graph_navigate(self, entity: str, max_hops: int) -> str:
|
||
"""图谱导航：查询实体的 N 跳关系网络。"""
|
||
if not self._client:
|
||
return json.dumps({"success": False, "error": "ZhiYi client not initialized"})
|
||
if not entity or len(entity.strip()) < 1:
|
||
return json.dumps({"success": False, "error": "entity is required"})
|
||
if max_hops < 1:
|
||
max_hops = 1
|
||
if max_hops > 5:
|
||
max_hops = 5
|
||
|
||
result = self._client.graph_navigate(entity.strip(), max_hops=max_hops)
|
||
if not result:
|
||
return json.dumps({"success": False, "error": "Graph navigate failed — check ZhiYi server logs"})
|
||
|
||
# 格式化路径输出，保留关键信息
|
||
paths = result.get("paths", [])
|
||
formatted_paths = []
|
||
for p in paths[:50]: # 最多返回50条路径，避免太长
|
||
formatted_paths.append({
|
||
"from": p.get("from", ""),
|
||
"relation": p.get("relation", ""),
|
||
"to": p.get("to", ""),
|
||
"hop": p.get("hop", 1),
|
||
"weight": round(p.get("weight", 0), 3),
|
||
})
|
||
|
||
return json.dumps({
|
||
"success": True,
|
||
"entity": result.get("entity", entity),
|
||
"bidirectional": result.get("bidirectional", False),
|
||
"count": result.get("count", len(formatted_paths)),
|
||
"paths": formatted_paths,
|
||
"message": f"Found {result.get('count', 0)} paths within {max_hops}-hop network",
|
||
})
|
||
|
||
def _tool_memory_graph_stats(self) -> str:
|
||
"""获取图谱统计：节点数、边数、密度。"""
|
||
if not self._client:
|
||
return json.dumps({"success": False, "error": "ZhiYi client not initialized"})
|
||
stats = self._client.graph_stats()
|
||
if not stats:
|
||
return json.dumps({"success": False, "error": "Could not reach ZhiYi /graph/stats"})
|
||
return json.dumps({
|
||
"success": True,
|
||
"stats": stats,
|
||
"message": f"Graph has {stats.get('node_count', 0)} nodes and {stats.get('edge_count', 0)} edges (density: {stats.get('density', 0):.4f})",
|
||
})
|
||
|
||
# ── Session management ────────────────────────────────────────────────────
|
||
|
||
def on_session_switch(self, new_session_id: str, *, parent_session_id: str = "", reset: bool = False, **kwargs) -> None:
|
||
self._session_id = new_session_id
|
||
self._turn_counter = 0
|
||
if reset:
|
||
self._flush_queue()
|
||
|
||
def on_turn_start(self, turn_number: int, message: str, **kwargs) -> None:
|
||
"""每轮开始时 flush 待写入的记忆，保证顺序。"""
|
||
self._flush_queue()
|
||
|
||
# ── System prompt ────────────────────────────────────────────────────────
|
||
|
||
def system_prompt_block(self) -> str:
|
||
return (
|
||
"\\n[ZhiYi Memory] You have access to ZhiYi MemoryWeave — a semantic memory system with knowledge graph.\\n"
|
||
"Your prefetch automatically retrieves: (1) relevant semantic memories, (2) Obsidian notes related via graph navigation.\\n"
|
||
"Tools available:\\n"
|
||
" memory_search — semantic search for relevant past information\\n"
|
||
" memory_write — save an important fact or preference\\n"
|
||
" memory_feedback — mark memories as useful/not-useful (drives self-optimization)\\n"
|
||
" memory_metrics — check memory system health (recall hit rate, gap closures, etc.)\\n"
|
||
" memory_stats — get memory statistics (total docs, vector dimensions)\\n"
|
||
" memory_graph_navigate — navigate the knowledge graph to understand entity relationships (N-hop network)\\n"
|
||
" memory_graph_stats — get knowledge graph statistics (nodes, edges, density)\\n"
|
||
)
|
||
|
||
|
||
# ── Plugin entry point ──────────────────────────────────────────────────────
|
||
|
||
|
||
def register(ctx) -> None:
|
||
"""Called by Hermes plugin system to register this memory provider."""
|
||
ctx.register_memory_provider(HermesZhiYiMemoryProvider()) |