1067 lines
55 KiB
Python
1067 lines
55 KiB
Python
"""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"<tool_call>\s*<function=([^>\n]+)>(.*?)</function>\s*</tool_call>", re.S)
|
|
PARAM_RE = re.compile(r"<parameter=([^>\n]+)>\n?(.*?)\n?</parameter>", 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 ``<think>…</think>`` 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 ``<think></think>`` 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("<think>"):
|
|
head, separator, tail = text.partition("</think>")
|
|
if not separator:
|
|
return head[len("<think>"):].strip(), ""
|
|
reasoning, text = head[len("<think>"):].strip(), tail
|
|
text = re.sub(r"<think>\s*</think>", "", 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("<tool_response>") and stripped.endswith("</tool_response>")
|
|
|
|
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 <think></think> 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())
|