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

evalkit/db/llm_providers.py

@@ -0,0 +1,176 @@
+from datetime import UTC, datetime
+
+from pydantic import BaseModel
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from evalkit.contracts.settings import JudgeSettings
+from evalkit.db import models
+from evalkit.db.locks import lock_name
+
+
+class UnknownLlmProvider(KeyError):
+    def __init__(self, name: str):
+        super().__init__(
+            f"unknown LLM provider: {name} — create it with evalkit llm-providers create"
+        )
+
+
+class LlmProviderExists(ValueError):
+    pass
+
+
+class ProviderInUse(ValueError):
+    pass
+
+
+class LlmProviderView(BaseModel):
+    name: str
+    base_url: str
+    has_api_key: bool
+    archived: bool
+
+
+def _view(row: models.LlmProvider) -> LlmProviderView:
+    return LlmProviderView(
+        name=row.name,
+        base_url=row.base_url,
+        has_api_key=row.api_key_encrypted is not None,
+        archived=row.archived_at is not None,
+    )
+
+
+class LlmProviderStore:
+    def __init__(self, session_factory: async_sessionmaker[AsyncSession]):
+        self._sessions = session_factory
+
+    async def _row(self, session: AsyncSession, name: str) -> models.LlmProvider:
+        row = await session.scalar(
+            select(models.LlmProvider).where(models.LlmProvider.name == name)
+        )
+        if row is None:
+            raise UnknownLlmProvider(name)
+        return row
+
+    async def _settings(self, session: AsyncSession) -> models.ServerSettings | None:
+        return await session.get(models.ServerSettings, 1)
+
+    async def _check_free(self, session: AsyncSession, name: str, base_url: str) -> None:
+        taken = await session.scalar(
+            select(models.LlmProvider.name).where(
+                (models.LlmProvider.name == name) | (models.LlmProvider.base_url == base_url)
+            )
+        )
+        if taken is not None:
+            raise LlmProviderExists(f"LLM provider exists with that name or base URL: {taken}")
+
+    async def create(
+        self, name: str, base_url: str, api_key_encrypted: bytes | None
+    ) -> LlmProviderView:
+        async with self._sessions() as session:
+            await lock_name(session, "llm-provider", name)
+            await self._check_free(session, name, base_url)
+            row = models.LlmProvider(
+                name=name, base_url=base_url, api_key_encrypted=api_key_encrypted
+            )
+            session.add(row)
+            await session.commit()
+            return _view(row)
+
+    async def edit(
+        self,
+        name: str,
+        *,
+        base_url: str | None = None,
+        api_key_encrypted: bytes | None = None,
+        clear_api_key: bool = False,
+    ) -> LlmProviderView:
+        async with self._sessions() as session:
+            row = await self._row(session, name)
+            if base_url is not None and base_url != row.base_url:
+                settings = await self._settings(session)
+                if (
+                    settings is not None
+                    and (settings.default_judge or {}).get("base_url") == row.base_url
+                ):
+                    raise ProviderInUse(
+                        f"{name} is the default judge's provider; "
+                        "set another default judge before changing its base URL"
+                    )
+                row.base_url = base_url
+                row.api_key_encrypted = None
+            if api_key_encrypted is not None or clear_api_key:
+                row.api_key_encrypted = api_key_encrypted
+            await session.commit()
+            return _view(row)
+
+    async def archive(self, name: str) -> LlmProviderView:
+        async with self._sessions() as session:
+            row = await self._row(session, name)
+            settings = await self._settings(session)
+            if (
+                settings is not None
+                and (settings.default_judge or {}).get("base_url") == row.base_url
+            ):
+                raise ProviderInUse(
+                    f"{name} is the default judge's provider; set another default judge first"
+                )
+            row.archived_at = datetime.now(UTC)
+            await session.commit()
+            return _view(row)
+
+    async def restore(self, name: str) -> LlmProviderView:
+        async with self._sessions() as session:
+            row = await self._row(session, name)
+            row.archived_at = None
+            await session.commit()
+            return _view(row)
+
+    async def get(self, name: str) -> LlmProviderView:
+        async with self._sessions() as session:
+            return _view(await self._row(session, name))
+
+    async def list_providers(self, include_archived: bool = False) -> list[LlmProviderView]:
+        query = select(models.LlmProvider).order_by(models.LlmProvider.name)
+        if not include_archived:
+            query = query.where(models.LlmProvider.archived_at.is_(None))
+        async with self._sessions() as session:
+            return [_view(row) for row in (await session.scalars(query)).all()]
+
+    async def active_by_base_url(self, base_url: str) -> LlmProviderView | None:
+        async with self._sessions() as session:
+            row = await session.scalar(
+                select(models.LlmProvider).where(
+                    models.LlmProvider.base_url == base_url,
+                    models.LlmProvider.archived_at.is_(None),
+                )
+            )
+        return _view(row) if row is not None else None
+
+    async def sealed_key(self, base_url: str) -> bytes | None:
+        async with self._sessions() as session:
+            return await session.scalar(
+                select(models.LlmProvider.api_key_encrypted).where(
+                    models.LlmProvider.base_url == base_url,
+                    models.LlmProvider.archived_at.is_(None),
+                )
+            )
+
+    async def default_judge(self) -> JudgeSettings | None:
+        async with self._sessions() as session:
+            settings = await self._settings(session)
+        if settings is None or settings.default_judge is None:
+            return None
+        return JudgeSettings.model_validate(settings.default_judge)
+
+    async def set_default_judge(self, judge: JudgeSettings) -> JudgeSettings:
+        if await self.active_by_base_url(judge.base_url) is None:
+            raise UnknownLlmProvider(judge.base_url)
+        async with self._sessions() as session:
+            settings = await self._settings(session)
+            if settings is None:
+                settings = models.ServerSettings(id=1)
+                session.add(settings)
+            settings.default_judge = judge.model_dump()
+            await session.commit()
+        return judge

evalkit/db/locks.py

@@ -0,0 +1,6 @@
+from sqlalchemy import func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+
+async def lock_name(session: AsyncSession, namespace: str, name: str = "") -> None:
+    await session.execute(select(func.pg_advisory_xact_lock(func.hashtext(f"{namespace}:{name}"))))

evalkit/db/mappings.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from typing import Any
+
+from pydantic import BaseModel
+from sqlalchemy import Select, func, select
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from evalkit.contracts.enums import MappingStatus
+from evalkit.db import models
+
+
+class MappingRecord(BaseModel):
+    id: int
+    version: int
+    status: str
+    document: dict[str, Any]
+    created_by: str | None
+    created_at: datetime
+    activated_at: datetime | None
+
+
+def _to_record(row: models.SpanMapping, author: str | None) -> MappingRecord:
+    return MappingRecord(
+        id=row.id,
+        version=row.version,
+        status=row.status,
+        document=row.mapping,
+        created_by=author,
+        created_at=row.created_at,
+        activated_at=row.activated_at,
+    )
+
+
+def _with_author(query: Select) -> Select:
+    return query.add_columns(models.User.username).outerjoin(
+        models.User, models.User.id == models.SpanMapping.created_by_id
+    )
+
+
+async def active_mapping(session: AsyncSession, agent_id: int) -> models.SpanMapping | None:
+    return await session.scalar(
+        select(models.SpanMapping).where(
+            models.SpanMapping.agent_id == agent_id,
+            models.SpanMapping.status == MappingStatus.ACTIVE.value,
+        )
+    )
+
+
+async def retire_active(session: AsyncSession, agent_id: int, keep: int | None = None) -> None:
+    current = await active_mapping(session, agent_id)
+    if current is not None and current.id != keep:
+        current.status = MappingStatus.RETIRED.value
+        await session.flush()
+
+
+async def next_version(session: AsyncSession, agent_id: int) -> int:
+    last = await session.scalar(
+        select(func.max(models.SpanMapping.version)).where(models.SpanMapping.agent_id == agent_id)
+    )
+    return (last or 0) + 1
+
+
+async def activate_document(
+    session: AsyncSession, agent_id: int, document: dict[str, Any], created_by_id: int | None
+) -> models.SpanMapping:
+    await retire_active(session, agent_id)
+    row = models.SpanMapping(
+        agent_id=agent_id,
+        version=await next_version(session, agent_id),
+        status=MappingStatus.ACTIVE.value,
+        mapping=document,
+        created_by_id=created_by_id,
+        activated_at=datetime.now(UTC),
+    )
+    session.add(row)
+    await session.flush()
+    return row
+
+
+class MappingStore:
+    def __init__(self, session_factory: async_sessionmaker[AsyncSession]):
+        self._sessions = session_factory
+
+    async def list(self, agent_id: int) -> list[MappingRecord]:
+        async with self._sessions() as session:
+            rows = (
+                await session.execute(
+                    _with_author(
+                        select(models.SpanMapping)
+                        .where(models.SpanMapping.agent_id == agent_id)
+                        .order_by(models.SpanMapping.version)
+                    )
+                )
+            ).all()
+        return [_to_record(row, author) for row, author in rows]
+
+    async def _with_status(self, agent_id: int, status: str) -> MappingRecord | None:
+        async with self._sessions() as session:
+            found = (
+                await session.execute(
+                    _with_author(
+                        select(models.SpanMapping).where(
+                            models.SpanMapping.agent_id == agent_id,
+                            models.SpanMapping.status == status,
+                        )
+                    )
+                )
+            ).one_or_none()
+        return _to_record(found[0], found[1]) if found is not None else None
+
+    async def active(self, agent_id: int) -> MappingRecord | None:
+        return await self._with_status(agent_id, "active")
+
+    async def draft(self, agent_id: int) -> MappingRecord | None:
+        return await self._with_status(agent_id, "draft")
+
+    async def put_draft(
+        self, agent_id: int, document: dict[str, Any], created_by: str
+    ) -> MappingRecord:
+        async with self._sessions() as session:
+            current = (
+                await session.execute(
+                    select(models.SpanMapping).where(
+                        models.SpanMapping.agent_id == agent_id,
+                        models.SpanMapping.status == "draft",
+                    )
+                )
+            ).scalar_one_or_none()
+            if current is not None:
+                current.status = "retired"
+                await session.flush()
+            row = models.SpanMapping(
+                agent_id=agent_id,
+                version=await next_version(session, agent_id),
+                status="draft",
+                mapping=document,
+                created_by_id=await session.scalar(
+                    select(models.User.id).where(models.User.username == created_by)
+                ),
+            )
+            session.add(row)
+            await session.commit()
+            await session.refresh(row)
+            return _to_record(row, created_by)
+
+    async def activate(self, agent_id: int, version: int) -> MappingRecord:
+        async with self._sessions() as session:
+            row = (
+                await session.execute(
+                    select(models.SpanMapping).where(
+                        models.SpanMapping.agent_id == agent_id,
+                        models.SpanMapping.version == version,
+                    )
+                )
+            ).scalar_one_or_none()
+            if row is None:
+                raise KeyError(f"no mapping version {version}")
+            await retire_active(session, agent_id, keep=row.id)
+            row.status = "active"
+            row.activated_at = datetime.now(UTC)
+            await session.commit()
+            await session.refresh(row)
+            author = await session.scalar(
+                select(models.User.username).where(models.User.id == row.created_by_id)
+            )
+            return _to_record(row, author)