@sovorn/pi-session-memory 0.1.0

package/worker/laya_runtime.py

@@ -0,0 +1,412 @@
+"""Pinned Laya decision adapter used by the JSONL worker."""
+
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+import hashlib
+import importlib.metadata
+import inspect
+import json
+import math
+import os
+from pathlib import Path
+import resource
+import subprocess
+import sys
+import time
+from urllib.parse import unquote, urlparse
+import warnings
+
+from worker.protocol import Gate, GateDecision, GateRequest, ProjectionDecision
+
+
+LAYA_SOURCE_COMMIT = "010bacef009c855ccba814b51f7c8e1d38ab5e3f"
+CHECKPOINT_REVISION = "f9ab0b228f0fc0f14d873dbc99038f135c2da1b2"
+CHECKPOINT_SHA256 = "4fa56de72383a9d3efa9cfa78955733c81b9fc8067a587ca4beb82c78107a24e"
+CHECKPOINT_REPO = "models--convaiinnovations--laya-typed-decisions"
+Digest = Callable[[Path], str]
+GitRunner = Callable[[Path, list[str]], subprocess.CompletedProcess[str]]
+Question = dict[str, str | dict[str, str]]
+MAX_PROJECTION_CANDIDATES = 8
+MAX_PROJECTION_NEED_CHARS = 800
+MAX_CANDIDATE_TEXT_CHARS = 500
+MAX_MEMORY_TEXT_CHARS = 1_000
+MAX_SUPERSESSION_SOURCES = 12
+MAX_ENTRY_ID_CHARS = 256
+MAX_SUPERSESSION_SOURCE_TEXT_CHARS = 5_000
+NO_CANDIDATE = "none"
+QUESTIONS: dict[Gate, Question] = {
+    "observation": {
+        "type": "noul",
+        "instructions": (
+            "Does this single session turn contain a durable user decision, constraint, unresolved commitment, "
+            "or project fact likely to matter later? Answer yes only for concrete reusable evidence; do not "
+            "retain small talk, transient status, or unsupported inference."
+        ),
+        "criteria": {
+            "true": "The turn contains a concrete durable fact, decision, constraint, or commitment worth retaining.",
+            "false": "The turn is transient, routine, unsupported, or has no likely future value.",
+        },
+    },
+    "reflection": {
+        "type": "noul",
+        "instructions": (
+            "Do these linked observations together establish a stable recurring decision or insight across turns "
+            "that is not already captured by any one observation? Answer yes only when the evidence supports a "
+            "cross-turn synthesis; unrelated facts or a single isolated observation are not mature."
+        ),
+        "criteria": {
+            "true": "Multiple observations support a stable, useful cross-turn synthesis.",
+            "false": "Evidence is unrelated, isolated, transient, or does not support a distinct synthesis.",
+        },
+    },
+    "resident": {
+        "type": "noul",
+        "instructions": (
+            "Is the already resident, linked session memory sufficient to address this current need without "
+            "retrieving another memory entry? Answer yes only when its evidence directly covers the request."
+        ),
+        "criteria": {
+            "true": "The resident source-linked memory directly covers the current need.",
+            "false": "The resident memory is insufficient or unrelated to the current need.",
+        },
+    },
+    "supersession": {
+        "type": "noul",
+        "instructions": (
+            "Does the exact linked newer session evidence explicitly establish that the old memory has been "
+            "reversed, corrected, or replaced? Do not treat recency, unrelated facts, an additive detail, or "
+            "uncertainty as supersession; the newer evidence must actually supersede the old claim."
+        ),
+        "criteria": {
+            "true": "The exact linked newer evidence explicitly contradicts or replaces the old memory; supersession is established.",
+            "false": "Supersession is not established: evidence is later only, unrelated, additive, ambiguous, or insufficient.",
+        },
+    },
+}
+
+
+@dataclass(frozen=True)
+class RuntimeContext:
+    python_version: str
+    laya_version: str
+    laya_source_commit: str
+    checkpoint_revision: str
+    checkpoint_sha256: str
+    device: str
+    mps_available: bool
+    load_ms: float
+    max_rss_bytes: int
+    warnings: tuple[str, ...]
+
+
+def _expanded(value: str, home: Path) -> Path:
+    text = os.path.expandvars(value)
+    if text == "~" or text.startswith("~/"):
+        text = f"{home}{text[1:]}"
+    return Path(text)
+
+
+def _env_text(env: Mapping[str, str], name: str) -> str | None:
+    value = env.get(name)
+    if not isinstance(value, str) or value.strip() == "":
+        return None
+    return value
+
+
+def checkpoint_directory(env: Mapping[str, str] | None = None, home: str | Path | None = None) -> Path:
+    if env is None:
+        env = os.environ
+    root = Path(home).expanduser() if home is not None else Path.home()
+    explicit = _env_text(env, "PI_SESSION_MEMORY_LAYA_CHECKPOINT")
+    if explicit is not None:
+        return _expanded(explicit, root)
+    hub = _env_text(env, "HF_HUB_CACHE")
+    if hub is None:
+        hf_home = _env_text(env, "HF_HOME")
+        if hf_home is not None:
+            hub_path = _expanded(hf_home, root) / "hub"
+        else:
+            xdg = _env_text(env, "XDG_CACHE_HOME")
+            cache_home = _expanded(xdg, root) if xdg is not None else root / ".cache"
+            hub_path = cache_home / "huggingface" / "hub"
+    else:
+        hub_path = _expanded(hub, root)
+    return hub_path / CHECKPOINT_REPO / "snapshots" / CHECKPOINT_REVISION
+
+
+def _sha256(path: Path) -> str:
+    with path.open("rb") as source:
+        return hashlib.file_digest(source, "sha256").hexdigest()
+
+
+def _verified_checkpoint(
+    env: Mapping[str, str] | None = None,
+    home: str | Path | None = None,
+    digest: Digest | None = None,
+) -> Path:
+    checkpoint = checkpoint_directory(env, home).resolve()
+    if checkpoint.name != CHECKPOINT_REVISION:
+        raise RuntimeError("checkpoint revision mismatch")
+    weights = checkpoint / "model.safetensors"
+    if not weights.is_file():
+        raise RuntimeError(f"checkpoint weights missing: {weights}")
+    if (digest or _sha256)(weights) != CHECKPOINT_SHA256:
+        raise RuntimeError("checkpoint weights mismatch")
+    return checkpoint
+
+
+def _pinned_commit(commit_id: object) -> bool:
+    if not isinstance(commit_id, str) or len(commit_id) != 40:
+        return False
+    if any(character not in "0123456789abcdefABCDEF" for character in commit_id):
+        return False
+    return commit_id.lower() == LAYA_SOURCE_COMMIT
+
+
+def _accept_vcs(vcs_info: object) -> None:
+    if not isinstance(vcs_info, dict) or vcs_info.get("vcs") != "git":
+        raise RuntimeError("Laya source is not the pinned local checkout")
+    commit_id = vcs_info.get("commit_id") if "commit_id" in vcs_info else None
+    if commit_id is None:
+        raise RuntimeError("Laya source is not the pinned local checkout")
+    if not _pinned_commit(commit_id):
+        raise RuntimeError("Laya source commit mismatch")
+
+
+def _verified_source(
+    distribution: importlib.metadata.Distribution,
+    loaded_agent_path: Path,
+    git_runner: GitRunner | None = None,
+) -> None:
+    metadata_text = distribution.read_text("direct_url.json")
+    if metadata_text is None:
+        raise RuntimeError("Laya source origin is unavailable")
+    try:
+        payload = json.loads(metadata_text)
+    except json.JSONDecodeError:
+        raise RuntimeError("Laya source origin is unavailable") from None
+    if not isinstance(payload, dict):
+        raise RuntimeError("Laya source is not the pinned local checkout")
+    if "vcs_info" in payload and payload.get("vcs_info") is not None:
+        _accept_vcs(payload.get("vcs_info"))
+        return
+    source_url = payload.get("url")
+    parsed_url = urlparse(source_url) if isinstance(source_url, str) else None
+    if parsed_url is None or parsed_url.scheme != "file":
+        raise RuntimeError("Laya source is not the pinned local checkout")
+    source_root = Path(unquote(parsed_url.path)).resolve()
+
+    if git_runner is None:
+        commit = subprocess.run(
+            ["git", "-C", str(source_root), "rev-parse", "HEAD"],
+            check=True,
+            capture_output=True,
+            text=True,
+        ).stdout.strip()
+        clean = subprocess.run(
+            ["git", "-C", str(source_root), "diff", "--quiet", "HEAD", "--", "laya"],
+            check=False,
+            capture_output=True,
+        )
+    else:
+        commit = git_runner(source_root, ["rev-parse", "HEAD"]).stdout.strip()
+        clean = git_runner(source_root, ["diff", "--quiet", "HEAD", "--", "laya"])
+    if commit != LAYA_SOURCE_COMMIT:
+        raise RuntimeError("Laya source commit mismatch")
+    if clean.returncode != 0:
+        raise RuntimeError("Laya source checkout is modified")
+
+    installed_root = loaded_agent_path.parent
+    for filename in ("__init__.py", "agent.py", "common.py"):
+        if (installed_root / filename).read_bytes() != (source_root / "laya" / filename).read_bytes():
+            raise RuntimeError("installed Laya code differs from the pinned source")
+
+
+def _probability(value: object) -> float:
+    if isinstance(value, bool) or not isinstance(value, (float, int)):
+        raise ValueError("Laya returned a non-numeric probability")
+    result = float(value)
+    if not math.isfinite(result) or not 0.0 <= result <= 1.0:
+        raise ValueError("Laya returned an invalid probability")
+    return result
+
+
+def _projection_state(state: str) -> tuple[str, list[tuple[str, str, str]]]:
+    try:
+        value: object = json.loads(state)
+    except json.JSONDecodeError:
+        raise ValueError("projection state is invalid JSON") from None
+    if not isinstance(value, dict) or set(value) != {"need", "candidates"}:
+        raise ValueError("projection state fields are invalid")
+    need = value["need"]
+    candidates = value["candidates"]
+    if not isinstance(need, str) or not need or len(need) > MAX_PROJECTION_NEED_CHARS:
+        raise ValueError("projection need is invalid")
+    if not isinstance(candidates, list) or not candidates or len(candidates) > MAX_PROJECTION_CANDIDATES:
+        raise ValueError("projection candidates are empty or exceed the bound")
+
+    parsed: list[tuple[str, str, str]] = []
+    seen: set[str] = set()
+    for candidate in candidates:
+        if not isinstance(candidate, dict) or set(candidate) != {"entryId", "kind", "text"}:
+            raise ValueError("projection candidate fields are invalid")
+        entry_id = candidate["entryId"]
+        kind = candidate["kind"]
+        text = candidate["text"]
+        if (
+            not isinstance(entry_id, str)
+            or not entry_id
+            or len(entry_id) > 256
+            or entry_id in seen
+            or kind not in ("observation", "reflection")
+            or not isinstance(text, str)
+            or not text
+            or len(text) > MAX_CANDIDATE_TEXT_CHARS
+        ):
+            raise ValueError("projection candidate is invalid")
+        seen.add(entry_id)
+        parsed.append((entry_id, kind, text))
+    return need, parsed
+
+
+def _supersession_state(state: str) -> None:
+    if len(state.encode("utf-8")) > 8_192:
+        raise ValueError("supersession state exceeds its byte bound")
+    try:
+        value: object = json.loads(state)
+    except json.JSONDecodeError:
+        raise ValueError("supersession state is invalid JSON") from None
+    if not isinstance(value, dict) or set(value) != {"oldMemory", "newEvidence"}:
+        raise ValueError("supersession state fields are invalid")
+    old_memory = value["oldMemory"]
+    evidence = value["newEvidence"]
+    if not isinstance(old_memory, dict) or set(old_memory) != {"entryId", "kind", "text"}:
+        raise ValueError("supersession old memory fields are invalid")
+    if (
+        not _valid_entry_id(old_memory["entryId"])
+        or old_memory["kind"] not in ("observation", "reflection")
+        or not _valid_memory_text(old_memory["text"])
+    ):
+        raise ValueError("supersession old memory is invalid")
+    if not isinstance(evidence, dict) or set(evidence) != {"entryId", "text", "sourceEntryIds", "sources"}:
+        raise ValueError("supersession newer evidence fields are invalid")
+    source_ids = evidence["sourceEntryIds"]
+    sources = evidence["sources"]
+    if (
+        not _valid_entry_id(evidence["entryId"])
+        or evidence["entryId"] == old_memory["entryId"]
+        or not _valid_memory_text(evidence["text"])
+        or not isinstance(source_ids, list)
+        or not source_ids
+        or len(source_ids) > MAX_SUPERSESSION_SOURCES
+        or not all(_valid_entry_id(entry_id) for entry_id in source_ids)
+        or len(set(source_ids)) != len(source_ids)
+        or evidence["entryId"] in source_ids
+        or old_memory["entryId"] in source_ids
+        or not isinstance(sources, list)
+        or len(sources) != len(source_ids)
+    ):
+        raise ValueError("supersession newer evidence is invalid")
+    for expected_id, source in zip(source_ids, sources, strict=True):
+        if (
+            not isinstance(source, dict)
+            or set(source) != {"entryId", "role", "text"}
+            or source["entryId"] != expected_id
+            or source["role"] not in ("user", "assistant", "toolResult")
+            or not isinstance(source["text"], str)
+            or len(source["text"]) > MAX_SUPERSESSION_SOURCE_TEXT_CHARS
+        ):
+            raise ValueError("supersession source link is invalid")
+
+
+def _valid_entry_id(value: object) -> bool:
+    return isinstance(value, str) and 0 < len(value) <= MAX_ENTRY_ID_CHARS
+
+
+def _valid_memory_text(value: object) -> bool:
+    return isinstance(value, str) and 0 < len(value) <= MAX_MEMORY_TEXT_CHARS
+
+
+class LayaEvaluator:
+    def __init__(self) -> None:
+        from laya.agent import Agent
+        import torch
+
+        distribution = importlib.metadata.distribution("laya")
+        if distribution.version != "0.3.7":
+            raise RuntimeError("Laya distribution version mismatch")
+        _verified_source(distribution, Path(inspect.getfile(Agent)).resolve())
+        checkpoint = _verified_checkpoint()
+
+        start = time.perf_counter()
+        with warnings.catch_warnings(record=True) as captured_warnings:
+            warnings.simplefilter("always")
+            self.agent = Agent(str(checkpoint))
+        load_ms = (time.perf_counter() - start) * 1000.0
+        max_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+        max_rss_bytes = int(max_rss if sys.platform == "darwin" else max_rss * 1024)
+        mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+        self.context = RuntimeContext(
+            python_version=sys.version.split()[0],
+            laya_version=distribution.version,
+            laya_source_commit=LAYA_SOURCE_COMMIT,
+            checkpoint_revision=CHECKPOINT_REVISION,
+            checkpoint_sha256=CHECKPOINT_SHA256,
+            device=str(self.agent.device),
+            mps_available=bool(mps_available),
+            load_ms=round(load_ms, 2),
+            max_rss_bytes=max_rss_bytes,
+            warnings=tuple(str(warning.message) for warning in captured_warnings),
+        )
+
+    def evaluate(self, request: GateRequest) -> GateDecision | ProjectionDecision:
+        if request.gate == "projection":
+            return self._select_projection(request)
+        if request.gate == "supersession":
+            _supersession_state(request.state)
+        result = self.agent.system_one(
+            {"gate": request.gate, "evidence": request.state},
+            {"warranted": QUESTIONS[request.gate]},
+        )
+        answers = result.get("answers")
+        if not isinstance(answers, dict):
+            raise ValueError("Laya returned no typed answers")
+        answer = answers.get("warranted")
+        if not isinstance(answer, dict) or answer.get("type") != "noul":
+            raise ValueError("Laya returned an unexpected decision type")
+        p_true = _probability(answer.get("noul"))
+        confidence = _probability(answer.get("confidence"))
+        return GateDecision(accepted=p_true >= 0.5, p_true=round(p_true, 4), confidence=round(confidence, 4))
+
+    def _select_projection(self, request: GateRequest) -> ProjectionDecision:
+        need, candidates = _projection_state(request.state)
+        labels = {f"candidate_{index}": f"{kind}: {text}" for index, (_, kind, text) in enumerate(candidates)}
+        labels[NO_CANDIDATE] = "None of these source-linked memories is relevant to the current need."
+        result = self.agent.system_one(
+            {"gate": request.gate, "current_need": need},
+            {
+                "selected": {
+                    "type": "choice",
+                    "instructions": "Which one bounded memory candidate is most relevant to the current need? Select none if no candidate is relevant.",
+                    "criteria": labels,
+                },
+            },
+        )
+        answers = result.get("answers")
+        if not isinstance(answers, dict):
+            raise ValueError("Laya returned no typed selection")
+        answer = answers.get("selected")
+        if not isinstance(answer, dict) or answer.get("type") != "choice":
+            raise ValueError("Laya returned an unexpected selection type")
+        selected = answer.get("choice")
+        if selected == NO_CANDIDATE:
+            return ProjectionDecision(selected_entry_id=None)
+        if not isinstance(selected, str) or not selected.startswith("candidate_"):
+            raise ValueError("Laya returned an unknown selection")
+        try:
+            index = int(selected.removeprefix("candidate_"))
+        except ValueError:
+            raise ValueError("Laya returned an unknown selection") from None
+        if index < 0 or index >= len(candidates) or selected != f"candidate_{index}":
+            raise ValueError("Laya returned an unknown selection")
+        return ProjectionDecision(selected_entry_id=candidates[index][0])

package/worker/protocol.py

@@ -0,0 +1,141 @@
+"""Bounded JSONL protocol for the pinned, session-local Laya gate worker."""
+
+from contextlib import redirect_stdout
+from dataclasses import asdict, dataclass
+import json
+import re
+from typing import BinaryIO, Callable, Literal, TextIO, TypeAlias
+
+
+Gate: TypeAlias = Literal["observation", "reflection", "resident", "projection", "supersession"]
+ErrorCode: TypeAlias = Literal["invalid_json", "invalid_request", "request_too_large", "laya_error"]
+PROTOCOL_VERSION = 1
+MAX_REQUEST_BYTES = 16_384
+MAX_STATE_BYTES = 8_192
+MAX_RESPONSE_BYTES = 2_048
+_REQUEST_ID = re.compile(r"[A-Za-z0-9._:-]{1,64}\Z")
+
+
+@dataclass(frozen=True)
+class GateRequest:
+    request_id: str
+    gate: Gate
+    state: str
+
+
+@dataclass(frozen=True)
+class GateDecision:
+    accepted: bool
+    p_true: float
+    confidence: float
+
+
+@dataclass(frozen=True)
+class ProjectionDecision:
+    selected_entry_id: str | None
+
+
+@dataclass(frozen=True)
+class ProtocolError(Exception):
+    code: ErrorCode
+
+
+def parse_request(line: bytes) -> GateRequest:
+    try:
+        payload: object = json.loads(line.decode("utf-8"))
+    except (UnicodeDecodeError, json.JSONDecodeError):
+        raise ProtocolError("invalid_json") from None
+
+    if not isinstance(payload, dict):
+        raise ProtocolError("invalid_request")
+
+    request_id = payload.get("request_id")
+    if not isinstance(request_id, str) or not _REQUEST_ID.fullmatch(request_id):
+        raise ProtocolError("invalid_request")
+    if set(payload) != {"protocol_version", "request_id", "gate", "state"}:
+        raise ProtocolError("invalid_request")
+    if type(payload["protocol_version"]) is not int or payload["protocol_version"] != PROTOCOL_VERSION:
+        raise ProtocolError("invalid_request")
+
+    gate = payload["gate"]
+    if gate == "observation":
+        typed_gate: Gate = "observation"
+    elif gate == "reflection":
+        typed_gate = "reflection"
+    elif gate == "resident":
+        typed_gate = "resident"
+    elif gate == "projection":
+        typed_gate = "projection"
+    elif gate == "supersession":
+        typed_gate = "supersession"
+    else:
+        raise ProtocolError("invalid_request")
+    state = payload["state"]
+    if not isinstance(state, str) or not state or len(state.encode("utf-8")) > MAX_STATE_BYTES:
+        raise ProtocolError("invalid_request")
+    return GateRequest(request_id=request_id, gate=typed_gate, state=state)
+
+
+def _error_response(code: ErrorCode, request_id: str | None = None) -> dict[str, object]:
+    messages = {
+        "invalid_json": "Request is not valid UTF-8 JSON.",
+        "invalid_request": "Request fields do not match protocol version 1.",
+        "request_too_large": "Request exceeds the configured byte limit.",
+        "laya_error": "The Laya decision could not be trusted.",
+    }
+    return {
+        "protocol_version": PROTOCOL_VERSION,
+        "request_id": request_id,
+        "status": "error",
+        "error": {"code": code, "message": messages[code]},
+    }
+
+
+def _write_response(output: BinaryIO, payload: dict[str, object]) -> None:
+    encoded = json.dumps(payload, ensure_ascii=False, separators=(",", ":"), allow_nan=False).encode("utf-8") + b"\n"
+    if len(encoded) > MAX_RESPONSE_BYTES:
+        encoded = json.dumps(_error_response("laya_error"), separators=(",", ":")).encode("utf-8") + b"\n"
+    output.write(encoded)
+    output.flush()
+
+
+def _discard_line_remainder(input_stream: BinaryIO) -> None:
+    while True:
+        remainder = input_stream.readline(MAX_REQUEST_BYTES + 1)
+        if not remainder or remainder.endswith(b"\n"):
+            return
+
+
+def serve(
+    input_stream: BinaryIO,
+    output_stream: BinaryIO,
+    diagnostics: TextIO,
+    evaluate: Callable[[GateRequest], GateDecision | ProjectionDecision],
+) -> None:
+    while line := input_stream.readline(MAX_REQUEST_BYTES + 1):
+        if len(line) > MAX_REQUEST_BYTES:
+            if not line.endswith(b"\n"):
+                _discard_line_remainder(input_stream)
+            _write_response(output_stream, _error_response("request_too_large"))
+            continue
+
+        try:
+            request = parse_request(line.rstrip(b"\r\n"))
+        except ProtocolError as error:
+            _write_response(output_stream, _error_response(error.code))
+            continue
+
+        try:
+            with redirect_stdout(diagnostics):
+                decision = evaluate(request)
+            response = {
+                "protocol_version": PROTOCOL_VERSION,
+                "request_id": request.request_id,
+                "gate": request.gate,
+                "status": "ok",
+                "decision": asdict(decision),
+            }
+        except Exception as error:
+            diagnostics.write(f"Laya gate failed ({type(error).__name__}).\n")
+            response = _error_response("laya_error", request.request_id)
+        _write_response(output_stream, response)