318 lines
18 KiB
Python
318 lines
18 KiB
Python
"""End-to-end self-test for the NM2 control plane.
|
|
|
|
Run this against a live `server.py` to verify, in one shot, that every NM2 surface
|
|
actually works on the trained package — the control plane was written against the
|
|
model's method signatures and must not be trusted until it has been exercised.
|
|
|
|
# terminal 1
|
|
set PYTHONPATH=H:\\Memory
|
|
C:\\Users\\Administrator\\miniconda3\\envs\\LLM\\python.exe agent_lab\\model_server\\server.py ^
|
|
--model-path qwen3_5_4b_natural_memory_v2_1 --port 8766
|
|
|
|
# terminal 2
|
|
python agent_lab\\model_server\\selftest.py --base-url http://127.0.0.1:8766/v1
|
|
|
|
Checks: health · coverage-gate state · reset · probe on an empty bank · trusted write ·
|
|
probe hit · coverage verdict · correction/versioning · retract · audit · KV compaction ·
|
|
session save/load/switch · config round-trip incl. rejection of a bad knob · memory-off
|
|
control arm · OpenAI chat completion with native tool calling.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
import time
|
|
import urllib.error
|
|
import urllib.parse
|
|
import urllib.request
|
|
|
|
RESULTS: list[tuple[str, bool, str]] = []
|
|
|
|
|
|
def record(name: str, ok: bool, detail: str = "") -> bool:
|
|
RESULTS.append((name, bool(ok), detail))
|
|
mark = "PASS" if ok else "FAIL"
|
|
print(f" [{mark}] {name}" + (f" -- {detail}" if detail else ""), flush=True)
|
|
return bool(ok)
|
|
|
|
|
|
class Client:
|
|
def __init__(self, base_url: str, timeout: float = 300.0) -> None:
|
|
self.base_url = base_url.rstrip("/")
|
|
self.root = self.base_url[:-3] if self.base_url.endswith("/v1") else self.base_url
|
|
self.timeout = timeout
|
|
|
|
def request(self, method: str, path: str, payload: dict | None = None, *, root: bool = False):
|
|
base = self.root if root else self.base_url
|
|
# Query strings may carry non-ASCII (Chinese) values; urllib encodes requests as
|
|
# ASCII, so the caller's path is percent-encoded here in one place.
|
|
safe_path = urllib.parse.quote(path, safe="/?=&%:")
|
|
url = f"{base}{safe_path}"
|
|
data = json.dumps(payload).encode("utf-8") if payload is not None else None
|
|
request = urllib.request.Request(url, data=data, method=method)
|
|
request.add_header("Content-Type", "application/json")
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=self.timeout) as response:
|
|
body = response.read().decode("utf-8")
|
|
return json.loads(body) if body else {}
|
|
except urllib.error.HTTPError as error:
|
|
detail = error.read().decode("utf-8", "replace")[:400]
|
|
raise RuntimeError(f"{method} {path} -> HTTP {error.code}: {detail}") from None
|
|
|
|
def get(self, path: str, *, root: bool = False):
|
|
return self.request("GET", path, root=root)
|
|
|
|
def post(self, path: str, payload: dict | None = None):
|
|
return self.request("POST", path, payload if payload is not None else {})
|
|
|
|
|
|
def wait_for_server(client: Client, seconds: float = 240.0) -> dict:
|
|
deadline = time.time() + seconds
|
|
last = ""
|
|
while time.time() < deadline:
|
|
try:
|
|
return client.get("/health", root=True)
|
|
except Exception as error: # still loading
|
|
last = str(error)
|
|
time.sleep(3)
|
|
raise SystemExit(f"server did not become ready: {last}")
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--base-url", default="http://127.0.0.1:8766/v1")
|
|
parser.add_argument("--wait", type=float, default=240.0)
|
|
args = parser.parse_args()
|
|
client = Client(args.base_url)
|
|
|
|
print("waiting for the model server ...", flush=True)
|
|
health = wait_for_server(client, args.wait)
|
|
print(json.dumps({"event": "healthy", "memory_mode": health.get("memory_mode"),
|
|
"records": (health.get("memory") or {}).get("records")}, ensure_ascii=False), flush=True)
|
|
record("health endpoint responds", True, f"mode={health.get('memory_mode')}")
|
|
|
|
# ---- coverage gate (NM2.1) -------------------------------------------
|
|
state = client.get("/nm2/state")
|
|
coverage = state["snapshot"].get("coverage") or {}
|
|
record("nm2/state exposes the coverage gate", "coverage" in state["snapshot"],
|
|
f"configured={coverage.get('configured')} loaded={coverage.get('loaded')} "
|
|
f"bound={coverage.get('bound_to_bank')}")
|
|
vocabulary = ((coverage.get("head") or {}) or {}).get("vocabulary") or []
|
|
record("attribute head loaded (NM2.1)", bool(coverage.get("loaded")),
|
|
f"attributes={len(vocabulary)}")
|
|
if not coverage.get("bound_to_bank"):
|
|
# A reset rebuilds the bank; rebind explicitly then re-check.
|
|
client.post("/nm2/config", {"knobs": {"coverage_gate": True}})
|
|
coverage = client.get("/nm2/state")["snapshot"]["coverage"]
|
|
record("coverage gate is bound to the live bank", bool(coverage.get("bound_to_bank")))
|
|
|
|
# ---- clean slate -------------------------------------------------------
|
|
client.post("/nm2/reset")
|
|
snapshot = client.get("/nm2/state")["snapshot"]
|
|
record("reset empties the bank", (snapshot["memory"].get("records") or 0) == 0,
|
|
f"records={snapshot['memory'].get('records')}")
|
|
|
|
empty = client.post("/nm2/probe", {"query": "我的项目代号是什么？"})
|
|
record("probe on an empty bank abstains", empty["need_memory"] is False,
|
|
f"stop={empty['stop_reason']}")
|
|
|
|
# ---- write / probe hit -------------------------------------------------
|
|
written = client.post("/nm2/write", {
|
|
"text": "记一下：我的项目代号是 蓝鲸-47，部署区域是 ap-southeast-3。",
|
|
"entity": "用户", "attribute": "项目代号", "value": "蓝鲸-47",
|
|
"importance": 0.95, "confidence": 0.99, "source": "selftest",
|
|
})
|
|
record("trusted write accepted", written.get("action") in {"inserted", "updated"},
|
|
f"action={written.get('action')}")
|
|
|
|
hit = client.post("/nm2/probe", {"query": "我的项目代号是什么？"})
|
|
matched = [r for r in hit["records"] if "蓝鲸-47" in (r.get("text") or "")]
|
|
record("probe retrieves the written record", hit["need_memory"] and bool(matched),
|
|
f"need={hit['need_memory']} records={len(hit['records'])} route_ms={hit['timing']['route_seconds']*1000:.2f}")
|
|
record("probe reports a router decision", hit["stop_reason"] is not None,
|
|
f"stop={hit['stop_reason']} hop={hit['hop_count']} margin={hit.get('score_margin')}")
|
|
record("probe reports the coverage verdict", hit.get("coverage") is not None,
|
|
json.dumps(hit.get("coverage"), ensure_ascii=False)[:120])
|
|
|
|
# ---- versioned correction ---------------------------------------------
|
|
record_id = matched[0]["record_id"] if matched else None
|
|
if record_id:
|
|
corrected = client.post("/nm2/correct", {
|
|
"record_id": record_id, "text": "更正：我的项目代号已改为 蓝鲸-99。",
|
|
"entity": "用户", "attribute": "项目代号", "value": "蓝鲸-99", "confidence": 1.0,
|
|
})
|
|
record("correction creates a new version", bool(corrected.get("corrected")))
|
|
active = client.get("/nm2/records?query=项目代号&status=active&limit=50")["records"]
|
|
superseded = client.get("/nm2/records?query=项目代号&status=superseded&limit=50")["records"]
|
|
record("old value superseded, new value active",
|
|
any("蓝鲸-99" in (r.get("text") or "") for r in active)
|
|
and any("蓝鲸-47" in (r.get("text") or "") for r in superseded),
|
|
f"active={len(active)} superseded={len(superseded)}")
|
|
|
|
# ---- retraction --------------------------------------------------------
|
|
active_now = client.get("/nm2/records?status=active&limit=50")["records"]
|
|
if active_now:
|
|
client.post("/nm2/retract", {"record_id": active_now[0]["record_id"]})
|
|
retracted = client.get("/nm2/records?status=retracted&limit=50")["records"]
|
|
record("retract moves a record out of active use", bool(retracted), f"retracted={len(retracted)}")
|
|
|
|
# ---- audit -------------------------------------------------------------
|
|
audit = client.get("/nm2/audit")
|
|
record("audit runs and reports health", isinstance(audit, dict) and bool(audit),
|
|
json.dumps({k: audit[k] for k in list(audit)[:4]}, ensure_ascii=False)[:140])
|
|
|
|
# ---- KV compaction -----------------------------------------------------
|
|
# Compaction is budget-driven: ``chunk_tokens`` only sets the chunk size once the
|
|
# trigger is reached. Drive it through the control plane by shrinking the budget
|
|
# instead of feeding a 30k-token context.
|
|
long_text = "这是一段用于压缩测试的长上下文。" * 200
|
|
idle = client.post("/nm2/compact", {"text": long_text})
|
|
record("KV compaction leaves a below-budget context intact",
|
|
idle.get("hot_tokens") == idle.get("input_tokens"),
|
|
f"{idle.get('input_tokens')} -> {idle.get('hot_tokens')} tokens (budget trigger not reached)")
|
|
applied_budget = client.post("/nm2/config", {"knobs": {"kv_budget_tokens": 1024, "context_chunk_tokens": 256}})
|
|
forced = client.post("/nm2/compact", {"text": long_text})
|
|
client.post("/nm2/config", {"knobs": {"kv_budget_tokens": 32768}})
|
|
report = forced.get("report") or {}
|
|
record("KV compaction chunks and archives once the budget is exceeded",
|
|
forced.get("hot_tokens", 0) < forced.get("input_tokens", 1) and (report.get("archived_records") or 0) > 0,
|
|
f"budget_applied={applied_budget.get('applied')} {forced.get('input_tokens')} -> {forced.get('hot_tokens')} tokens; "
|
|
f"archived={report.get('archived_records')} report={json.dumps(report, ensure_ascii=False)[:130]}")
|
|
|
|
# ---- session isolation -------------------------------------------------
|
|
client.post("/nm2/write", {"text": "用户 A 的专属事实：密钥是 A-ONLY-1。",
|
|
"entity": "用户A", "attribute": "密钥", "value": "A-ONLY-1"})
|
|
client.post("/nm2/session", {"action": "save", "user": "selftest_a"})
|
|
client.post("/nm2/reset")
|
|
after_reset = client.get("/nm2/records?status=active&limit=50")["records"]
|
|
record("reset clears before session load", len(after_reset) == 0, f"active={len(after_reset)}")
|
|
loaded = client.post("/nm2/session", {"action": "load", "user": "selftest_a"})
|
|
restored = client.get("/nm2/records?status=active&limit=50")["records"]
|
|
record("session load restores the user's memory",
|
|
bool(restored) and any("A-ONLY-1" in (r.get("text") or "") for r in restored),
|
|
f"active={len(restored)} loaded={loaded.get('loaded')}")
|
|
|
|
# ---- coverage gate (NM2.1 refusal mechanism) ---------------------------
|
|
# The gate only applies while the bank populates >=90% of the head's vocabulary,
|
|
# so the scenario is: fill the vocabulary except one attribute, then ask about the
|
|
# one that is missing (must be refused) and about a held one (must not be).
|
|
vocabulary = [str(v) for v in (vocabulary or []) if str(v).strip()]
|
|
if vocabulary:
|
|
hold_out = vocabulary[-1]
|
|
client.post("/nm2/reset")
|
|
client.post("/nm2/config", {"knobs": {"coverage_gate": True}})
|
|
for index, name in enumerate(vocabulary[:-1]):
|
|
client.post("/nm2/write", {
|
|
"text": f"记住：我的{name}是 V-{index:05d}。",
|
|
"entity": "覆盖率用户", "attribute": name, "value": f"V-{index:05d}",
|
|
})
|
|
state_after = client.get("/nm2/state")["snapshot"]["coverage"]
|
|
stats = ((state_after.get("head") or {}).get("coverage_stats") or {})
|
|
record("bank populated across the head vocabulary",
|
|
len(vocabulary) > 1,
|
|
f"wrote {len(vocabulary) - 1}/{len(vocabulary)} attributes; stats={stats}")
|
|
missing = client.post("/nm2/probe", {"query": f"我的{hold_out}是什么？"})
|
|
verdict = missing.get("coverage") or {}
|
|
record("coverage gate engages instead of standing down",
|
|
verdict.get("gate") != "bypassed", json.dumps(verdict, ensure_ascii=False)[:120])
|
|
|
|
# (a) A held attribute must still be retrievable while the gate is active.
|
|
held_records = 0
|
|
held_detail = ""
|
|
for phrasing in (f"我之前告诉过你的{vocabulary[0]}是什么？", f"我的{vocabulary[0]}的值是什么？"):
|
|
held = client.post("/nm2/probe", {"query": phrasing})
|
|
held_records = len(held["records"])
|
|
held_detail = (f"attr={vocabulary[0]} predicted={(held.get('coverage') or {}).get('attribute')} "
|
|
f"records={held_records} stop={held.get('stop_reason')}")
|
|
if held["need_memory"] and held_records:
|
|
break
|
|
record("coverage gate does not block an attribute that IS held", bool(held_records), held_detail)
|
|
|
|
# (b) Deterministic invariant, independent of how well the head names attributes:
|
|
# once the attribute the head predicts is absent from the bank, the gate must
|
|
# refuse. Force that state by retracting exactly the predicted attribute.
|
|
predicted = verdict.get("attribute")
|
|
active_rows = client.get("/nm2/records?status=active&limit=200")["records"]
|
|
bank_attributes = {row.get("attribute") for row in active_rows}
|
|
print(f" (note) head predicted {predicted!r} for {hold_out!r}; "
|
|
f"{'held -> gate allows it' if predicted in bank_attributes else 'absent -> gate should refuse'}", flush=True)
|
|
if predicted:
|
|
for row in active_rows:
|
|
if row.get("attribute") == predicted:
|
|
client.post("/nm2/retract", {"record_id": row["record_id"]})
|
|
after_retract = client.post("/nm2/probe", {"query": f"我的{hold_out}是什么？"})
|
|
record("coverage gate refuses once the predicted attribute is absent",
|
|
(not after_retract["need_memory"]) and after_retract["stop_reason"] == "attribute_not_covered",
|
|
f"retracted {predicted!r} -> stop={after_retract['stop_reason']} need={after_retract['need_memory']}")
|
|
else:
|
|
record("coverage gate refuses once the predicted attribute is absent", False,
|
|
"the head predicted no attribute for this phrasing")
|
|
|
|
# ---- config round-trip -------------------------------------------------
|
|
before = client.get("/nm2/config")["knobs"]
|
|
applied = client.post("/nm2/config", {"knobs": {"read_threshold": 0.42, "top_k_records": 3}})
|
|
after = client.get("/nm2/config")["knobs"]
|
|
record("config set/get round-trip",
|
|
abs(float(after["read_threshold"]) - 0.42) < 1e-6 and int(after["top_k_records"]) == 3,
|
|
f"read_threshold {before['read_threshold']} -> {after['read_threshold']}")
|
|
rejected = client.post("/nm2/config", {"knobs": {"definitely_not_a_knob": 1, "read_threshold": 9}})
|
|
record("invalid knobs are rejected, not silently applied",
|
|
"definitely_not_a_knob" in (rejected.get("rejected") or {})
|
|
and "read_threshold" in (rejected.get("rejected") or {}),
|
|
json.dumps(rejected.get("rejected"), ensure_ascii=False)[:140])
|
|
client.post("/nm2/config", {"knobs": {"read_threshold": 0.65, "top_k_records": 8}})
|
|
|
|
# ---- control arm -------------------------------------------------------
|
|
client.post("/nm2/mode", {"mode": "off"})
|
|
offline = client.post("/nm2/probe", {"query": "我的项目代号是什么？"})
|
|
record("memory_mode=off does not retrieve", not offline["records"],
|
|
f"records={len(offline['records'])}")
|
|
client.post("/nm2/mode", {"mode": "on"})
|
|
|
|
# ---- chat completion with native tool calling --------------------------
|
|
tools = [{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "nm2_read",
|
|
"description": "Inspect the model's own NM2 long-term memory (action=probe).",
|
|
"parameters": {"type": "object", "properties": {"action": {"type": "string"}}, "required": ["action"]},
|
|
},
|
|
}]
|
|
chat = client.post("/chat/completions", {
|
|
"model": "natural-memory-v2",
|
|
"messages": [{"role": "user", "content": "请调用 nm2_read 工具，action=probe，查询我的项目代号。"}],
|
|
"tools": tools,
|
|
"max_tokens": 128,
|
|
})
|
|
message = chat["choices"][0]["message"]
|
|
calls = message.get("tool_calls") or []
|
|
record("OpenAI chat completion returns a well-formed response",
|
|
chat.get("object") == "chat.completion" and "content" in message,
|
|
f"finish={chat['choices'][0]['finish_reason']} tokens={chat['usage']['total_tokens']}")
|
|
record("native tool calling works",
|
|
bool(calls) or bool(message.get("content", "").strip()),
|
|
f"tool_calls={[c['function']['name'] for c in calls]}")
|
|
record("per-turn NM2 diagnostics ride along with the completion",
|
|
isinstance(chat.get("nm2"), dict) and "router_decisions" in chat["nm2"],
|
|
f"prefix_tokens={chat.get('nm2', {}).get('prefix_tokens')} "
|
|
f"read_ms={(chat.get('nm2', {}).get('read_seconds') or 0) * 1000:.2f}")
|
|
|
|
traces = client.get("/nm2/trace?limit=5")["traces"]
|
|
record("trace history is queryable", len(traces) > 0, f"turns={len(traces)}")
|
|
|
|
# ---- summary -----------------------------------------------------------
|
|
passed = sum(1 for _, ok, _ in RESULTS if ok)
|
|
total = len(RESULTS)
|
|
print(f"\n{passed}/{total} checks passed", flush=True)
|
|
failures = [name for name, ok, _ in RESULTS if not ok]
|
|
if failures:
|
|
print("failed: " + "; ".join(failures), flush=True)
|
|
print(json.dumps({"passed": passed, "total": total, "failures": failures}, ensure_ascii=False), flush=True)
|
|
return 0 if not failures else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|