evalkit-py 0.2.0.dev12__py3-none-any.whl

evalkit/db/promotion.py

@@ -0,0 +1,271 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any
+
+from sqlalchemy import delete, func, select
+from sqlalchemy.dialects.postgresql import insert
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from evalkit.contracts.enums import SyncSource
+from evalkit.core.access import WorkspaceRef
+from evalkit.core.promotion import (
+    MissingTargetDataset,
+    OverlayCopy,
+    PromotionPlan,
+    twin_summary,
+)
+from evalkit.core.review import CaseEvent, CaseEventKind, CaseReview
+from evalkit.db import models
+from evalkit.db.cases import find_dataset_row
+from evalkit.db.review_state import latest_events
+from evalkit.db.workspaces import load_ref
+from evalkit.store.base import CaseView, DatasetRow
+
+
+async def overlays_extending(
+    session: AsyncSession, workspace_id: int, case_refs: list[int]
+) -> dict[int, list[OverlayCopy]]:
+    if not case_refs:
+        return {}
+    rows = (
+        await session.execute(
+            select(models.Overlay, models.Dataset.name)
+            .join(models.Dataset, models.Dataset.id == models.Overlay.source_dataset_id)
+            .where(
+                models.Dataset.workspace_id == workspace_id,
+                models.Overlay.target_case.in_(case_refs),
+            )
+            .order_by(models.Dataset.name)
+        )
+    ).all()
+    out: dict[int, list[OverlayCopy]] = {}
+    for overlay, source in rows:
+        out.setdefault(overlay.target_case, []).append(
+            OverlayCopy(
+                source_dataset=source,
+                agents=tuple(overlay.agents),
+                tags=tuple(overlay.tags),
+                evaluators=overlay.evaluators,
+                source_file=overlay.source_file,
+            )
+        )
+    return out
+
+
+async def _workspace(session: AsyncSession, workspace_id: int) -> WorkspaceRef:
+    return await load_ref(session, await session.get_one(models.Workspace, workspace_id))
+
+
+async def _username(session: AsyncSession, user_id: int | None) -> str | None:
+    if user_id is None:
+        return None
+    return await session.scalar(select(models.User.username).where(models.User.id == user_id))
+
+
+async def _placement(
+    session: AsyncSession, dataset_id: int, case_ref: int
+) -> tuple[list[str], str | None]:
+    row = (
+        await session.execute(
+            select(models.DatasetCase.path_tags, models.DatasetCase.source).where(
+                models.DatasetCase.dataset_id == dataset_id,
+                models.DatasetCase.case_ref == case_ref,
+            )
+        )
+    ).one()
+    return list(row[0]), row[1]
+
+
+async def plan_promotion(
+    session: AsyncSession, source_ws: WorkspaceRef, dataset: DatasetRow, case: CaseView
+) -> PromotionPlan:
+    assert source_ws.promotes_to is not None
+    target_ws = await load_ref(
+        session, await session.get_one(models.Workspace, source_ws.promotes_to)
+    )
+    target = await find_dataset_row(session, target_ws.name, dataset.name)
+    if target is None or target.kind != "suite":
+        raise MissingTargetDataset(f"{dataset.name}@{target_ws.name}")
+    overlays = (await overlays_extending(session, source_ws.id, [case.case_ref])).get(
+        case.case_ref, []
+    )
+    overlay_targets: dict[str, int] = {}
+    for name in sorted({o.source_dataset for o in overlays}):
+        found = await find_dataset_row(session, target_ws.name, name)
+        if found is None:
+            raise MissingTargetDataset(f"{name}@{target_ws.name}")
+        overlay_targets[name] = found.id
+    path_tags, source = await _placement(session, dataset.id, case.case_ref)
+    return PromotionPlan(
+        source_ref=dataset.ref,
+        target_ref=target.ref,
+        target_workspace_id=target_ws.id,
+        target_dataset_id=target.id,
+        case_ref=case.case_ref,
+        case_id=case.case_id,
+        content_hash=case.content_hash,
+        path_tags=path_tags,
+        source=source,
+        overlays=overlays,
+        overlay_targets=overlay_targets,
+    )
+
+
+async def execute_promotion(
+    session: AsyncSession, plan: PromotionPlan, *, actor_id: int, event_id: int
+) -> int:
+    sync = models.DatasetSync(
+        dataset_id=plan.target_dataset_id,
+        actor_id=actor_id,
+        source=SyncSource.PROMOTION.value,
+        detail={
+            "from": plan.source_ref,
+            "case_id": plan.case_id,
+            "content_hash": plan.content_hash,
+            "event_id": event_id,
+            "overlays": len(plan.overlays),
+        },
+    )
+    session.add(sync)
+    await session.flush()
+    current = (
+        await session.execute(
+            select(models.DatasetCase.content_hash, models.DatasetCase.removed_at).where(
+                models.DatasetCase.dataset_id == plan.target_dataset_id,
+                models.DatasetCase.case_ref == plan.case_ref,
+            )
+        )
+    ).one_or_none()
+    if current is None or current[0] != plan.content_hash or current[1] is not None:
+        statement = insert(models.DatasetCase).values(
+            dataset_id=plan.target_dataset_id,
+            case_ref=plan.case_ref,
+            content_hash=plan.content_hash,
+            path_tags=plan.path_tags,
+            source=plan.source,
+            removed_at=None,
+        )
+        await session.execute(
+            statement.on_conflict_do_update(
+                index_elements=["dataset_id", "case_ref"],
+                set_={
+                    "content_hash": statement.excluded.content_hash,
+                    "path_tags": statement.excluded.path_tags,
+                    "source": statement.excluded.source,
+                    "removed_at": None,
+                    "updated_at": func.now(),
+                },
+            )
+        )
+        for kind, sync_id in (
+            (CaseEventKind.SYNCED, sync.id),
+            (CaseEventKind.APPROVED, None),
+        ):
+            session.add(
+                models.CaseEvent(
+                    dataset_id=plan.target_dataset_id,
+                    case_ref=plan.case_ref,
+                    content_hash=plan.content_hash,
+                    kind=kind.value,
+                    sync_id=sync_id,
+                    actor_id=actor_id,
+                )
+            )
+            await session.flush()
+    await session.execute(
+        delete(models.Overlay).where(
+            models.Overlay.target_case == plan.case_ref,
+            models.Overlay.source_dataset_id.in_(
+                select(models.Dataset.id).where(
+                    models.Dataset.workspace_id == plan.target_workspace_id
+                )
+            ),
+        )
+    )
+    for overlay in plan.overlays:
+        session.add(
+            models.Overlay(
+                source_dataset_id=plan.overlay_targets[overlay.source_dataset],
+                target_case=plan.case_ref,
+                agents=list(overlay.agents),
+                tags=list(overlay.tags),
+                evaluators=overlay.evaluators,
+                source_file=overlay.source_file,
+            )
+        )
+    await session.flush()
+    return sync.id
+
+
+async def _review_of(
+    session: AsyncSession, dataset: DatasetRow, case_ref: int, case_id: str, content_hash: str
+) -> CaseReview:
+    events = await latest_events(session, dataset.id, {case_ref: content_hash})
+    review = CaseReview.unreviewed(dataset.ref, case_id, content_hash)
+    row = events.decisive.get(case_ref)
+    if row is not None:
+        actor = await _username(session, row.actor_id)
+        assert actor is not None
+        review.last_event = CaseEvent(
+            id=row.id,
+            dataset=dataset.ref,
+            case_id=case_id,
+            content_hash=row.content_hash,
+            kind=CaseEventKind(row.kind),
+            actor=actor,
+            target=await _username(session, row.target_id),
+            created_at=row.created_at,
+        )
+    review.requested_to = await _username(session, events.requested_to.get(case_ref))
+    return review
+
+
+async def twin_status(
+    session: AsyncSession, ws: WorkspaceRef, dataset: DatasetRow, cases: list[CaseView]
+) -> dict[str, dict[str, Any] | None]:
+    other_id = ws.promotes_to if ws.promotes_to is not None else ws.promoted_from
+    if other_id is None or dataset.kind != "suite" or not cases:
+        return {case.case_id: None for case in cases}
+    other_ws = await _workspace(session, other_id)
+    other = await find_dataset_row(session, other_ws.name, dataset.name)
+    refs = [case.case_ref for case in cases]
+    theirs: dict[int, tuple[str, datetime | None]] = {}
+    if other is not None:
+        theirs = {
+            ref: (content_hash, removed_at)
+            for ref, content_hash, removed_at in (
+                await session.execute(
+                    select(
+                        models.DatasetCase.case_ref,
+                        models.DatasetCase.content_hash,
+                        models.DatasetCase.removed_at,
+                    ).where(
+                        models.DatasetCase.dataset_id == other.id,
+                        models.DatasetCase.case_ref.in_(refs),
+                    )
+                )
+            ).all()
+        }
+    mine_overlays = await overlays_extending(session, ws.id, refs)
+    their_overlays = await overlays_extending(session, other_ws.id, refs)
+    out: dict[str, dict[str, Any] | None] = {}
+    for case in cases:
+        found = theirs.get(case.case_ref)
+        other_hash = found[0] if found is not None and found[1] is None else None
+        review = (
+            await _review_of(session, other, case.case_ref, case.case_id, other_hash)
+            if other is not None and other_hash is not None
+            else None
+        )
+        out[case.case_id] = twin_summary(
+            workspace=other_ws.name,
+            dataset=dataset.name,
+            case_id=case.case_id,
+            mine=case.content_hash,
+            other=other_hash,
+            review=review,
+            my_overlays=mine_overlays.get(case.case_ref, []),
+            other_overlays=their_overlays.get(case.case_ref, []),
+        )
+    return out

evalkit/db/rekey.py

@@ -0,0 +1,46 @@
+from sqlalchemy import func, select
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from evalkit.db import models
+from evalkit.db.crypto import SecretBox
+
+type Sealed = models.TraceSource | models.Agent | models.RemoteEvaluator | models.LlmProvider
+
+SEALED: tuple[type[Sealed], ...] = (
+    models.TraceSource,
+    models.Agent,
+    models.RemoteEvaluator,
+    models.LlmProvider,
+)
+
+
+async def count_sealed(session_factory: async_sessionmaker[AsyncSession]) -> int:
+    count = 0
+    async with session_factory() as session:
+        for model in SEALED:
+            count += int(
+                await session.scalar(
+                    select(func.count()).where(model.api_key_encrypted.is_not(None))
+                )
+                or 0
+            )
+    return count
+
+
+async def _rekey[T: Sealed](
+    session: AsyncSession, model: type[T], old: SecretBox, new: SecretBox
+) -> int:
+    rows = (await session.scalars(select(model).where(model.api_key_encrypted.is_not(None)))).all()
+    for row in rows:
+        assert row.api_key_encrypted is not None
+        row.api_key_encrypted = new.seal(old.open(row.api_key_encrypted))
+    return len(rows)
+
+
+async def rekey_all(
+    session_factory: async_sessionmaker[AsyncSession], old: SecretBox, new: SecretBox
+) -> int:
+    async with session_factory() as session:
+        count = sum([await _rekey(session, model, old, new) for model in SEALED])
+        await session.commit()
+    return count

evalkit/db/review.py

@@ -0,0 +1,250 @@
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+
+from sqlalchemy import select, update
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from evalkit.core.review import (
+    CaseComment,
+    CaseEvent,
+    CaseEventKind,
+    CaseReview,
+    CaseVersion,
+    NewCaseComment,
+    NewCaseEvent,
+)
+from evalkit.core.workspace import Reference, slug
+from evalkit.db import models
+from evalkit.db.cases import dataset_by_ref
+from evalkit.db.review_state import latest_events
+
+AfterDecision = Callable[[AsyncSession, int], Awaitable[None]]
+
+
+async def _case_refs(session: AsyncSession, case_ids: list[str]) -> dict[str, int]:
+    rows = await session.execute(
+        select(models.Case.case_id, models.Case.id).where(models.Case.case_id.in_(case_ids))
+    )
+    return {case_id: ref for case_id, ref in rows.all()}
+
+
+async def _user_ids(session: AsyncSession, names: set[str]) -> dict[str, int]:
+    rows = await session.execute(
+        select(models.User.username, models.User.id).where(models.User.username.in_(names))
+    )
+    return {name: uid for name, uid in rows.all()}
+
+
+async def _usernames(session: AsyncSession, ids: set[int]) -> dict[int, str]:
+    rows = await session.execute(
+        select(models.User.id, models.User.username).where(models.User.id.in_(ids))
+    )
+    return {uid: name for uid, name in rows.all()}
+
+
+class ReviewLog:
+    def __init__(self, session_factory: async_sessionmaker[AsyncSession]):
+        self._session_factory = session_factory
+
+    async def _events(
+        self, session: AsyncSession, dataset: str, case_id: str, rows: list[models.CaseEvent]
+    ) -> list[CaseEvent]:
+        users = {r.actor_id for r in rows} | {r.target_id for r in rows if r.target_id}
+        names = await _usernames(session, users)
+        return [
+            CaseEvent(
+                id=r.id,
+                dataset=dataset,
+                case_id=case_id,
+                content_hash=r.content_hash,
+                kind=CaseEventKind(r.kind),
+                actor=names[r.actor_id],
+                target=names[r.target_id] if r.target_id else None,
+                created_at=r.created_at,
+            )
+            for r in rows
+        ]
+
+    async def _comment(
+        self, session: AsyncSession, row: models.CaseComment, dataset: str, case_id: str
+    ) -> CaseComment:
+        names = await _usernames(session, {row.author_id})
+        return CaseComment(
+            id=row.id,
+            dataset=dataset,
+            case_id=case_id,
+            content_hash=row.content_hash,
+            parent_id=row.parent_id,
+            author=names[row.author_id],
+            body=row.body,
+            resolved=row.resolved,
+            created_at=row.created_at,
+        )
+
+    async def load(self, dataset: str, versions: list[CaseVersion]) -> dict[str, CaseReview]:
+        if not versions:
+            return {}
+        reviews = {
+            v.case_id: CaseReview.unreviewed(dataset, v.case_id, v.content_hash) for v in versions
+        }
+        async with self._session_factory() as session:
+            dataset_id = (await dataset_by_ref(session, dataset)).id
+            refs = await _case_refs(session, [v.case_id for v in versions])
+            ids = {ref: case_id for case_id, ref in refs.items()}
+            latest = await latest_events(
+                session,
+                dataset_id,
+                {refs[v.case_id]: v.content_hash for v in versions if v.case_id in refs},
+            )
+            for ref, row in latest.decisive.items():
+                [event] = await self._events(session, dataset, ids[ref], [row])
+                reviews[ids[ref]].last_event = event
+            names = await _usernames(session, set(latest.requested_to.values()))
+        for ref, target in latest.requested_to.items():
+            reviews[ids[ref]].requested_to = names[target]
+        return reviews
+
+    async def latest_hashes(self, dataset: str, case_ids: list[str]) -> dict[str, str]:
+        if not case_ids:
+            return {}
+        async with self._session_factory() as session:
+            dataset_id = (await dataset_by_ref(session, dataset)).id
+            rows = await session.execute(
+                select(models.Case.case_id, models.DatasetCase.content_hash)
+                .join(models.Case, models.Case.id == models.DatasetCase.case_ref)
+                .where(
+                    models.DatasetCase.dataset_id == dataset_id,
+                    models.Case.case_id.in_(case_ids),
+                )
+            )
+            return {case_id: content_hash for case_id, content_hash in rows.all()}
+
+    async def _add_events(
+        self, session: AsyncSession, events: list[NewCaseEvent]
+    ) -> list[models.CaseEvent]:
+        refs = await _case_refs(session, [e.case_id for e in events])
+        users = await _user_ids(
+            session, {e.actor for e in events} | {e.target for e in events if e.target}
+        )
+        rows: list[models.CaseEvent] = []
+        for event in events:
+            dataset_id = (await dataset_by_ref(session, event.dataset)).id
+            row = models.CaseEvent(
+                dataset_id=dataset_id,
+                case_ref=refs[event.case_id],
+                content_hash=event.content_hash,
+                kind=event.kind.value,
+                actor_id=users[event.actor],
+                target_id=users[event.target] if event.target else None,
+            )
+            session.add(row)
+            await session.flush()
+            rows.append(row)
+        return rows
+
+    async def _comment_row(
+        self, session: AsyncSession, comment: NewCaseComment
+    ) -> models.CaseComment:
+        dataset_id = (await dataset_by_ref(session, comment.dataset)).id
+        refs = await _case_refs(session, [comment.case_id])
+        users = await _user_ids(session, {comment.author})
+        return models.CaseComment(
+            dataset_id=dataset_id,
+            case_ref=refs[comment.case_id],
+            content_hash=comment.content_hash,
+            parent_id=comment.parent_id,
+            author_id=users[comment.author],
+            body=comment.body,
+        )
+
+    async def append(self, events: list[NewCaseEvent]) -> None:
+        if not events:
+            return
+        async with self._session_factory() as session:
+            await self._add_events(session, events)
+            await session.commit()
+
+    async def decide(
+        self,
+        event: NewCaseEvent,
+        comment: NewCaseComment | None,
+        after: AfterDecision | None = None,
+    ) -> None:
+        async with self._session_factory() as session:
+            if comment is not None:
+                session.add(await self._comment_row(session, comment))
+            [row] = await self._add_events(session, [event])
+            if after is not None:
+                await after(session, row.id)
+            await session.commit()
+
+    async def history(self, dataset: str, case_id: str) -> list[CaseEvent]:
+        async with self._session_factory() as session:
+            dataset_id = (await dataset_by_ref(session, dataset)).id
+            rows = (
+                await session.scalars(
+                    select(models.CaseEvent)
+                    .join(models.Case, models.Case.id == models.CaseEvent.case_ref)
+                    .where(
+                        models.CaseEvent.dataset_id == dataset_id, models.Case.case_id == case_id
+                    )
+                    .order_by(models.CaseEvent.created_at.desc(), models.CaseEvent.id.desc())
+                )
+            ).all()
+            return await self._events(session, dataset, case_id, list(rows))
+
+    async def comments(self, dataset: str, case_id: str) -> list[CaseComment]:
+        async with self._session_factory() as session:
+            dataset_id = (await dataset_by_ref(session, dataset)).id
+            rows = (
+                await session.scalars(
+                    select(models.CaseComment)
+                    .join(models.Case, models.Case.id == models.CaseComment.case_ref)
+                    .where(
+                        models.CaseComment.dataset_id == dataset_id,
+                        models.Case.case_id == case_id,
+                    )
+                    .order_by(models.CaseComment.created_at.asc(), models.CaseComment.id.asc())
+                )
+            ).all()
+            return [await self._comment(session, row, dataset, case_id) for row in rows]
+
+    async def add_comment(self, comment: NewCaseComment) -> CaseComment:
+        async with self._session_factory() as session:
+            row = await self._comment_row(session, comment)
+            session.add(row)
+            await session.commit()
+            await session.refresh(row)
+            return await self._comment(session, row, comment.dataset, comment.case_id)
+
+    async def _describe(self, session: AsyncSession, row: models.CaseComment) -> CaseComment:
+        name, workspace = (
+            await session.execute(
+                select(models.Dataset.name, models.Workspace.name)
+                .join(models.Workspace, models.Workspace.id == models.Dataset.workspace_id)
+                .where(models.Dataset.id == row.dataset_id)
+            )
+        ).one()
+        case_id = await session.scalar(
+            select(models.Case.case_id).where(models.Case.id == row.case_ref)
+        )
+        assert case_id is not None
+        return await self._comment(session, row, slug(Reference(name, workspace)), case_id)
+
+    async def get_comment(self, comment_id: int) -> CaseComment | None:
+        async with self._session_factory() as session:
+            row = await session.get(models.CaseComment, comment_id)
+            return None if row is None else await self._describe(session, row)
+
+    async def set_resolved(self, comment_id: int, resolved: bool) -> CaseComment:
+        async with self._session_factory() as session:
+            await session.execute(
+                update(models.CaseComment)
+                .where(models.CaseComment.id == comment_id)
+                .values(resolved=resolved)
+            )
+            await session.commit()
+            row = await session.get(models.CaseComment, comment_id)
+            assert row is not None
+            return await self._describe(session, row)

evalkit/db/review_state.py

@@ -0,0 +1,50 @@
+from dataclasses import dataclass, field
+
+from sqlalchemy import select, tuple_
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from evalkit.core.review import DECISIVE_EVENT_KINDS, CaseEventKind
+from evalkit.db import models
+
+
+@dataclass
+class LatestEvents:
+    decisive: dict[int, models.CaseEvent] = field(default_factory=dict)
+    requested_to: dict[int, int] = field(default_factory=dict)
+
+
+async def _latest(
+    session: AsyncSession, dataset_id: int, versions: dict[int, str], kinds: list[str]
+) -> dict[int, models.CaseEvent]:
+    rows = await session.scalars(
+        select(models.CaseEvent)
+        .where(
+            models.CaseEvent.dataset_id == dataset_id,
+            tuple_(models.CaseEvent.case_ref, models.CaseEvent.content_hash).in_(
+                list(versions.items())
+            ),
+            models.CaseEvent.kind.in_(kinds),
+        )
+        .distinct(models.CaseEvent.case_ref)
+        .order_by(
+            models.CaseEvent.case_ref,
+            models.CaseEvent.created_at.desc(),
+            models.CaseEvent.id.desc(),
+        )
+    )
+    return {row.case_ref: row for row in rows.all()}
+
+
+async def latest_events(
+    session: AsyncSession, dataset_id: int, versions: dict[int, str]
+) -> LatestEvents:
+    if not versions:
+        return LatestEvents()
+    decisive = await _latest(session, dataset_id, versions, [k.value for k in DECISIVE_EVENT_KINDS])
+    requests = await _latest(session, dataset_id, versions, [CaseEventKind.REVIEW_REQUESTED.value])
+    return LatestEvents(
+        decisive=decisive,
+        requested_to={
+            ref: event.target_id for ref, event in requests.items() if event.target_id is not None
+        },
+    )