cluster-knaster/core/src/knaster/store.py

"""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("canary_model", String(200)),
    Column("canary_share", Float),
    Column("canary_cluster", 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),
)

ontology_decisions = Table(
    "ontology_decisions",
    metadata,
    Column("key", String(300), primary_key=True),
    Column("decision", String(20), nullable=False),  # accept | map | reject
    Column("params", JSON),
    Column("decided_by", String(100)),
    Column("decided_at", Float),
    Column("applied_at", Float),
    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,
                  ts: float | None = None) -> None:
        with self.engine.begin() as conn:
            conn.execute(insert(events).values(ts=ts or 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)]

    # --- Ontologie-Entscheidungen ----------------------------------------------------------
    def set_ontology_decision(self, key: str, decision: str, params: dict[str, Any], decided_by: str) -> None:
        self._upsert(ontology_decisions, "key", key, {"decision": decision, "params": params,
                                                      "decided_by": decided_by, "decided_at": now(),
                                                      "applied_at": None})

    def list_ontology_decisions(self, *, applied: bool | None = None) -> list[dict[str, Any]]:
        q = select(ontology_decisions).order_by(ontology_decisions.c.decided_at)
        if applied is True:
            q = q.where(ontology_decisions.c.applied_at.is_not(None))
        elif applied is False:
            q = q.where(ontology_decisions.c.applied_at.is_(None))
        with self.engine.connect() as conn:
            return [dict(r._mapping) for r in conn.execute(q)]

    def mark_ontology_applied(self, keys: list[str]) -> None:
        with self.engine.begin() as conn:
            conn.execute(update(ontology_decisions).where(ontology_decisions.c.key.in_(keys))
                         .values(applied_at=now()))

    # --- 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))