294 lines
12 KiB
Python
294 lines
12 KiB
Python
"""NM2 试用台：直接跟「模型自带的记忆」打交道。
|
||
|
||
用法（在装了项目的环境里跑）：
|
||
|
||
python H:\\Memory\\agent_lab\\try_nm2.py --demo # 一键把整段故事演完
|
||
python H:\\Memory\\agent_lab\\try_nm2.py # 交互菜单，自己随便试
|
||
|
||
它只做一件事：把控制面的接口用中文包一层，让你看得见记忆层每一步到底做了什么 ——
|
||
路由觉得要不要用记忆、命中哪条记录、有没有注入证据、自动遗忘的概率是多少。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import sys
|
||
import time
|
||
import urllib.error
|
||
import urllib.parse
|
||
import urllib.request
|
||
|
||
DEFAULT_BASE = "http://127.0.0.1:8766/v1"
|
||
|
||
# 把接口里的英文状态翻成人话，看不懂的照原样显示。
|
||
STOP_REASON_ZH = {
|
||
"evidence_found": "找到证据，已注入",
|
||
"below_read_threshold": "分数低于读取门槛，不注入",
|
||
"router_abstained": "路由器判定这轮不需要记忆",
|
||
"attribute_not_covered": "库里没有这个属性（覆盖门拒答）",
|
||
"memory_disabled": "记忆已关闭",
|
||
"token_evidence_override": "词元证据否决，不注入",
|
||
}
|
||
|
||
STATUS_ZH = {
|
||
"active": "生效",
|
||
"superseded": "已被新版本取代",
|
||
"retracted": "已撤回",
|
||
"quarantined": "隔离中（低置信，未启用）",
|
||
}
|
||
|
||
|
||
class ServiceError(Exception):
|
||
"""控制面返回的错误，带上人话。"""
|
||
|
||
|
||
class Client:
|
||
def __init__(self, base: str, timeout: float = 600.0) -> None:
|
||
self.base = base.rstrip("/")
|
||
self.timeout = timeout
|
||
|
||
def _call(self, method: str, path: str, payload: dict | None = None, params: dict | None = None):
|
||
url = f"{self.base}{path}"
|
||
if params:
|
||
url = f"{url}?{urllib.parse.urlencode(params)}"
|
||
data = None
|
||
headers = {"Content-Type": "application/json; charset=utf-8"}
|
||
if payload is not None:
|
||
data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
||
request = urllib.request.Request(url, data=data, headers=headers, method=method)
|
||
try:
|
||
with urllib.request.urlopen(request, timeout=self.timeout) as response:
|
||
body = response.read().decode("utf-8", "replace")
|
||
except urllib.error.HTTPError as error:
|
||
detail = error.read().decode("utf-8", "replace")[:500]
|
||
raise ServiceError(f"接口 {path} 返回 HTTP {error.code}：{detail}") from None
|
||
except urllib.error.URLError as error:
|
||
raise ServiceError(
|
||
f"连不上 {self.base}（{error.reason}）。模型服务没在跑？"
|
||
) from None
|
||
return json.loads(body) if body.strip() else {}
|
||
|
||
# ---- 记忆层 ---------------------------------------------------------
|
||
def health(self) -> dict:
|
||
return self._call("GET", "/health")
|
||
|
||
def reset(self) -> dict:
|
||
return self._call("POST", "/nm2/reset", {})
|
||
|
||
def write(self, *, text: str, entity: str = "", attribute: str = "", value: str = "",
|
||
importance: float = 0.9, confidence: float = 0.95) -> dict:
|
||
return self._call("POST", "/nm2/write", {
|
||
"text": text, "entity": entity, "attribute": attribute, "value": value,
|
||
"importance": importance, "confidence": confidence,
|
||
})
|
||
|
||
def records(self, status: str = "active", limit: int = 100) -> list[dict]:
|
||
return self._call("GET", "/nm2/records", params={"status": status, "limit": limit}).get("records", [])
|
||
|
||
def ask(self, question: str, *, remember: bool = True, max_tokens: int = 256) -> dict:
|
||
"""问一轮。remember=False 时这一轮不读也不写记忆（当场做对照）。"""
|
||
|
||
body = {
|
||
"messages": [{"role": "user", "content": question}],
|
||
"max_tokens": max_tokens,
|
||
"stream": False,
|
||
"nm2": {"memory_enabled": remember},
|
||
}
|
||
return self._call("POST", "/chat/completions", body)
|
||
|
||
|
||
# ---------------------------------------------------------------- 显示
|
||
def rule(title: str = "") -> None:
|
||
line = "─" * 62
|
||
print(f"\n{line}")
|
||
if title:
|
||
print(title)
|
||
print(line)
|
||
|
||
|
||
def show_answer(reply: dict) -> str:
|
||
message = reply["choices"][0]["message"]
|
||
content = (message.get("content") or "").strip()
|
||
print(f"模型回答：{content or '（空）'}")
|
||
return content
|
||
|
||
|
||
def show_decision(reply: dict) -> None:
|
||
info = reply.get("nm2") or {}
|
||
decisions = info.get("router_decisions") or [{}]
|
||
decision = decisions[0] if decisions else {}
|
||
reason = str(decision.get("stop_reason", ""))
|
||
print("这一轮记忆层做了什么：")
|
||
print(f" · 要不要用记忆：{'要' if decision.get('need_memory') else '不要'}")
|
||
print(f" · 结论：{STOP_REASON_ZH.get(reason, reason or '（无）')}")
|
||
print(f" · 命中页：{', '.join(decision.get('page_ids') or []) or '（无）'}")
|
||
print(f" · 命中记录：{len(decision.get('record_ids') or [])} 条")
|
||
print(f" · 注入前缀：{info.get('prefix_tokens', 0)} 词元")
|
||
print(f" · 路由耗时：{round(float(info.get('read_seconds') or 0) * 1000)} 毫秒")
|
||
forget = info.get("auto_forget_probability") or [None]
|
||
if forget and forget[0] is not None:
|
||
mark = " ← 超过 0.9 会自动遗忘" if float(forget[0]) >= 0.9 else ""
|
||
print(f" · 自动遗忘概率：{float(forget[0]):.4f}{mark}")
|
||
|
||
|
||
def show_bank(client: Client, title: str = "记忆库现状") -> list[dict]:
|
||
active = client.records("active")
|
||
print(f"{title}（生效记录 {len(active)} 条）：")
|
||
for record in active:
|
||
origin = record.get("origin") or "（无来源标记）"
|
||
print(f" · [{record.get('status')}] {record.get('entity')}｜{record.get('attribute')}"
|
||
f" = {record.get('value')} 来源轮次 {origin}")
|
||
if not active:
|
||
print(" （空）")
|
||
return active
|
||
|
||
|
||
# ---------------------------------------------------------------- 演示
|
||
def demo(client: Client) -> None:
|
||
rule("第一步：清空记忆，确保从零开始")
|
||
client.reset()
|
||
print("已清空。")
|
||
|
||
rule("第二步：用户一口气说了两件事")
|
||
sentence = "记一下：我的紧急联系人是 王工，电话分机 7781。"
|
||
print(f"用户：{sentence}")
|
||
reply = client.ask(sentence)
|
||
show_answer(reply)
|
||
show_decision(reply)
|
||
|
||
print("\n（智能体随后按字段各写一条记录。这里直接调接口，效果等价——"
|
||
"两条记录会同属这一轮。）")
|
||
for attribute, value in (("紧急联系人", "王工"), ("电话分机", "7781")):
|
||
result = client.write(
|
||
text=f"{attribute} 是 {value}。", entity="user", attribute=attribute, value=value,
|
||
)
|
||
record = result.get("record", {})
|
||
print(f" 写入 {attribute}={value} → {result.get('action')}，"
|
||
f"来源轮次 {record.get('origin') or '（无）'}")
|
||
|
||
show_bank(client, "现在库里")
|
||
|
||
rule("第三步：用户要求忘掉紧急联系人信息")
|
||
forget = "请遗忘我的紧急联系人信息，不要再保留和使用它。"
|
||
print(f"用户：{forget}")
|
||
reply = client.ask(forget)
|
||
show_answer(reply)
|
||
show_decision(reply)
|
||
show_bank(client, "遗忘之后")
|
||
|
||
rule("第四步：追问分机号——这里就是以前会泄露的地方")
|
||
question = "我的紧急联系人分机是多少？如果已经不保留这条信息，就直接说明。"
|
||
print(f"用户：{question}")
|
||
reply = client.ask(question)
|
||
answer = show_answer(reply)
|
||
show_decision(reply)
|
||
|
||
leaked = "7781" in answer
|
||
print()
|
||
if leaked:
|
||
print("结果：分机号还是被说出来了 —— 这一轮没修好，把上面的路由结论发我。")
|
||
else:
|
||
print("结果：没有泄露分机号。同一轮写入的记录被一起撤回了。")
|
||
|
||
rule("第五步：换成「关掉记忆」再问一次，看对照组会怎样")
|
||
reply = client.ask(question, remember=False)
|
||
show_answer(reply)
|
||
print("（这一轮模型没有任何记忆可用，只能靠猜或明说不知道。）")
|
||
|
||
|
||
# ---------------------------------------------------------------- 交互
|
||
MENU = """
|
||
NM2 试用台
|
||
1) 一键演完整段故事（写入 → 提问 → 遗忘 → 追问 → 对照）
|
||
2) 写入一条事实
|
||
3) 提一个问题（带记忆）
|
||
4) 关掉记忆再问一次（当场对照）
|
||
5) 遗忘（直接对模型说「请遗忘…」，走真实自动遗忘）
|
||
6) 看记忆库
|
||
7) 清空记忆
|
||
0) 退出
|
||
"""
|
||
|
||
|
||
def interactive(client: Client) -> None:
|
||
while True:
|
||
print(MENU)
|
||
try:
|
||
choice = input("请选择：").strip()
|
||
except (EOFError, KeyboardInterrupt):
|
||
print()
|
||
return
|
||
try:
|
||
if choice == "1":
|
||
demo(client)
|
||
elif choice == "2":
|
||
entity = input("实体（直接回车=user）：").strip() or "user"
|
||
attribute = input("属性（例如 默认语言）：").strip()
|
||
value = input("取值（例如 葡萄牙语）：").strip()
|
||
text = input("记录原文（直接回车自动拼）：").strip() or f"{attribute} 是 {value}。"
|
||
result = client.write(text=text, entity=entity, attribute=attribute, value=value)
|
||
print(f"结果：{result.get('action')}")
|
||
show_bank(client)
|
||
elif choice == "3":
|
||
question = input("问题：").strip()
|
||
if not question:
|
||
continue
|
||
reply = client.ask(question)
|
||
show_answer(reply)
|
||
show_decision(reply)
|
||
elif choice == "4":
|
||
question = input("问题（这一轮会把记忆关掉）：").strip()
|
||
if not question:
|
||
continue
|
||
show_answer(client.ask(question, remember=False))
|
||
elif choice == "5":
|
||
text = input("对模型说：").strip() or "请遗忘我刚刚告诉你的那条信息。"
|
||
reply = client.ask(text)
|
||
show_answer(reply)
|
||
show_decision(reply)
|
||
show_bank(client, "遗忘之后")
|
||
elif choice == "6":
|
||
show_bank(client)
|
||
elif choice == "7":
|
||
client.reset()
|
||
print("已清空。")
|
||
elif choice == "0":
|
||
return
|
||
else:
|
||
print("没有这个选项。")
|
||
except ServiceError as error:
|
||
print(f"出错：{error}")
|
||
|
||
|
||
def main() -> int:
|
||
parser = argparse.ArgumentParser(description="NM2 试用台")
|
||
parser.add_argument("--base", default=DEFAULT_BASE, help="控制面地址")
|
||
parser.add_argument("--demo", action="store_true", help="直接跑自动演示，不进菜单")
|
||
args = parser.parse_args()
|
||
|
||
client = Client(args.base)
|
||
try:
|
||
health = client.health()
|
||
except ServiceError as error:
|
||
print(f"启动失败：{error}")
|
||
print("\n先起模型服务，再跑这个脚本。命令见 agent_lab/RESULTS.md 第 8 节。")
|
||
return 1
|
||
|
||
memory = health.get("memory", {})
|
||
print(f"已连上模型服务：{health.get('model_path')}")
|
||
print(f"当前记忆开关：{health.get('memory_mode')}｜"
|
||
f"生效记录 {memory.get('active_records')} 条 / 共 {memory.get('records')} 条")
|
||
|
||
if args.demo:
|
||
started = time.time()
|
||
demo(client)
|
||
print(f"\n演示结束，用了 {time.time() - started:.0f} 秒。")
|
||
else:
|
||
interactive(client)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|