| """Base contract for memory methods. |
| |
| A method implements memorize() for each visit and build_memory_context() for each question. |
| The runner owns the answering call, so every method is answered with the same prompt, except |
| methods whose own agent produces the answer (ANSWERS_INTERNALLY). memorize() receives the |
| rendered visit text with its visit number and date. |
| """ |
| from __future__ import annotations |
|
|
| from abc import ABC, abstractmethod |
| from dataclasses import dataclass, field |
|
|
|
|
| def as_list(x) -> list: |
| """Coerce anything list-like (a numpy array included) to a plain list.""" |
| if x is None: |
| return [] |
| try: |
| return list(x) |
| except Exception: |
| return [] |
|
|
|
|
| class Tokenizer: |
| """Token counter backed by the tiktoken encoding named in backbone.yaml (`answer.tokenizer` for context budgets |
| and chunk sizes, `embedding.tokenizer` for embedding inputs); run_eval sets it before any method runs.""" |
|
|
| def __init__(self, role: str): |
| self.role, self._enc = role, None |
|
|
| def set(self, name: str | None) -> None: |
| if name and not (isinstance(name, str) and name.startswith("<")): |
| import tiktoken |
| self._enc = tiktoken.get_encoding(name) |
| else: |
| raise SystemExit(f"set {self.role}.tokenizer in backbone.yaml") |
|
|
| def _get(self): |
| if self._enc is None: |
| raise RuntimeError(f"{self.role}.tokenizer is not set (backbone.yaml)") |
| return self._enc |
|
|
| def encode(self, text: str, **kw) -> list[int]: |
| return self._get().encode(text, **kw) |
|
|
| def decode(self, tokens) -> str: |
| return self._get().decode(tokens) |
|
|
|
|
| ENC = Tokenizer("answer") |
| EMBED_ENC = Tokenizer("embedding") |
|
|
|
|
| @dataclass |
| class MemoryBuildResult: |
| """Result of one memorize() call.""" |
| success: bool = True |
| method: str = "" |
| action: str = "" |
| visit_number: int = -1 |
| chunk_count: int = 0 |
| added_chunks: int = 0 |
| time_cost: float = 0.0 |
| extra: dict = field(default_factory=dict) |
|
|
| def to_dict(self) -> dict: |
| return {"success": self.success, "method": self.method, "action": self.action, |
| "visit_number": self.visit_number, "chunk_count": self.chunk_count, |
| "added_chunks": self.added_chunks, "time_cost": round(self.time_cost, 4), |
| **self.extra} |
|
|
|
|
| @dataclass |
| class MemoryContext: |
| """The MEMORY block handed to the runner, or the answer of an internally answering method.""" |
| mem: str |
| n_tok: int = 0 |
| retrieved_count: int = 0 |
| retrieved_visits: list = field(default_factory=list) |
| retrieved_memories: list = field(default_factory=list) |
| context_truncated: bool = False |
| error: bool = False |
| answer: str | None = None |
| extra: dict = field(default_factory=dict) |
|
|
|
|
| class BaseAgent(ABC): |
| ANSWERS_INTERNALLY = False |
| MEMORIZE_WITH_LLM = False |
|
|
| def __init__(self, **params): |
| self.params = params |
| self._context_id: str | None = None |
|
|
| def set_trace_sink(self, sink) -> None: |
| """Install the callback that appends one event per dict to `memory_trace.jsonl`.""" |
| self._trace_sink = sink |
|
|
| def _trace(self, event: str, **fields) -> None: |
| sink = getattr(self, "_trace_sink", None) |
| if sink is None: |
| return |
| try: |
| from ..llm_client import USAGE |
| sink({"method": getattr(self, "METHOD_NAME", type(self).__name__), |
| "event": event, "phase": getattr(USAGE, "phase", None), **fields}) |
| except Exception as e: |
| n = getattr(self, "_trace_fail", 0) + 1 |
| self._trace_fail = n |
| if n == 1: |
| try: |
| sink({"method": getattr(self, "METHOD_NAME", type(self).__name__), |
| "event": "trace_error", "error": f"{type(e).__name__}: {e}"[:300], |
| "failed_event": event}) |
| except Exception: |
| pass |
|
|
| def consolidate(self) -> dict | None: |
| """Sleep-time memory maintenance, run at each checkpoint before answering.""" |
| return None |
|
|
| def dump_memory(self, out_dir) -> dict | None: |
| """Export the current memory state into out_dir; return a small stats dict, or None.""" |
| return None |
|
|
| def set_context_id(self, context_id: str) -> None: |
| """Set the per-patient namespace.""" |
| self._context_id = str(context_id) |
|
|
| @abstractmethod |
| def memorize(self, text: str, visit_number: int, date: str) -> MemoryBuildResult: |
| """Store one rendered visit into memory.""" |
|
|
| def memorize_session(self, session: dict, |
| rendered_text: str | None = None) -> MemoryBuildResult: |
| """Memorize one session from dataset.load_sessions().""" |
| if rendered_text is None: |
| from ..dataset import render_visit |
| rendered_text = render_visit(session) |
| return self.memorize(rendered_text, session["visit_number"], session["date"]) |
|
|
| @abstractmethod |
| def build_memory_context(self, question: str, as_of: str = "", |
| display_question: str | None = None) -> MemoryContext: |
| """Produce the MEMORY block for one question.""" |
|
|
| @staticmethod |
| def count_tokens(text: str) -> int: |
| return len(ENC.encode(text)) if text else 0 |
|
|
| @staticmethod |
| def cut_tokens(text: str, max_tokens: int) -> str: |
| """The first `max_tokens` tokens of `text`.""" |
| toks = ENC.encode(text) |
| return text if len(toks) <= max_tokens else ENC.decode(toks[:max_tokens]) |
|
|
| @staticmethod |
| def split_tokens(text: str, chunk_size: int, chunk_overlap: int, |
| header: str = "") -> list[str]: |
| """Split text into token windows; continuation chunks repeat `header` + '(cont.)'.""" |
| if chunk_overlap >= chunk_size: |
| raise ValueError(f"chunk_overlap ({chunk_overlap}) must be smaller " |
| f"than chunk_size ({chunk_size})") |
| if not text.strip(): |
| return [] |
| tokens = ENC.encode(text) |
| if len(tokens) <= chunk_size: |
| return [text] |
| chunks, start = [], 0 |
| while start < len(tokens): |
| end = min(start + chunk_size, len(tokens)) |
| piece = ENC.decode(tokens[start:end]) |
| if start > 0 and header: |
| piece = f"{header} (cont.)\n{piece}" |
| chunks.append(piece) |
| if end == len(tokens): |
| break |
| start = end - chunk_overlap |
| return chunks |
|
|
|
|
| def snapshot_jsonl_gz(path, rows) -> int: |
| """Write rows as gzipped JSONL; returns row count.""" |
| import gzip |
| import json as _json |
| with gzip.open(path, "wt", encoding="utf-8") as fh: |
| n = 0 |
| for r in rows: |
| fh.write(_json.dumps(r, ensure_ascii=False, default=str) + "\n") |
| n += 1 |
| return n |
|
|
|
|
| def snapshot_sqlite(src, dst) -> None: |
| """Copy an open SQLite database with the backup API.""" |
| import sqlite3 |
| a = sqlite3.connect(f"file:{src}?mode=ro", uri=True) |
| try: |
| b = sqlite3.connect(str(dst)) |
| try: |
| with b: |
| a.backup(b) |
| finally: |
| b.close() |
| finally: |
| a.close() |
|
|
|
|
| def snapshot_vectors_delta(out_dir, pairs, seen: set) -> int: |
| """Write the vectors this checkpoint added, and record them as seen.""" |
| import numpy as np |
| ids, vecs = [], [] |
| for mid, v in pairs: |
| mid = str(mid) |
| if not mid or mid in seen or v is None: |
| continue |
| seen.add(mid) |
| ids.append(mid) |
| vecs.append(np.asarray(v, dtype="float32")) |
| if not ids: |
| np.savez_compressed(out_dir / "vectors.npz", |
| ids=np.array([], dtype=object), vecs=np.zeros((0, 0), dtype="float32")) |
| return 0 |
| np.savez_compressed(out_dir / "vectors.npz", |
| ids=np.array(ids, dtype=object), vecs=np.vstack(vecs)) |
| return len(ids) |
|
|
|
|
| def snapshot_tree(src, dst) -> int: |
| """Copy a store directory, SQLite files through the backup API, skipping lock and tmp files; return the count.""" |
| import shutil |
| from pathlib import Path as _P |
| src, dst = _P(src), _P(dst) |
| n = 0 |
| for f in src.rglob("*"): |
| if not f.is_file() or f.name.endswith((".lock", "-wal", "-shm")): |
| continue |
| if "tmp" in f.relative_to(src).parts[:-1]: |
| continue |
| rel = dst / f.relative_to(src) |
| rel.parent.mkdir(parents=True, exist_ok=True) |
| if f.suffix in (".db", ".sqlite") or f.name.endswith(".sqlite3") or f.name == "storage.sqlite": |
| try: |
| snapshot_sqlite(f, rel) |
| except Exception: |
| shutil.copy2(f, rel) |
| else: |
| shutil.copy2(f, rel) |
| n += 1 |
| return n |
|
|
