File size: 8,824 Bytes
750b88d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 | """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
|