"""Persistenz (SQLAlchemy Core). SQLite für Tests/kleine Setups, Postgres im Betrieb.""" from __future__ import annotations import uuid from typing import Any from sqlalchemy import ( JSON, Boolean, Column, Float, Integer, MetaData, String, Table, Text, create_engine, delete, func, insert, select, update, ) from sqlalchemy.engine import Engine from sqlalchemy.pool import StaticPool from .util import now metadata = MetaData() cases = Table( "cases", metadata, Column("id", String(36), primary_key=True), Column("ts", Float, nullable=False, index=True), Column("run_id", String(36), index=True), Column("input", Text, nullable=False), Column("skill", String(200), index=True), Column("skill_version", Integer), Column("cluster", String(200), index=True), Column("stage", Integer, nullable=False), Column("slot", String(100)), Column("model", String(200)), Column("adapter", String(200)), Column("messages", JSON), Column("output", Text), Column("passed", Boolean), Column("signals", JSON), Column("prompt_tokens", Integer), Column("split", String(10), nullable=False), Column("origin", String(20), nullable=False), ) preference_pairs = Table( "preference_pairs", metadata, Column("id", String(36), primary_key=True), Column("ts", Float, nullable=False), Column("run_id", String(36)), Column("skill", String(200), index=True), Column("cluster", String(200)), Column("prompt", JSON), Column("chosen", Text), Column("rejected", Text), Column("split", String(10)), ) proposals = Table( "proposals", metadata, Column("seq", Integer, primary_key=True, autoincrement=True), Column("kind", String(50), nullable=False), Column("params", JSON), Column("summary", Text), Column("reason", Text), Column("risk", String(10)), Column("state", String(20), nullable=False, index=True), Column("source", String(20)), Column("dedup_key", String(300), index=True), Column("created_at", Float), Column("decided_at", Float), Column("decided_by", String(100)), Column("result", JSON), ) slot_state = Table( "slot_state", metadata, Column("slot", String(100), primary_key=True), Column("status", String(20)), Column("effective_context", Integer), Column("base_model", String(200)), Column("reason", Text), Column("set_by", String(50)), Column("updated_at", Float), ) adapter_state = Table( "adapter_state", metadata, Column("name", String(200), primary_key=True), Column("base_model", String(200)), Column("cluster", String(200)), Column("version", Integer), Column("status", String(20)), Column("updated_at", Float), ) skill_state = Table( "skill_state", metadata, Column("name", String(200), primary_key=True), Column("status", String(20)), Column("stage", Integer), Column("reason", Text), Column("updated_at", Float), ) candidates = Table( "candidates", metadata, Column("id", String(36), primary_key=True), Column("kind", String(30)), Column("slot", String(100)), Column("target", JSON), Column("baseline", JSON), Column("state", String(30), index=True), Column("metrics", JSON), Column("history", JSON), Column("snapshot", JSON), Column("created_at", Float), Column("updated_at", Float), ) healer_state = Table( "healer_state", metadata, Column("key", String(300), primary_key=True), Column("data", JSON), Column("updated_at", Float), ) events = Table( "events", metadata, Column("id", Integer, primary_key=True, autoincrement=True), Column("ts", Float, nullable=False, index=True), Column("kind", String(50), nullable=False, index=True), Column("subject", String(200), index=True), Column("data", JSON), ) def new_id() -> str: return str(uuid.uuid4()) class Store: """Dünne Repository-Schicht. Jede Methode öffnet ihre eigene Transaktion.""" def __init__(self, url: str): kwargs: dict[str, Any] = {} if url.startswith("sqlite"): kwargs["connect_args"] = {"check_same_thread": False} if url in {"sqlite://", "sqlite:///:memory:"}: kwargs["poolclass"] = StaticPool self.engine: Engine = create_engine(url, future=True, **kwargs) metadata.create_all(self.engine) # --- Fälle ----------------------------------------------------------------------------- def add_case(self, **row: Any) -> str: row.setdefault("id", new_id()) row.setdefault("ts", now()) with self.engine.begin() as conn: conn.execute(insert(cases).values(**row)) return row["id"] def list_cases( self, *, skill: str | None = None, cluster: str | None = None, min_stage: int | None = None, stage: int | None = None, split: str | None = None, passed: bool | None = None, since: float | None = None, limit: int | None = None, ) -> list[dict[str, Any]]: q = select(cases).order_by(cases.c.ts.desc()) if skill is not None: q = q.where(cases.c.skill == skill) if cluster is not None: q = q.where(cases.c.cluster == cluster) if min_stage is not None: q = q.where(cases.c.stage >= min_stage) if stage is not None: q = q.where(cases.c.stage == stage) if split is not None: q = q.where(cases.c.split == split) if passed is not None: q = q.where(cases.c.passed == passed) if since is not None: q = q.where(cases.c.ts >= since) if limit is not None: q = q.limit(limit) with self.engine.connect() as conn: return [dict(r._mapping) for r in conn.execute(q)] def add_preference_pair(self, **row: Any) -> str: row.setdefault("id", new_id()) row.setdefault("ts", now()) with self.engine.begin() as conn: conn.execute(insert(preference_pairs).values(**row)) return row["id"] def list_preference_pairs(self, *, cluster: str | None = None) -> list[dict[str, Any]]: q = select(preference_pairs).order_by(preference_pairs.c.ts) if cluster is not None: q = q.where(preference_pairs.c.cluster == cluster) with self.engine.connect() as conn: return [dict(r._mapping) for r in conn.execute(q)] # --- Ereignisse ------------------------------------------------------------------------ def add_event(self, kind: str, subject: str | None = None, data: dict | None = None) -> None: with self.engine.begin() as conn: conn.execute(insert(events).values(ts=now(), kind=kind, subject=subject, data=data or {})) def count_events(self, kind: str, *, subject: str | None = None, since: float | None = None) -> int: q = select(func.count()).select_from(events).where(events.c.kind == kind) if subject is not None: q = q.where(events.c.subject == subject) if since is not None: q = q.where(events.c.ts >= since) with self.engine.connect() as conn: return int(conn.execute(q).scalar_one()) def event_subjects(self, kind: str, *, since: float | None = None) -> dict[str, int]: q = select(events.c.subject, func.count()).where(events.c.kind == kind).group_by(events.c.subject) if since is not None: q = q.where(events.c.ts >= since) with self.engine.connect() as conn: return {r[0]: int(r[1]) for r in conn.execute(q) if r[0] is not None} # --- generische Zustandstabellen ----------------------------------------------------- def _upsert(self, table: Table, key_col: str, key: str, values: dict[str, Any]) -> None: values = {**values, "updated_at": now()} with self.engine.begin() as conn: exists = conn.execute(select(table).where(table.c[key_col] == key)).first() if exists: conn.execute(update(table).where(table.c[key_col] == key).values(**values)) else: conn.execute(insert(table).values(**{key_col: key, **values})) def _get(self, table: Table, key_col: str, key: str) -> dict[str, Any] | None: with self.engine.connect() as conn: row = conn.execute(select(table).where(table.c[key_col] == key)).first() return dict(row._mapping) if row else None def _all(self, table: Table) -> list[dict[str, Any]]: with self.engine.connect() as conn: return [dict(r._mapping) for r in conn.execute(select(table))] def set_slot_state(self, slot: str, **values: Any) -> None: self._upsert(slot_state, "slot", slot, values) def get_slot_state(self, slot: str) -> dict[str, Any] | None: return self._get(slot_state, "slot", slot) def all_slot_states(self) -> list[dict[str, Any]]: return self._all(slot_state) def set_adapter_state(self, name: str, **values: Any) -> None: self._upsert(adapter_state, "name", name, values) def all_adapter_states(self) -> list[dict[str, Any]]: return self._all(adapter_state) def set_skill_state(self, name: str, **values: Any) -> None: self._upsert(skill_state, "name", name, values) def all_skill_states(self) -> list[dict[str, Any]]: return self._all(skill_state) # --- Vorschläge ------------------------------------------------------------------------ def add_proposal(self, **row: Any) -> int: row.setdefault("created_at", now()) row.setdefault("state", "open") with self.engine.begin() as conn: res = conn.execute(insert(proposals).values(**row)) return int(res.inserted_primary_key[0]) def get_proposal(self, seq: int) -> dict[str, Any] | None: return self._get(proposals, "seq", seq) # type: ignore[arg-type] def update_proposal(self, seq: int, **values: Any) -> None: with self.engine.begin() as conn: conn.execute(update(proposals).where(proposals.c.seq == seq).values(**values)) def list_proposals(self, *, state: str | None = None) -> list[dict[str, Any]]: q = select(proposals).order_by(proposals.c.seq) if state is not None: q = q.where(proposals.c.state == state) with self.engine.connect() as conn: return [dict(r._mapping) for r in conn.execute(q)] def open_proposal_with_key(self, dedup_key: str) -> dict[str, Any] | None: q = select(proposals).where(proposals.c.dedup_key == dedup_key, proposals.c.state == "open") with self.engine.connect() as conn: row = conn.execute(q).first() return dict(row._mapping) if row else None # --- Kandidaten ------------------------------------------------------------------------ def add_candidate(self, **row: Any) -> str: row.setdefault("id", new_id()) row.setdefault("created_at", now()) row.setdefault("updated_at", row["created_at"]) with self.engine.begin() as conn: conn.execute(insert(candidates).values(**row)) return row["id"] def get_candidate(self, cid: str) -> dict[str, Any] | None: return self._get(candidates, "id", cid) def update_candidate(self, cid: str, **values: Any) -> None: values["updated_at"] = now() with self.engine.begin() as conn: conn.execute(update(candidates).where(candidates.c.id == cid).values(**values)) def list_candidates(self, *, state: str | None = None) -> list[dict[str, Any]]: q = select(candidates).order_by(candidates.c.created_at) if state is not None: q = q.where(candidates.c.state == state) with self.engine.connect() as conn: return [dict(r._mapping) for r in conn.execute(q)] # --- Healer-Zustand -------------------------------------------------------------------- def get_healer_state(self, key: str) -> dict[str, Any] | None: row = self._get(healer_state, "key", key) return row["data"] if row else None def set_healer_state(self, key: str, data: dict[str, Any]) -> None: self._upsert(healer_state, "key", key, {"data": data}) def delete_healer_state(self, key: str) -> None: with self.engine.begin() as conn: conn.execute(delete(healer_state).where(healer_state.c.key == key))