349 lines
12 KiB
Python
349 lines
12 KiB
Python
"""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))
|