"""Complete NM2 control plane + OpenAI-compatible model server.
The project's `natural_memory_service.py` exposes a single-message chat endpoint and
a small record-admin API. An agent that must *operate* NM2 rather than merely talk
through it needs three things that endpoint cannot provide:
1. **Observation** -- the router's actual decision per turn. The model already records
it (`runtime.v2_last_decisions`, and the policy head's write/forget probabilities,
injected prefix length, routing-only latency, and KV-compaction result) but nothing
exposes it. This server snapshots that runtime state after every call and keeps a
per-turn trace history.
2. **Control** -- read/write thresholds, top-k pages/records, hop limit, grounding
guards, gpu-cache and KV budgets, mutable at runtime; plus explicit
write / correct / retract / approve(quarantine) / reset and per-user session
save/load so multi-user isolation can be tested.
3. **Routing as a first-class operation** -- a probe endpoint that runs *only* the
bounded coarse-index -> page -> record -> Top-K route and returns the decision plus
the retrieved records, with no generation. That turns "what would memory give me
here" into a tool the agent can call and a signal the harness can score.
Chat completions additionally accept an `nm2` block for per-request overrides, and
`memory_mode` (on / read_only / off) still provides the control arm of the A/B.
Language/process isolation is intentional: this file runs in the model's own conda
environment (torch/bitsandbytes must not be touched by agent dependencies); the agent
side is a separate Node project speaking HTTP.
"""
from __future__ import annotations
import argparse
import json
import re
import sys
import threading
import time
import uuid
from collections import deque
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any
from urllib.parse import parse_qs, urlparse
import torch
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from V2_dpskw.memory_os_v2 import MemoryOSV2 # noqa: E402
from V2_dpskw.natural_memory_service import NaturalMemoryService # noqa: E402
from V2_dpskw.qwen_integration import compose_corrected_evidence, memory_origin # noqa: E402
from V2_dpskw.stream_chat_qwen_memory import _write_turn # noqa: E402
TOOL_CALL_RE = re.compile(r"\s*\n]+)>(.*?)\s*", re.S)
PARAM_RE = re.compile(r"\n]+)>\n?(.*?)\n?", re.S)
MEMORY_MODES = ("on", "read_only", "off")
#: POST endpoints that store, edit or import long-term memory. They are refused unless
#: the mode permits writing, so the A/B control arm really has no memory layer: the mode
#: switch itself (`/v1/memory/mode`, `/v1/nm2/mode`) and reset are deliberately excluded.
#: Read paths stay open -- with every mutation refused, the bank is empty by
#: construction, and keeping reads available lets diagnostics prove that.
_MEMORY_MUTATIONS = frozenset({
"/v1/nm2/write",
"/v1/nm2/correct",
"/v1/nm2/retract",
"/v1/nm2/approve",
"/v1/nm2/compact",
"/v1/nm2/session",
"/v1/memory",
})
#: Runtime-mutable knobs. Each entry maps the JSON key an agent sends to where the
#: value actually lives, plus a validator. Keeping an explicit allowlist means an
#: agent can steer NM2 without being able to corrupt model state.
CONFIG_KNOBS: dict[str, dict[str, Any]] = {
"top_k_pages": {"config": "memory_top_k_pages", "type": int, "min": 1, "max": 64},
"top_k_records": {"config": "memory_top_k_records", "type": int, "min": 1, "max": 128},
"max_hops": {"config": "memory_max_hops", "type": int, "min": 1, "max": 3},
"hot_pages": {"config": "memory_hot_pages", "type": int, "min": 0, "max": 256},
"read_threshold": {"runtime": "read_threshold", "type": float, "min": 0.0, "max": 1.0},
"write_threshold": {"runtime": "write_threshold", "type": float, "min": 0.0, "max": 1.0},
"min_read_margin": {"runtime": "min_read_margin", "type": float, "min": 0.0, "max": 10.0},
"require_evidence": {"runtime": "require_evidence", "type": bool},
"auto_memory_threshold": {"config": "auto_memory_threshold", "type": float, "min": 0.0, "max": 1.0},
"auto_forget_threshold": {"config": "auto_forget_threshold", "type": float, "min": 0.0, "max": 1.0},
"gpu_cache_records": {"config": "memory_gpu_cache_records", "type": int, "min": 0, "max": 4096},
"gpu_cache_tokens": {"config": "memory_gpu_cache_tokens", "type": int, "min": 0, "max": 1_048_576},
"context_chunk_tokens": {"config": "context_chunk_tokens", "type": int, "min": 64, "max": 8192},
"kv_budget_tokens": {"config": "kv_budget_tokens", "type": int, "min": 1024, "max": 262144},
"kv_keep_recent_tokens": {"config": "kv_keep_recent_tokens", "type": int, "min": 128, "max": 131072},
# NM2.1 coverage gate: refuses a question whose asked-about attribute the bank
# does not hold. Toggling needs the head (re)loaded, so it is handled specially.
"coverage_gate": {"special": "coverage_gate", "type": bool},
"coverage_vocabulary_fraction": {"config": "memory_coverage_vocabulary_fraction", "type": float, "min": 0.0, "max": 1.0},
}
def _jsonable(value: Any) -> Any:
if isinstance(value, torch.Tensor):
return value.detach().cpu().tolist()
if isinstance(value, dict):
return {str(key): _jsonable(item) for key, item in value.items()}
if isinstance(value, (list, tuple)):
return [_jsonable(item) for item in value]
if isinstance(value, (str, int, float, bool)) or value is None:
return value
return str(value)
def split_thinking(raw: str) -> tuple[str, str]:
"""Split a leading ``…`` block into (reasoning, answer).
This model emits a visible reasoning trace before its answer. The chat template
itself treats a leading think block as ``reasoning_content`` for the *next* turn, so
the server must do the same: otherwise an agent receives a wall of "Thinking Process:"
as the assistant message and the actual answer (or tool call) is buried — or, when the
token budget runs out mid-thought, never produced at all.
With thinking disabled the template pre-closes an empty ```` block, and
the checkpoint sometimes echoes the enclosing role turn back out; those artifacts are
stripped here so what reaches an agent is the answer.
"""
text = raw.lstrip()
reasoning = ""
if text.startswith(""):
head, separator, tail = text.partition("")
if not separator:
return head[len(""):].strip(), ""
reasoning, text = head[len(""):].strip(), tail
text = re.sub(r"\s*", "", text)
text = re.sub(r"^(assistant|user|system)\s*\n", "", text.lstrip())
return reasoning, text.strip()
def parse_qwen_tool_calls(text: str) -> tuple[str, list[dict[str, Any]]]:
"""Split model output into visible content and OpenAI-shaped tool calls."""
calls: list[dict[str, Any]] = []
for name, body in TOOL_CALL_RE.findall(text):
arguments: dict[str, Any] = {}
for arg_name, raw_value in PARAM_RE.findall(body):
value = raw_value.strip()
try:
arguments[arg_name.strip()] = json.loads(value)
except Exception:
arguments[arg_name.strip()] = value
calls.append({
"id": f"call_{uuid.uuid4().hex[:16]}",
"type": "function",
"function": {"name": name.strip(), "arguments": json.dumps(arguments, ensure_ascii=False)},
})
return TOOL_CALL_RE.sub("", text).strip(), calls
def to_template_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
converted: list[dict[str, Any]] = []
for message in messages:
role = str(message.get("role", "user"))
content = message.get("content")
if isinstance(content, list):
content = "".join(str(part.get("text", "")) for part in content if isinstance(part, dict))
entry: dict[str, Any] = {"role": role, "content": "" if content is None else str(content)}
if role == "assistant" and message.get("tool_calls"):
calls = []
for call in message["tool_calls"]:
function = call.get("function", call)
raw_arguments = function.get("arguments", {})
if isinstance(raw_arguments, str):
try:
raw_arguments = json.loads(raw_arguments)
except Exception:
raw_arguments = {"value": raw_arguments}
calls.append({"function": {"name": function.get("name", ""), "arguments": raw_arguments}})
entry["tool_calls"] = calls
converted.append(entry)
return converted
def tool_functions(tools: Any) -> list[dict[str, Any]] | None:
if not isinstance(tools, list) or not tools:
return None
functions = [tool["function"] if isinstance(tool, dict) and isinstance(tool.get("function"), dict) else tool for tool in tools]
return [f for f in functions if isinstance(f, dict)] or None
class NM2AgentService(NaturalMemoryService):
"""The NM2 model plus a complete observation/control surface over its memory."""
def __init__(self, *args: Any, memory_mode: str = "on", session_dir: str | Path | None = None, trace_limit: int = 500, enable_thinking: bool = False, **kwargs: Any) -> None:
super().__init__(*args, **kwargs)
if memory_mode not in MEMORY_MODES:
raise ValueError(f"memory_mode must be one of {MEMORY_MODES}")
self.memory_mode = memory_mode
# Default matches the project's own chat path; see the note at apply_chat_template.
self.enable_thinking = bool(enable_thinking)
self.model_name = "natural-memory-v2"
self.session_dir = Path(session_dir) if session_dir else Path(__file__).resolve().parent / "sessions"
self.session_dir.mkdir(parents=True, exist_ok=True)
self._written_user_turns: set[str] = set()
# Origin of the user turn currently being served, so explicit writes issued by
# the agent during that turn join the same forgettable group as the turn's own
# automatic write. Tool calls arrive as separate requests, so the value has to
# outlive the completion request that started the turn.
self._turn_origin = ""
self._turn_counter = 0
self._traces: deque[dict[str, Any]] = deque(maxlen=trace_limit)
self._request_overrides: dict[str, Any] = {}
self._trace_lock = threading.RLock()
# ------------------------------------------------------------ observation
def coverage_state(self) -> dict[str, Any]:
"""The NM2.1 coverage gate: whether it is loaded, its head, and live counters."""
head = getattr(self.model, "_attribute_head_meta", None)
return {
"configured": bool(getattr(self.model.memory_config, "memory_coverage_gate", False)),
"loaded": getattr(self.model, "_attribute_coverage", None) is not None,
"bound_to_bank": bool(
self.model.memory_os_v2 is not None
and getattr(self.model.memory_os_v2, "attribute_coverage", None) is not None
),
"vocabulary_fraction": getattr(self.model.memory_config, "memory_coverage_vocabulary_fraction", None),
"head": _jsonable(head),
}
def set_coverage_gate(self, enabled: bool) -> dict[str, Any]:
"""Load or unload the attribute head and (re)bind it to the live bank."""
self.model.memory_config.memory_coverage_gate = bool(enabled)
if enabled:
self.model.setup_attribute_coverage(self.model_path)
else:
self.model._attribute_coverage = None
self.model._attribute_head_meta = None
if self.model.memory_os_v2 is not None:
self.model.memory_os_v2.attribute_coverage = getattr(self.model, "_attribute_coverage", None)
return self.coverage_state()
def runtime_snapshot(self) -> dict[str, Any]:
"""Everything NM2 recorded during the last read, in JSON-safe form."""
runtime = self.model.runtime
snapshot: dict[str, Any] = {
"router_decisions": _jsonable(list(getattr(runtime, "v2_last_decisions", []) or [])),
"no_evidence": bool(getattr(runtime, "v2_no_evidence", False)),
"prefix_tokens": int(getattr(runtime, "text_prefix_tokens", 0) or 0),
"prefix_used": bool(getattr(runtime, "text_prefix_used", False) or False),
"read_seconds": float(getattr(runtime, "text_read_seconds", 0.0) or 0.0),
"auto_memory_probability": _jsonable(getattr(runtime, "auto_memory_probability", None)),
"auto_forget_probability": _jsonable(getattr(runtime, "auto_memory_forget_probability", None)),
"context_compaction": _jsonable(getattr(runtime, "context_compaction", None)),
"coverage": self.coverage_state(),
"memory": _jsonable(self.model.memory_v2_stats()),
}
return snapshot
def effective_config(self) -> dict[str, Any]:
config = self.model.memory_config
os_v2 = self.model.memory_os_v2
values: dict[str, Any] = {}
for name, spec in CONFIG_KNOBS.items():
if spec.get("special") == "coverage_gate":
values[name] = bool(getattr(config, "memory_coverage_gate", False))
elif "config" in spec:
values[name] = getattr(config, spec["config"], None)
elif os_v2 is not None:
values[name] = getattr(os_v2, spec["runtime"], None)
return {
"knobs": values,
"fixed": {
"page_capacity": getattr(config, "memory_page_capacity", None),
"max_pages": getattr(config, "memory_max_pages", None),
"router_dim": getattr(config, "memory_router_dim", None),
"router_heads": getattr(config, "memory_router_heads", None),
"coarse_index_bits": getattr(config, "memory_coarse_index_bits", None),
"storage_mode": getattr(config, "memory_storage_mode", None),
"auto_compact_context": getattr(config, "auto_compact_context", None),
},
"memory_mode": self.memory_mode,
}
def update_config(self, payload: dict[str, Any]) -> dict[str, Any]:
applied: dict[str, Any] = {}
rejected: dict[str, str] = {}
for name, raw in (payload.get("knobs") or payload or {}).items():
spec = CONFIG_KNOBS.get(name)
if spec is None:
rejected[name] = "unknown knob"
continue
try:
if spec["type"] is bool:
value: Any = bool(raw) if not isinstance(raw, str) else raw.strip().lower() in {"1", "true", "yes", "on"}
else:
value = spec["type"](raw)
if "min" in spec and value < spec["min"] or "max" in spec and value > spec["max"]:
rejected[name] = f"out of range [{spec.get('min')}, {spec.get('max')}]"
continue
except Exception as error:
rejected[name] = f"invalid value: {error}"
continue
if spec.get("special") == "coverage_gate":
self.set_coverage_gate(bool(value))
applied[name] = bool(value)
continue
if "config" in spec:
setattr(self.model.memory_config, spec["config"], value)
else:
setattr(self.model.memory_os_v2, spec["runtime"], value)
applied[name] = value
# The KV budget lives in a dataclass built once at load time, so mirroring the
# config value is not enough on its own: without this the knob reported success
# while ``hot_limit`` stayed at the value captured at construction.
if {"kv_budget_tokens", "kv_keep_recent_tokens"} & set(applied):
budget = getattr(self.model.memory_os_v2, "kv_budget", None)
if budget is not None:
budget.max_tokens = int(self.model.memory_config.kv_budget_tokens)
budget.keep_recent_tokens = min(
int(self.model.memory_config.kv_keep_recent_tokens), budget.max_tokens
)
# Keep the bank's read threshold mirror in sync (the model sets it at load).
if self.model.memory_os_v2 is not None:
self.model.memory_os_v2.read_threshold = float(getattr(self.model.memory_os_v2, "read_threshold", 0.65))
return {"applied": applied, "rejected": rejected, "config": self.effective_config()}
def traces(self, limit: int = 50) -> list[dict[str, Any]]:
with self._trace_lock:
return list(self._traces)[-limit:]
def _push_trace(self, kind: str, payload: dict[str, Any]) -> dict[str, Any]:
with self._trace_lock:
self._turn_counter += 1
trace = {"turn": self._turn_counter, "kind": kind, "at": time.time(), **payload}
self._traces.append(trace)
return trace
# --------------------------------------------------------------- routing
def probe(self, payload: dict[str, Any]) -> dict[str, Any]:
"""Run only the bounded NM2 route for a query: no generation, full decision."""
query_text = str(payload.get("query", "")).strip()
if not query_text:
raise ValueError("query is required")
started = time.perf_counter()
if self.memory_mode == "off":
# The control arm must hold on every endpoint, not just on chat completions:
# otherwise "memory off" still retrieves through the probe.
trace = self._push_trace("probe", {
"query": query_text[:400], "memory_mode": self.memory_mode,
"decision": {"need_memory": False, "stop_reason": "memory_disabled", "record_ids": []},
"records": [],
})
return {
"query": query_text, "need_memory": False, "stop_reason": "memory_disabled",
"confidence": 0.0, "hop_count": 0, "hop_trace": [], "score_margin": None,
"evidence_score": None, "page_ids": [], "records": [], "coverage": None,
"timing": {"encode_seconds": 0.0, "route_seconds": 0.0,
"total_seconds": round(time.perf_counter() - started, 4)},
"trace_turn": trace["turn"],
}
with self.lock, torch.inference_mode():
if self.model.memory_os_v2 is None:
raise RuntimeError("NM2 hierarchical memory is not enabled in this package")
ids, mask = self._encode_plain(query_text)
key = self.model._encode_model_key(ids, mask)
encode_seconds = time.perf_counter() - started
# The NM2.1 coverage gate is the abstention mechanism: ask it directly so a
# probe can report "which attribute did the question ask about, and does the
# bank hold it" even when the router would have abstained for another reason.
coverage_verdict: dict[str, Any] | None = None
coverage_fn = getattr(self.model.memory_os_v2, "attribute_coverage", None)
if coverage_fn is not None:
try:
verdict = coverage_fn(key[0])
coverage_verdict = (
{"attribute": verdict[0], "probability": verdict[1]}
if verdict
else {"attribute": None, "probability": None, "gate": "bypassed"}
)
except Exception as error: # pragma: no cover - diagnostics only
coverage_verdict = {"error": str(error)}
route_started = time.perf_counter()
records, decision = self.model.read_hierarchical_memory(
key[0],
query_text=query_text,
query_token_ids=ids[0],
top_k_pages=payload.get("top_k_pages"),
top_k_records=payload.get("top_k_records"),
max_hops=payload.get("max_hops"),
)
route_seconds = time.perf_counter() - route_started
decision_dict = _jsonable(decision.__dict__ if hasattr(decision, "__dict__") else decision)
record_scores = decision_dict.get("record_scores") or []
records_out = []
for index, record in enumerate(records):
records_out.append({
"record_id": getattr(record, "record_id", None),
"text": getattr(record, "text", None),
"entity": getattr(record, "entity", None),
"attribute": getattr(record, "attribute", None),
"value": getattr(record, "value", None),
"status": getattr(record, "status", None),
"version": getattr(record, "version", None),
"score": record_scores[index] if index < len(record_scores) else None,
})
trace = self._push_trace("probe", {
"query": query_text[:400],
"decision": decision_dict,
"records": records_out,
"coverage": coverage_verdict,
"encode_seconds": round(encode_seconds, 4),
"route_seconds": round(route_seconds, 4),
"snapshot": self.runtime_snapshot(),
})
return {
"query": query_text,
"need_memory": bool(decision_dict.get("need_memory")),
"stop_reason": decision_dict.get("stop_reason"),
"confidence": decision_dict.get("confidence"),
"hop_count": decision_dict.get("hop_count"),
"hop_trace": decision_dict.get("hop_trace"),
"score_margin": decision_dict.get("score_margin"),
"evidence_score": decision_dict.get("evidence_score"),
"page_ids": decision_dict.get("page_ids"),
"records": records_out,
"coverage": coverage_verdict,
"timing": {"encode_seconds": round(encode_seconds, 4), "route_seconds": round(route_seconds, 4),
"total_seconds": round(time.perf_counter() - started, 4)},
"trace_turn": trace["turn"],
}
# ------------------------------------------------------------- lifecycle
def apply_overrides(self, overrides: dict[str, Any] | None) -> dict[str, Any]:
"""Temporarily steer per-request NM2 behaviour; returns the previous values."""
if not overrides:
return {}
previous: dict[str, Any] = {}
for name, value in overrides.items():
if name not in CONFIG_KNOBS:
continue
spec = CONFIG_KNOBS[name]
target = self.model.memory_config if "config" in spec else self.model.memory_os_v2
attribute = spec.get("config") or spec.get("runtime")
previous[name] = getattr(target, attribute, None)
setattr(target, attribute, value)
return previous
def restore_overrides(self, previous: dict[str, Any]) -> None:
for name, value in previous.items():
spec = CONFIG_KNOBS.get(name)
if spec is None:
continue
target = self.model.memory_config if "config" in spec else self.model.memory_os_v2
setattr(target, spec.get("config") or spec.get("runtime"), value)
# ------------------------------------------------------------ completions
def _is_tool_response(self, text: str) -> bool:
stripped = text.strip()
return stripped.startswith("") and stripped.endswith("")
def _new_write_target(self, messages: list[dict[str, Any]]) -> str | None:
for message in reversed(messages):
if str(message.get("role")) != "user":
continue
content = message.get("content")
if isinstance(content, list):
content = "".join(str(p.get("text", "")) for p in content if isinstance(p, dict))
text = str(content or "").strip()
if not text or self._is_tool_response(text):
return None
key = f"{len(text)}:{hash(text) & 0xFFFFFFFF:x}"
if key in self._written_user_turns:
return None
self._written_user_turns.add(key)
return text
return None
def complete(self, payload: dict[str, Any]) -> dict[str, Any]:
messages = payload.get("messages")
if not isinstance(messages, list) or not messages:
raise ValueError("messages is required")
max_new_tokens = min(2048, max(1, int(payload.get("max_tokens") or 512)))
tools = tool_functions(payload.get("tools"))
template_messages = to_template_messages(messages)
overrides = payload.get("nm2") if isinstance(payload.get("nm2"), dict) else {}
memory_enabled = self.memory_mode != "off" and overrides.get("memory_enabled", True) is not False
forced_query = overrides.get("query_text")
started = time.perf_counter()
previous = {}
write_target_text = ""
with self.lock, torch.inference_mode():
previous = self.apply_overrides({k: v for k, v in overrides.items() if k in CONFIG_KNOBS})
try:
write_target = self._new_write_target(messages) if (self.memory_mode == "on" and memory_enabled) else None
write_target_text = write_target or ""
# Record which turn any explicit memory write belongs to before the model
# starts generating: its tool calls are what will issue those writes.
self._turn_origin = memory_origin(write_target_text) if write_target_text else ""
memory_changed = False
reset_triggered = False
reset_id = self.model.memory_config.reset_token_id
if write_target is not None:
if reset_id is not None:
probe_ids = self.tokenizer(write_target, add_special_tokens=False)["input_ids"]
if reset_id in probe_ids:
self.model.reset_memory(batch_size=1, device=self.device)
reset_triggered = True
if not reset_triggered and self.model.memory_config.native_mode:
memory_changed = _write_turn(self.model, self.tokenizer, write_target, self.device)
if memory_changed or reset_triggered:
self._persist()
prompt = self.tokenizer.apply_chat_template(
template_messages,
tools=tools,
tokenize=False,
add_generation_prompt=True,
# With thinking enabled this checkpoint emits an unbounded
# "Thinking Process:" monologue and never reaches an answer, even at a
# 1024-token budget; the template closes an empty block
# when thinking is disabled, so the project's own chat path calls it
# with enable_thinking=False. Keep that default, allow an override.
enable_thinking=bool(payload.get("thinking", self.enable_thinking)),
)
encoded = self.tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
encoded = {key: value.to(self.device) for key, value in encoded.items()}
query_text = str(forced_query) if forced_query else write_target_text
if memory_enabled and query_text:
query_ids, query_mask = self._encode_plain(query_text)
else:
query_ids = encoded["input_ids"][:, -1:]
query_mask = torch.ones_like(query_ids)
generation_started = time.perf_counter()
output = self.model.generate(
**encoded,
max_new_tokens=max_new_tokens,
do_sample=False,
update_memory=False,
memory_query_input_ids=query_ids if memory_enabled else None,
memory_query_attention_mask=query_mask if memory_enabled else None,
memory_query_text=query_text if memory_enabled else "",
use_cache=True,
pad_token_id=self.tokenizer.pad_token_id,
)
generation_seconds = time.perf_counter() - generation_started
finally:
self.restore_overrides(previous)
raw = self.tokenizer.decode(output[0, encoded["input_ids"].shape[1]:], skip_special_tokens=True)
prompt_tokens = int(encoded["input_ids"].shape[1])
completion_tokens = int(output.shape[1] - prompt_tokens)
snapshot = self.runtime_snapshot()
content, tool_calls = parse_qwen_tool_calls(raw)
reasoning, content = split_thinking(content)
# A greedy decode that stops exactly at the cap was truncated, not finished.
finish_reason = "tool_calls" if tool_calls else ("length" if completion_tokens >= max_new_tokens else "stop")
trace = self._push_trace("chat", {
"user_message": (write_target_text or "")[:400],
"assistant_text": content[:400],
"reasoning_chars": len(reasoning),
# Hoisted so a consumer does not have to reach into `snapshot` for the
# fields it needs on every turn.
"router_decisions": snapshot["router_decisions"],
"prefix_tokens": snapshot["prefix_tokens"],
"read_seconds": snapshot["read_seconds"],
"coverage": snapshot["coverage"],
"no_evidence": snapshot["no_evidence"],
"memory_mode": self.memory_mode,
"memory_changed": bool(memory_changed),
"reset_triggered": bool(reset_triggered),
"tool_calls": [call["function"]["name"] for call in tool_calls],
"overrides": {k: v for k, v in overrides.items() if k in CONFIG_KNOBS},
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"generation_seconds": round(generation_seconds, 3),
"total_seconds": round(time.perf_counter() - started, 3),
"snapshot": snapshot,
})
return {
"id": f"chatcmpl-{uuid.uuid4().hex[:16]}",
"object": "chat.completion",
"created": int(time.time()),
"model": self.model_name,
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": content, **({"tool_calls": tool_calls} if tool_calls else {})},
"finish_reason": finish_reason,
}],
"usage": {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens},
"nm2": {
"turn": trace["turn"],
"memory_mode": self.memory_mode,
"memory_changed": bool(memory_changed),
"reset_triggered": bool(reset_triggered),
"router_decisions": snapshot["router_decisions"],
"prefix_tokens": snapshot["prefix_tokens"],
"read_seconds": snapshot["read_seconds"],
"auto_memory_probability": snapshot["auto_memory_probability"],
"auto_forget_probability": snapshot["auto_forget_probability"],
"context_compaction": snapshot["context_compaction"],
"coverage": snapshot["coverage"],
"memory": snapshot["memory"],
},
"reasoning_content": reasoning,
# Opt-in raw view: when a client asks for `debug`, the decoded model output
# and prompt size come back so an "empty answer" can be diagnosed instead of
# guessed at.
**({"debug": {"raw_text": raw[:2000], "content_chars": len(content),
"reasoning_chars": len(reasoning),
"prompt_chars": len(prompt), "prompt_tokens": prompt_tokens,
"finish_reason": finish_reason}} if payload.get("debug") else {}),
}
# ------------------------------------------------------------- governance
def write_record(self, payload: dict[str, Any]) -> dict[str, Any]:
text = str(payload.get("text", "")).strip()
if not text:
raise ValueError("text is required")
# Only forward importance/confidence when the caller actually supplies them. The
# previous defaults (0.9/0.99) were a hidden constant chosen by this server rather
# than a judgement from the model, which already carries a learned importance head;
# absent values must fall through to the memory layer's own gate instead.
graded: dict[str, Any] = {}
for knob in ("importance", "confidence"):
if payload.get(knob) is not None:
graded[knob] = float(payload[knob])
with self.lock, torch.inference_mode():
ids, mask = self._encode_plain(text)
key = self.model._encode_model_key(ids, mask)[0]
record, action = self.model.write_hierarchical_memory(
text=text,
key=key,
summary=key,
token_ids=ids[0].detach().cpu(),
token_mask=mask[0].detach().cpu().bool(),
entity=str(payload.get("entity", "")),
attribute=str(payload.get("attribute", "")),
value=str(payload.get("value", "")),
source=str(payload.get("source", "agent")),
trusted=bool(payload.get("trusted", True)),
force=bool(payload.get("force", True)),
# Tie this record to the user turn being served. A turn often carries
# more than one fact, and the agent writes one record per fact, so
# without this the records of a single turn are unrelated and an
# explicit "forget X" can only ever retire the one that matched.
origin=str(payload.get("origin") or self._turn_origin or ""),
**graded,
)
self._persist()
return {"action": action, "record": _jsonable(self.model.get_memory_record(record.record_id))}
def correct_record(self, payload: dict[str, Any]) -> dict[str, Any]:
record_id = str(payload.get("record_id", "")).strip()
if not record_id:
raise ValueError("record_id is required")
with self.lock, torch.inference_mode():
text = payload.get("text")
entity = payload.get("entity")
attribute = payload.get("attribute")
value = payload.get("value")
token_ids = token_mask = None
if not (isinstance(text, str) and text.strip()):
# Same reasoning as NaturalMemoryService.edit_memory: what the model reads is
# the stored evidence card, so a fields-only correction must rebuild that card
# (and re-encode its tokens) or the correction silently keeps answering the old
# value.
rebuilt = compose_corrected_evidence(
self.model.get_memory_record(record_id),
entity=entity if isinstance(entity, str) else None,
attribute=attribute if isinstance(attribute, str) else None,
value=value if isinstance(value, str) else None,
)
if rebuilt is not None:
text = rebuilt
if isinstance(text, str) and text.strip():
ids, mask = self._encode_plain(text)
token_ids, token_mask = ids[0].detach().cpu(), mask[0].detach().cpu().bool()
result = self.model.edit_memory_record(
record_id,
text=text if isinstance(text, str) else None,
entity=entity,
attribute=attribute,
value=value,
importance=payload.get("importance"),
confidence=payload.get("confidence"),
evidence=payload.get("evidence") if isinstance(payload.get("evidence"), list) else None,
token_ids=token_ids,
token_mask=token_mask,
)
self._persist()
return {"corrected": True, "record": _jsonable(result)}
def approve_record(self, record_id: str) -> dict[str, Any]:
with self.lock:
if self.model.memory_os_v2 is None:
raise RuntimeError("NM2 memory is not enabled")
record = self.model.memory_os_v2.approve(record_id)
self._persist()
return {"approved": True, "record": _jsonable(self.model.get_memory_record(record.record_id))}
def compact(self, payload: dict[str, Any]) -> dict[str, Any]:
text = str(payload.get("text", "")).strip()
if not text:
raise ValueError("text is required")
with self.lock, torch.inference_mode():
ids, mask = self._encode_plain(text)
hot_ids, hot_mask, report = self.model.compact_context_for_kv(
ids, mask, archive=bool(payload.get("archive", True)),
chunk_tokens=payload.get("chunk_tokens"),
)
return {
"input_tokens": int(ids.shape[1]),
"hot_tokens": int(hot_ids.shape[1]),
"report": _jsonable(report),
"memory": _jsonable(self.model.memory_v2_stats()),
}
def session_action(self, payload: dict[str, Any]) -> dict[str, Any]:
"""Per-user memory isolation: save / load / list / delete / reset."""
action = str(payload.get("action", "")).strip()
user = re.sub(r"[^A-Za-z0-9_.-]", "_", str(payload.get("user", "default")).strip() or "default")
path = self.session_dir / f"{user}.pt"
if action == "list":
return {"sessions": [f.stem for f in sorted(self.session_dir.glob("*.pt"))]}
if action == "delete":
path.unlink(missing_ok=True)
return {"deleted": user}
if action == "save":
with self.lock:
payload_out = self.model.memory_os_v2.export_payload() if self.model.memory_os_v2 else {}
torch.save({"user": user, "saved_at": time.time(), "nm2": payload_out}, path)
return {"saved": user, "path": str(path), "bytes": path.stat().st_size,
"stats": _jsonable(self.model.memory_v2_stats())}
if action == "load":
if not path.exists():
raise FileNotFoundError(f"no saved session for {user}")
with self.lock:
saved = torch.load(path, map_location="cpu", weights_only=False)
payload = saved.get("nm2") or {}
if not payload:
raise ValueError(f"saved session {user} has no NM2 payload")
current = self.model.memory_os_v2
if current is None:
raise RuntimeError("NM2 memory is not enabled")
bank = current.bank
restored = MemoryOSV2.from_payload(
payload,
router=current.router,
runtime_device=self.device,
gpu_cache_records=int(getattr(bank, "gpu_cache_records", 256)),
gpu_cache_tokens=int(getattr(bank, "gpu_cache_tokens", 131072)),
gpu_cache_reserve_mb=int(getattr(bank, "gpu_cache_reserve_mb", 2048)),
gpu_cache_adaptive=bool(getattr(bank, "gpu_cache_adaptive", True)),
min_read_margin=float(payload.get("min_read_margin", 0.0)),
require_evidence=bool(payload.get("require_evidence", False)),
)
# Thresholds travel with the session, not with the process.
restored.read_threshold = float(payload.get("read_threshold", 0.65))
restored.write_threshold = float(payload.get("write_threshold", 0.50))
for name, spec in CONFIG_KNOBS.items():
if "config" in spec:
continue
saved_value = payload.get(name)
if saved_value is not None:
setattr(restored, spec["runtime"], saved_value)
self.model.memory_os_v2 = restored
self._written_user_turns.clear()
return {
"loaded": user,
"path": str(path),
"saved_at": saved.get("saved_at"),
"read_threshold": restored.read_threshold,
"write_threshold": restored.write_threshold,
"memory": _jsonable(self.model.memory_v2_stats()),
}
if action == "reset":
with self.lock:
self.model.reset_memory(batch_size=1, device=self.device)
self._written_user_turns.clear()
self._persist()
return {"reset": True, "memory": _jsonable(self.model.memory_v2_stats())}
raise ValueError(f"unknown session action {action!r}")
class _Handler(BaseHTTPRequestHandler):
server_version = "NM2AgentControlPlane/1.0"
service: NM2AgentService
def log_message(self, fmt: str, *args: Any) -> None:
sys.stderr.write("[nm2] " + (fmt % args) + "\n")
def _send(self, status: int, payload: Any) -> None:
body = json.dumps(payload, ensure_ascii=False, default=str).encode("utf-8")
# A client can vanish between the request and the response (the harness aborts a
# turn by closing the socket). That is not a server fault and there is nobody left
# to answer, so the write is best-effort rather than an exception that would bubble
# into the generic handler and try to send a 500 down the same dead socket.
try:
self.send_response(status)
self.send_header("Content-Type", "application/json; charset=utf-8")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError, OSError):
self.close_connection = True
def _send_stream(self, completion: dict[str, Any]) -> None:
"""Emit one completion as an OpenAI SSE stream.
Streaming clients (pi-ai among them) parse the ``chat.completion.chunk`` event
shape; a single JSON body leaves them with an empty message and no tool calls.
Generation here is not incremental, so the whole content is sent as one delta —
the wire format is what the client needs, not sub-token granularity.
"""
# No Content-Length is sent for a stream, so the client needs the socket to close
# to see EOF. Incremental SSE parsers do not care, but simple HTTP clients hang
# on a kept-alive connection.
self.close_connection = True
def chunk(delta: dict[str, Any], finish: str | None = None) -> dict[str, Any]:
return {
"id": completion["id"],
"object": "chat.completion.chunk",
"created": completion["created"],
"model": completion["model"],
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
}
message = completion["choices"][0]["message"]
finish_reason = completion["choices"][0]["finish_reason"]
events: list[dict[str, Any]] = [chunk({"role": "assistant", "content": ""})]
if message.get("content"):
events.append(chunk({"content": message["content"]}))
for index, call in enumerate(message.get("tool_calls") or []):
events.append(chunk({"tool_calls": [{
"index": index, "id": call["id"], "type": "function",
"function": {"name": call["function"]["name"], "arguments": call["function"]["arguments"]},
}]}))
events.append(chunk({}, finish_reason))
# The headers belong inside the guard too: a client that disconnects while the
# stream is being set up reset the connection during ``end_headers`` and the
# resulting ConnectionResetError escaped into the generic 500 path, logging a
# misleading ``POST /v1/chat/completions 500`` and a socketserver traceback for
# what was only a client hang-up.
try:
self.send_response(200)
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
self.send_header("Cache-Control", "no-cache")
self.send_header("Connection", "close")
self.end_headers()
for event in events:
self.wfile.write(f"data: {json.dumps(event, ensure_ascii=False, default=str)}\n\n".encode("utf-8"))
self.wfile.write(b"data: [DONE]\n\n")
self.wfile.flush()
except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError, OSError):
pass # client aborted mid-stream
def _read_json(self) -> dict[str, Any]:
length = int(self.headers.get("Content-Length") or 0)
if length <= 0:
return {}
return json.loads(self.rfile.read(length).decode("utf-8"))
def _route(self, method: str) -> None:
parsed = urlparse(self.path)
params = parse_qs(parsed.query)
service = self.service
path = parsed.path
try:
if method == "GET":
if path == "/health":
return self._send(200, {**service.health(), "memory_mode": service.memory_mode})
if path == "/v1/models":
return self._send(200, {"object": "list", "data": [
{"id": service.model_name, "object": "model", "owned_by": "local"}]})
if path == "/v1/nm2/state":
return self._send(200, {"config": service.effective_config(), "snapshot": service.runtime_snapshot()})
if path == "/v1/nm2/config":
return self._send(200, service.effective_config())
if path == "/v1/nm2/trace":
return self._send(200, {"traces": service.traces(int(params.get("limit", ["50"])[0]))})
if path == "/v1/nm2/records":
return self._send(200, service.list_memory(params))
if path == "/v1/nm2/audit":
return self._send(200, _jsonable(service.audit()))
if path == "/v1/nm2/export":
return self._send(200, _jsonable(service.export_memory(params)))
if path.startswith("/v1/nm2/records/"):
return self._send(200, _jsonable(service.get_memory(path.rsplit("/", 1)[-1])))
if path == "/v1/memory":
return self._send(200, service.list_memory(params))
if path == "/v1/memory/audit":
return self._send(200, _jsonable(service.audit()))
if path == "/v1/diagnostics":
return self._send(200, {"turns": service.traces(int(params.get("limit", ["50"])[0]))})
return self._send(404, {"error": {"message": f"unknown path {path}", "type": "not_found"}})
payload = self._read_json()
if path == "/v1/chat/completions":
completion = service.complete(payload)
if payload.get("stream"):
return self._send_stream(completion)
return self._send(200, completion)
# The control arm is only a control if the memory layer is genuinely
# unavailable. `memory_mode=off` used to disable only the automatic path
# (injected prefix + auto-write) while an explicit nm2_write tool call still
# stored records and nm2_read still retrieved them -- so an agent told to use
# its memory tools could pass a "no memory" arm outright.
if service.memory_mode != "on" and path in _MEMORY_MUTATIONS:
return self._send(409, {"error": {
"message": f"memory_mode={service.memory_mode}: this arm cannot write long-term memory",
"type": "memory_disabled" if service.memory_mode == "off" else "memory_read_only",
}})
if path == "/v1/nm2/probe":
return self._send(200, service.probe(payload))
if path == "/v1/nm2/config":
return self._send(200, service.update_config(payload))
if path == "/v1/nm2/mode":
mode = str(payload.get("mode", "")).strip()
if mode not in MEMORY_MODES:
raise ValueError(f"mode must be one of {MEMORY_MODES}")
service.memory_mode = mode
return self._send(200, {"memory_mode": mode})
if path == "/v1/nm2/write":
return self._send(200, service.write_record(payload))
if path == "/v1/nm2/correct":
return self._send(200, service.correct_record(payload))
if path == "/v1/nm2/retract":
record_id = str(payload.get("record_id", "")).strip()
if not record_id:
raise ValueError("record_id is required")
return self._send(200, _jsonable(service.retract_memory(record_id)))
if path == "/v1/nm2/approve":
record_id = str(payload.get("record_id", "")).strip()
if not record_id:
raise ValueError("record_id is required")
return self._send(200, service.approve_record(record_id))
if path == "/v1/nm2/compact":
return self._send(200, service.compact(payload))
if path == "/v1/nm2/session":
return self._send(200, service.session_action(payload))
if path == "/v1/nm2/reset":
service._written_user_turns.clear()
return self._send(200, _jsonable(service.reset_memory()))
if path == "/v1/memory/mode":
mode = str(payload.get("mode", "")).strip()
if mode not in MEMORY_MODES:
raise ValueError(f"mode must be one of {MEMORY_MODES}")
service.memory_mode = mode
return self._send(200, {"memory_mode": mode})
if path == "/v1/memory/reset":
service._written_user_turns.clear()
return self._send(200, _jsonable(service.reset_memory()))
if path == "/v1/memory":
return self._send(200, service.write_record(payload))
return self._send(404, {"error": {"message": f"unknown path {path}", "type": "not_found"}})
except KeyError as error:
# Unknown record ids arrive from a model that guessed one; answer with a
# readable 404 instead of a raw KeyError wrapped in a 500.
self._send(404, {"error": {"message": f"unknown record id: {error}", "type": "unknown_record"}})
except ValueError as error:
message = str(error)
if message.startswith("only active records can be edited"):
# A domain conflict, not a server fault. NM2's own auto-forget retracts a
# fact in the very turn where the user revokes it, so a model that then
# asks to "correct" the record it memorised last turn is addressing
# something that is no longer active. A 500 taught the model nothing and
# it looped; 409 plus the exact recovery keeps the turn recoverable.
return self._send(409, {
"error": {
"message": message,
"type": "record_not_active",
"hint": (
"this record is no longer active, so it cannot be edited in place; "
"store the new value with a plain write "
"(action=write, entity+attribute+value) -- NM2 versions by "
"entity::attribute automatically"
),
},
})
self._send(400, {"error": {"message": message, "type": "bad_request"}})
except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError) as error:
# The client hung up (a harness aborting a turn does exactly this). There is
# no one to answer and nothing went wrong on this side, so this must not be
# logged as a 500 the way it was before.
self.close_connection = True
sys.stderr.write(f"[nm2] client disconnected: {type(error).__name__}\n")
except Exception as error:
self._send(500, {"error": {"message": str(error), "type": type(error).__name__}})
def do_GET(self) -> None: # noqa: N802
self._route("GET")
def do_POST(self) -> None: # noqa: N802
self._route("POST")
class _Server(ThreadingHTTPServer):
daemon_threads = True
allow_reuse_address = True
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model-path", default="qwen3_5_4b_natural_memory_v2")
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, default=8766)
parser.add_argument("--no-4bit", action="store_true")
parser.add_argument("--no-auto-persist", action="store_true")
parser.add_argument("--memory-mode", choices=MEMORY_MODES, default="on")
parser.add_argument("--session-dir", default="")
parser.add_argument("--trace-limit", type=int, default=500)
parser.add_argument("--model-name", default="natural-memory-v2")
args = parser.parse_args()
model_path = Path(args.model_path)
if not model_path.is_absolute() and not model_path.exists():
model_path = Path(__file__).resolve().parents[2] / "V2_dpskw" / args.model_path
print(json.dumps({"event": "loading", "model_path": str(model_path)}, ensure_ascii=False), flush=True)
service = NM2AgentService(
model_path,
no_4bit=args.no_4bit,
auto_persist=not args.no_auto_persist,
memory_mode=args.memory_mode,
session_dir=args.session_dir or None,
trace_limit=args.trace_limit,
)
service.model_name = args.model_name
service.model.eval()
handler = type("Handler", (_Handler,), {"service": service})
server = _Server((args.host, args.port), handler)
print(json.dumps({
"event": "ready",
"url": f"http://{args.host}:{args.port}",
"openai_base_url": f"http://{args.host}:{args.port}/v1",
"memory_mode": service.memory_mode,
"nm2_endpoints": [
"GET /v1/nm2/state", "GET /v1/nm2/config", "POST /v1/nm2/config",
"GET /v1/nm2/trace", "POST /v1/nm2/probe", "GET /v1/nm2/records",
"POST /v1/nm2/write", "POST /v1/nm2/correct", "POST /v1/nm2/retract",
"POST /v1/nm2/approve", "GET /v1/nm2/audit", "GET /v1/nm2/export",
"POST /v1/nm2/compact", "POST /v1/nm2/session", "POST /v1/nm2/reset",
],
"memory": _jsonable(service.model.memory_v2_stats()),
}, ensure_ascii=False, default=str), flush=True)
try:
server.serve_forever()
except KeyboardInterrupt:
pass
finally:
server.server_close()
service.close()
return 0
if __name__ == "__main__":
raise SystemExit(main())