Initial commit: Natural Memory Agent Lab：模型服务 + TypeScript agent + 记忆能力场景评测（冲突更新、跨会话回忆、遗忘撤回、多跳、未知拒答）
This commit is contained in:
WpyQwq
+293
@@ -0,0 +1,293 @@
|
||||
"""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())
|
||||
Reference in New Issue
Block a user