"""Event-to-state orchestration, exact continuation snapshots, and model sessions.""" from __future__ import annotations import copy import json import os import tempfile from dataclasses import asdict from pathlib import Path from threading import RLock from .appraisal import appraise, EVENT_FACTS from .channels import ChannelDynamics from .config import ProjectConfig, finite from .core import AffectCore, EventSpec, Variable from .memory import EpisodicMemory EXTRA_EVENTS = { "warning_cue": EventSpec({}, 0.0), "uncertain_wait": EventSpec({"load": 0.03}, 0.0), "safety_confirmed": EventSpec({"autonomy": 0.08, "load": -0.15}, 0.0), "trade_small_loss": EventSpec({"integrity": -0.04, "load": 0.08}, 0.12), "trade_large_loss": EventSpec( {"integrity": -0.18, "load": 0.22, "autonomy": -0.08}, 0.5 ), "trade_gain": EventSpec({"integrity": 0.08, "load": -0.12}, 0.0), } class AffectEngine: def __init__(self, config=None): self.config = config or ProjectConfig() self.core = AffectCore() self.core.events.update(copy.deepcopy(EXTRA_EVENTS)) self.channels = ChannelDynamics() self.memory = EpisodicMemory(capacity=self.config.affect.memory_capacity) self.history = [] self._lock = RLock() @property def levels(self): with self._lock: return self.channels.read(self.core.pain) def advance(self, seconds): seconds = finite(seconds, "seconds") with self._lock: self.core.advance(seconds) self.channels.advance(seconds, self.config.affect) return self.levels def observe(self, event, intensity=1.0, cue=None, elapsed=0.0, appraisal=None): """Read a cue prediction BEFORE learning the outcome; never use future data.""" assessment = appraise(event, intensity, appraisal) finite(elapsed, "elapsed") if event not in self.core.events: raise KeyError(event) cue = cue or event if not isinstance(cue, str) or not cue.strip() or len(cue) > 256: raise ValueError("cue must contain 1-256 characters") with self._lock: self.advance(elapsed) before = self.levels expected = self.memory.predict(cue, self.core.clock, self.config.affect) self.core.apply(event, intensity) self.channels.update(assessment, expected, before.pain, self.core.pain) episode = self.memory.observe( cue, assessment.harm, self.core.clock, self.config.affect, event, self.levels.as_dict(), ) record = { "sequence": episode["id"], "t": self.core.clock, "event": event, "intensity": intensity, "cue": cue, "appraisal": asdict(assessment), "before": before.as_dict(), "levels": self.levels.as_dict(), "expected_harm_before": expected, "prediction_error": episode["prediction_error"], } self.history.append(record) self.history = self.history[-self.config.affect.history_capacity :] self.core.history = self.core.history[ -self.config.affect.history_capacity : ] return copy.deepcopy(record) def snapshot(self): with self._lock: return copy.deepcopy( { "schema_version": 2, "config": self.config.to_dict(), "core": self.core.snapshot(), "channels": asdict(self.channels), "memory": self.memory.to_dict(), "history": self.history, } ) def request_context(self, limit=6): """Atomically read levels and factual observations, without affect labels.""" if not isinstance(limit, int) or limit < 0: raise ValueError("Observation limit must be a nonnegative integer") with self._lock: records = self.history[-limit:] if limit else [] facts = [ { "sequence": row["sequence"], "t": row["t"], "description": EVENT_FACTS[row["event"]], } for row in records ] return self.levels, facts @classmethod def from_snapshot(cls, state): state = copy.deepcopy(state) if state.get("schema_version") != 2: raise ValueError("Unsupported session schema") engine = cls(ProjectConfig.from_dict(state["config"])) core = state["core"] engine.core = AffectCore( variables=[Variable(**v) for v in core["variables"]], events={k: EventSpec(**v) for k, v in core["events"].items()}, **core.get("parameters", {}), ) engine.core.clock = finite(core["clock"], "clock") engine.core.transient = finite(core["transient"], "transient", 0, 1) engine.core.last_event = core["last_event"] engine.core.expectations = { k: finite(v, "expectation", 0, 1) for k, v in core["expectations"].items() } engine.channels = ChannelDynamics( **{k: finite(v, k, 0, 1) for k, v in state["channels"].items()} ) engine.memory = EpisodicMemory.from_dict(state["memory"]) engine.history = state["history"][-engine.config.affect.history_capacity :] if engine.memory.capacity != engine.config.affect.memory_capacity: raise ValueError("Memory capacity differs from saved configuration") return engine def save(self, path): path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) handle, temporary = tempfile.mkstemp( prefix=path.name, suffix=".tmp", dir=path.parent ) try: with os.fdopen(handle, "w", encoding="utf-8") as stream: json.dump( self.snapshot(), stream, ensure_ascii=False, indent=2, allow_nan=False, ) os.replace(temporary, path) finally: if os.path.exists(temporary): os.unlink(temporary) @classmethod def load(cls, path): return cls.from_snapshot(json.loads(Path(path).read_text(encoding="utf-8"))) class AffectSession: """A single conversation; the model's system persona is state-independent.""" def __init__( self, model, tokenizer, bank, engine=None, history_turns=12, include_event_facts=True, ): from .steering import MultiChannelSteerer if history_turns < 1: raise ValueError("history_turns must be positive") self.model, self.tokenizer = model, tokenizer self.engine = engine or AffectEngine() self.history, self.history_turns = [], history_turns self.include_event_facts = include_event_facts self.last_request = None self.controller = MultiChannelSteerer( model, bank, lambda: self.engine.levels, self.engine.config.steering ) self._lock = RLock() def respond(self, user, **generation): import hashlib from .model import generate, format_prompt from .persona import CODELOCK_SYSTEM with self._lock: settings = asdict(self.engine.config.generation) settings.update(generation) levels, facts = self.engine.request_context() if not self.include_event_facts: facts = [] system = settings.pop("system", CODELOCK_SYSTEM) if facts: system += "\nObserved event record (authoritative facts; do not add unobserved consequences):\n" system += "\n".join( f"- At t={row['t']:g}s: {row['description']}" for row in facts ) prompt = format_prompt(self.tokenizer, user, self.history, system) self.last_request = { "levels": levels.as_dict(), "observations": facts, "prompt_sha256": hashlib.sha256(prompt.encode()).hexdigest(), } with self.controller, self.controller.frozen(levels): text = generate( self.model, self.tokenizer, user, self.history, system=system, **settings, ) self.history.extend( [ {"role": "user", "content": user}, {"role": "assistant", "content": text}, ] ) self.history = self.history[-2 * self.history_turns :] return text