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

evalkit/api/trace_imports.py

@@ -0,0 +1,173 @@
+import asyncio
+import logging
+
+from fastapi import APIRouter, HTTPException, Request
+from pydantic import AwareDatetime, BaseModel, model_validator
+
+from evalkit.api.access import ActorDep
+from evalkit.api.deps import AgentStoreDep, HolderDep, MappingStoreDep, SessionDep
+from evalkit.api.experiments import task_holder
+from evalkit.api.traces import agent_or_404, require_otel, require_write
+from evalkit.contracts.enums import TraceImportStatus
+from evalkit.contracts.esf import MappingDocument
+from evalkit.core.agents import AgentDefinition
+from evalkit.core.runner import error_text
+from evalkit.core.traces.importer import IMPORT_STOPPED_ERROR, guarded_import
+from evalkit.core.traces.sources import client_for
+from evalkit.db.mappings import MappingStore
+from evalkit.db.trace_imports import TraceImportStore, TraceImportView
+from evalkit.db.trace_sources import ResolvedTraceSource
+
+router = APIRouter(prefix="/api/agents", tags=["traces"])
+LOGGER = logging.getLogger(__name__)
+
+ImportTasks = dict[int, tuple[str, asyncio.Task]]
+NO_SESSION_MAPPING = (
+    "activate a mapping with session_id before importing a time range; "
+    "import single sessions to try mappings first"
+)
+
+
+class RangeImportRequest(BaseModel):
+    source: str
+    project: str
+    since: AwareDatetime
+    until: AwareDatetime
+
+    @model_validator(mode="after")
+    def ordered(self) -> "RangeImportRequest":
+        if self.since >= self.until:
+            raise ValueError("since must be before until")
+        return self
+
+
+def _imports(request: Request) -> TraceImportStore:
+    return request.app.state.trace_imports
+
+
+async def _session_mapping(mappings: MappingStore, agent: AgentDefinition) -> MappingDocument:
+    active = await mappings.active(agent.id)
+    document = MappingDocument.model_validate(active.document) if active else None
+    if document is None or document.session_id is None:
+        raise HTTPException(status_code=409, detail=NO_SESSION_MAPPING)
+    return document
+
+
+def _log_failure(task: asyncio.Task) -> None:
+    if task.cancelled():
+        return
+    exc = task.exception()
+    if exc is not None:
+        LOGGER.error("trace import task %s failed: %s", task.get_name(), error_text(exc))
+
+
+def _launch(
+    request: Request,
+    import_id: int,
+    holder: str,
+    source: ResolvedTraceSource,
+    mapping: MappingDocument,
+) -> None:
+    client = client_for(
+        source.provider,
+        source.url,
+        source.api_key.get_secret_value() if source.api_key else None,
+        request.app.state.cfg.otel.import_page_traces,
+    )
+    store = _imports(request)
+    task = asyncio.create_task(
+        guarded_import(import_id, holder, client, store, request.app.state.spans, mapping),
+        name=f"import:{import_id}",
+    )
+    registry: ImportTasks = request.app.state.import_tasks
+    registry[import_id] = (holder, task)
+    task.add_done_callback(lambda t: registry.pop(import_id, None))
+    task.add_done_callback(_log_failure)
+
+
+async def interrupt_imports(store: TraceImportStore, registry: ImportTasks) -> None:
+    tracked = list(registry.items())
+    for _, (_, task) in tracked:
+        task.cancel()
+    await asyncio.gather(*(task for _, (_, task) in tracked), return_exceptions=True)
+    for import_id, (holder, task) in tracked:
+        if task.cancelled():
+            await store.finish(import_id, holder, TraceImportStatus.FAILED, IMPORT_STOPPED_ERROR)
+
+
+@router.post("/{agent}/imports", status_code=202)
+async def start_import(
+    agent: str,
+    body: RangeImportRequest,
+    request: Request,
+    agents: AgentStoreDep,
+    mappings: MappingStoreDep,
+    session: SessionDep,
+    actor: ActorDep,
+    pod: HolderDep,
+) -> TraceImportView:
+    _, record = await agent_or_404(agents, agent)
+    await require_write(session, actor, record)
+    require_otel(record)
+    mapping = await _session_mapping(mappings, record)
+    source = await request.app.state.trace_sources.resolve(body.source)
+    holder = task_holder(pod)
+    view = await _imports(request).create(
+        record.id, source.id, body.project, body.since, body.until, actor.id, holder
+    )
+    _launch(request, view.id, holder, source, mapping)
+    return view
+
+
+@router.get("/{agent}/imports")
+async def list_imports(
+    agent: str, request: Request, agents: AgentStoreDep
+) -> list[TraceImportView]:
+    _, record = await agent_or_404(agents, agent)
+    return await _imports(request).recent(record.id)
+
+
+@router.get("/{agent}/imports/{import_id}")
+async def get_import(
+    agent: str, import_id: int, request: Request, agents: AgentStoreDep
+) -> TraceImportView:
+    _, record = await agent_or_404(agents, agent)
+    return await _imports(request).get(record.id, import_id)
+
+
+@router.post("/{agent}/imports/{import_id}/stop")
+async def stop_import(
+    agent: str,
+    import_id: int,
+    request: Request,
+    agents: AgentStoreDep,
+    session: SessionDep,
+    actor: ActorDep,
+) -> TraceImportView:
+    _, record = await agent_or_404(agents, agent)
+    await require_write(session, actor, record)
+    return await _imports(request).stop(record.id, import_id)
+
+
+@router.post("/{agent}/imports/{import_id}/continue")
+async def continue_import(
+    agent: str,
+    import_id: int,
+    request: Request,
+    agents: AgentStoreDep,
+    mappings: MappingStoreDep,
+    session: SessionDep,
+    actor: ActorDep,
+    pod: HolderDep,
+) -> TraceImportView:
+    _, record = await agent_or_404(agents, agent)
+    await require_write(session, actor, record)
+    require_otel(record)
+    mapping = await _session_mapping(mappings, record)
+    store = _imports(request)
+    current = await store.get(record.id, import_id)
+    source = await request.app.state.trace_sources.resolve(current.source)
+    holder = task_holder(pod)
+    view = await store.restart(record.id, import_id, holder)
+    _launch(request, view.id, holder, source, mapping)
+    return view

evalkit/api/trace_sources.py

@@ -0,0 +1,49 @@
+from fastapi import APIRouter, Request
+from pydantic import BaseModel
+
+from evalkit.api.auth import AdminUser
+from evalkit.contracts.enums import TraceProvider
+from evalkit.db.trace_sources import TraceSourceStore
+
+router = APIRouter(prefix="/api/trace-sources", tags=["trace sources"])
+
+
+def _store(request: Request) -> TraceSourceStore:
+    return request.app.state.trace_sources
+
+
+@router.get("")
+async def list_trace_sources(request: Request, include_archived: bool = False) -> list[dict]:
+    rows = await _store(request).list_sources(include_archived=include_archived)
+    return [row.model_dump(mode="json") for row in rows]
+
+
+class CreateTraceSource(BaseModel):
+    name: str
+    provider: TraceProvider
+    url: str
+    api_key: str | None = None
+
+
+@router.post("")
+async def create_trace_source(body: CreateTraceSource, request: Request, admin: AdminUser) -> dict:
+    created = await _store(request).create(body.name, body.provider, body.url, body.api_key)
+    return created.model_dump(mode="json")
+
+
+class EditTraceSource(BaseModel):
+    url: str | None = None
+    api_key: str | None = None
+
+
+@router.patch("/{name}")
+async def edit_trace_source(
+    name: str, body: EditTraceSource, request: Request, admin: AdminUser
+) -> dict:
+    edited = await _store(request).edit(name, url=body.url, api_key=body.api_key)
+    return edited.model_dump(mode="json")
+
+
+@router.post("/{name}/archive")
+async def archive_trace_source(name: str, request: Request, admin: AdminUser) -> dict:
+    return (await _store(request).archive(name)).model_dump(mode="json")

evalkit/api/traces.py

@@ -0,0 +1,196 @@
+import logging
+from collections.abc import AsyncIterator
+from typing import Literal
+
+import httpx
+from fastapi import APIRouter, HTTPException, Request
+from fastapi.responses import JSONResponse, Response, StreamingResponse
+from pydantic import BaseModel, ValidationError
+from pydantic_core import to_json
+
+from evalkit.api.access import ActorDep, ensure_active, require, workspace_ref
+from evalkit.api.deps import (
+    AgentStoreDep,
+    MappingStoreDep,
+    SessionDep,
+    SessionReaderDep,
+    SpanStoreDep,
+)
+from evalkit.api.experiments import slug_or_400
+from evalkit.contracts.esf import MappingDocument
+from evalkit.core.access import can_write
+from evalkit.core.agents import AgentDefinition, UnknownAgent, explicit
+from evalkit.core.traces.mapping import session_id_of
+from evalkit.core.traces.raw import RawSpan
+from evalkit.core.traces.sessions import SessionTrajectory
+from evalkit.core.traces.sources import client_for
+from evalkit.core.workspace import Reference
+from evalkit.db.agents import AgentStore
+
+router = APIRouter(prefix="/api", tags=["traces"])
+LOGGER = logging.getLogger(__name__)
+
+
+async def session_too_large(request: Request, exc: Exception) -> JSONResponse:
+    LOGGER.warning("refused session read %s: %s", request.url.path, exc)
+    return JSONResponse(status_code=413, content={"detail": str(exc)})
+
+
+async def agent_or_404(agents: AgentStore, agent: str) -> tuple[Reference, AgentDefinition]:
+    ref = slug_or_400(agent)
+    return ref, await agents.get(explicit(ref))
+
+
+async def require_write(session: SessionDep, actor: ActorDep, agent: AgentDefinition) -> None:
+    ws = await workspace_ref(session, agent.workspace)
+    require(can_write(actor, ws), f"you cannot change {agent.ref}'s mappings or traces")
+    ensure_active(ws)
+
+
+def require_otel(agent: AgentDefinition) -> None:
+    if not agent.is_otel:
+        raise HTTPException(
+            status_code=400,
+            detail=f"{agent.ref} takes ATIF trajectories; only OTel agents import traces",
+        )
+
+
+@router.get("/agents/{agent}/sessions/{session_id}/trajectory")
+async def session_trajectory(
+    agent: str,
+    session_id: str,
+    sessions: SessionReaderDep,
+    mappings: MappingStoreDep,
+    mapping: Literal["active", "draft"] = "active",
+    esf: bool = False,
+) -> Response:
+    if mapping == "draft":
+        record = await sessions.agent(agent)
+        if await mappings.draft(record.id) is None:
+            raise HTTPException(
+                status_code=404,
+                detail=f"no draft mapping for {agent}; push one with evalkit traces mapping push",
+            )
+    result = await sessions.read(agent, session_id, mapping=mapping, include_esf=esf)
+    if result is None:
+        raise UnknownAgent(agent)
+    if result.esf is None:
+        return Response(content=to_json(result), media_type="application/json")
+    return StreamingResponse(_with_esf(result), media_type="application/json")
+
+
+async def _with_esf(result: SessionTrajectory) -> AsyncIterator[bytes]:
+    esf = result.esf or []
+    yield to_json(result, exclude={"esf"})[:-1] + b',"esf":['
+    for i, span in enumerate(esf):
+        yield (b"," if i else b"") + to_json(span)
+    yield b"]}"
+
+
+async def _json_array(spans: AsyncIterator[RawSpan]) -> AsyncIterator[bytes]:
+    yield b"["
+    first = True
+    async for span in spans:
+        yield (b"" if first else b",") + to_json(span)
+        first = False
+    yield b"]"
+
+
+@router.get("/agents/{agent}/sessions/{session_id}/spans")
+async def session_spans(
+    agent: str, session_id: str, sessions: SessionReaderDep
+) -> StreamingResponse:
+    record = await sessions.agent(agent)
+    if record is None:
+        raise UnknownAgent(agent)
+    return StreamingResponse(
+        _json_array(sessions.raw(record, session_id)), media_type="application/json"
+    )
+
+
+@router.get("/agents/{agent}/mappings")
+async def list_mappings(agent: str, agents: AgentStoreDep, mappings: MappingStoreDep) -> list[dict]:
+    _, record = await agent_or_404(agents, agent)
+    return [m.model_dump(mode="json") for m in await mappings.list(record.id)]
+
+
+@router.put("/agents/{agent}/mappings/draft")
+async def put_draft(
+    agent: str,
+    body: dict,
+    agents: AgentStoreDep,
+    mappings: MappingStoreDep,
+    session: SessionDep,
+    actor: ActorDep,
+    activate: bool = False,
+) -> dict:
+    _, record = await agent_or_404(agents, agent)
+    await require_write(session, actor, record)
+    try:
+        document = MappingDocument.model_validate(body)
+    except ValidationError as exc:
+        first = exc.errors()[0]
+        location = ".".join(str(p) for p in first["loc"])
+        raise HTTPException(status_code=400, detail=f"{location}: {first['msg']}") from exc
+    draft = await mappings.put_draft(
+        record.id, document.model_dump(by_alias=True, exclude_none=True), actor.username
+    )
+    if activate:
+        draft = await mappings.activate(record.id, draft.version)
+    return draft.model_dump(mode="json")
+
+
+@router.post("/agents/{agent}/mappings/{version}/activate")
+async def activate_mapping(
+    agent: str,
+    version: int,
+    agents: AgentStoreDep,
+    mappings: MappingStoreDep,
+    session: SessionDep,
+    actor: ActorDep,
+) -> dict:
+    _, record = await agent_or_404(agents, agent)
+    await require_write(session, actor, record)
+    try:
+        return (await mappings.activate(record.id, version)).model_dump(mode="json")
+    except KeyError as exc:
+        raise HTTPException(status_code=404, detail=str(exc)) from exc
+
+
+class ImportRequest(BaseModel):
+    source: str
+    project: str
+    session_id: str
+
+
+@router.post("/agents/{agent}/sessions/import")
+async def import_session(
+    agent: str,
+    body: ImportRequest,
+    request: Request,
+    agents: AgentStoreDep,
+    spans: SpanStoreDep,
+    mappings: MappingStoreDep,
+    session: SessionDep,
+    actor: ActorDep,
+) -> dict:
+    _, record = await agent_or_404(agents, agent)
+    await require_write(session, actor, record)
+    require_otel(record)
+    source = await request.app.state.trace_sources.resolve(body.source)
+    client = client_for(
+        source.provider,
+        source.url,
+        source.api_key.get_secret_value() if source.api_key else None,
+        request.app.state.cfg.otel.import_page_traces,
+    )
+    try:
+        fetched = await client.session_spans(body.project, body.session_id)
+    except httpx.HTTPError as exc:
+        raise HTTPException(status_code=502, detail=f"{body.source}: {exc}") from exc
+    active = await mappings.active(record.id)
+    document = MappingDocument.model_validate(active.document) if active else None
+    for span in fetched:
+        span.session_id_hint = session_id_of(span.attributes, document) or body.session_id
+    imported = await spans.upsert(record.id, fetched, external_trace_source_id=source.id)
+    return {"imported": imported}

evalkit/api/users.py

@@ -0,0 +1,72 @@
+from fastapi import APIRouter, Request
+from pydantic import BaseModel
+
+from evalkit.api.auth import AdminUser
+from evalkit.db.users import UserStore
+
+router = APIRouter(prefix="/api/users", tags=["users"])
+
+
+def _users(request: Request) -> UserStore:
+    return UserStore(request.app.state.session_factory)
+
+
+@router.get("")
+async def list_users(request: Request, admin: AdminUser) -> list[dict]:
+    return [
+        {
+            "id": u.id,
+            "username": u.username,
+            "is_admin": u.is_admin,
+            "archived": u.archived_at is not None,
+            "created_at": u.created_at.isoformat(),
+        }
+        for u in await _users(request).list()
+    ]
+
+
+class CreateUserRequest(BaseModel):
+    username: str
+    password: str
+
+
+@router.post("")
+async def create_user(body: CreateUserRequest, request: Request, admin: AdminUser) -> dict:
+    await _users(request).create(body.username, body.password)
+    return {"username": body.username}
+
+
+@router.post("/{username}/archive")
+async def archive_user(username: str, request: Request, admin: AdminUser) -> dict:
+    await _users(request).archive(username)
+    return {"archived": username}
+
+
+@router.post("/{username}/restore")
+async def restore_user(username: str, request: Request, admin: AdminUser) -> dict:
+    await _users(request).restore(username)
+    return {"restored": username}
+
+
+class SetPasswordRequest(BaseModel):
+    password: str
+
+
+@router.post("/{username}/password")
+async def set_password(
+    username: str, body: SetPasswordRequest, request: Request, admin: AdminUser
+) -> dict:
+    await _users(request).set_password(username, body.password)
+    return {"password_set": username}
+
+
+@router.post("/{username}/promote")
+async def promote_user(username: str, request: Request, admin: AdminUser) -> dict:
+    await _users(request).set_admin(username, True)
+    return {"promoted": username}
+
+
+@router.post("/{username}/demote")
+async def demote_user(username: str, request: Request, admin: AdminUser) -> dict:
+    await _users(request).set_admin(username, False)
+    return {"demoted": username}

evalkit/api/workspaces.py

@@ -0,0 +1,131 @@
+from datetime import UTC, datetime
+
+from fastapi import APIRouter, HTTPException, Request
+from pydantic import BaseModel
+
+from evalkit.api.access import ActorDep, require
+from evalkit.api.auth import AdminUser
+from evalkit.contracts.enums import WorkspaceRole
+from evalkit.core.access import Actor, WorkspaceRef
+from evalkit.db.workspaces import WorkspaceStore
+
+router = APIRouter(prefix="/api/workspaces", tags=["workspaces"])
+
+
+def _store(request: Request) -> WorkspaceStore:
+    return WorkspaceStore(request.app.state.session_factory)
+
+
+def _view(ws: WorkspaceRef, names: dict[int, str]) -> dict:
+    return {
+        "id": ws.id,
+        "name": ws.name,
+        "type": "shared" if ws.is_shared else "personal",
+        "owner": ws.name if not ws.is_shared else None,
+        "archived": ws.archived,
+        "promotes_to": names.get(ws.promotes_to) if ws.promotes_to is not None else None,
+        "promoted_from": names.get(ws.promoted_from) if ws.promoted_from is not None else None,
+    }
+
+
+async def _one(store: WorkspaceStore, ws: WorkspaceRef) -> dict:
+    linked = [i for i in (ws.promotes_to, ws.promoted_from) if i is not None]
+    names = {i: (await store.get_by_id(i)).name for i in linked}
+    return _view(ws, names)
+
+
+@router.get("")
+async def list_workspaces(request: Request, include_archived: bool = False) -> list[dict]:
+    rows = await _store(request).list(include_archived=True)
+    names = {ws.id: ws.name for ws in rows}
+    return [_view(ws, names) for ws in rows if include_archived or not ws.archived]
+
+
+class CreateWorkspace(BaseModel):
+    name: str
+
+
+@router.post("")
+async def create_workspace(body: CreateWorkspace, request: Request, admin: AdminUser) -> dict:
+    store = _store(request)
+    return await _one(store, await store.create_shared(body.name, admin.id))
+
+
+@router.get("/{ws}")
+async def get_workspace(ws: str, request: Request) -> dict:
+    store = _store(request)
+    return await _one(store, await store.get_by_name(ws))
+
+
+class UpdateWorkspace(BaseModel):
+    name: str | None = None
+    promotes_to: str | None = None
+
+
+@router.patch("/{ws}")
+async def update_workspace(
+    ws: str, body: UpdateWorkspace, request: Request, admin: AdminUser
+) -> dict:
+    store = _store(request)
+    update: list[str] = []
+    if "promotes_to" in body.model_fields_set:
+        await store.set_promotes_to(ws, body.promotes_to)
+        if body.promotes_to is not None:
+            update.append(
+                f"CI that syncs into @{body.promotes_to} must sync into @{ws} instead; "
+                f"evalkit sync --workspace {body.promotes_to} now fails"
+            )
+    if body.name is not None:
+        await store.rename(ws, body.name)
+        update += [
+            f"developer evalkit.yaml credentials keys ending in @{ws}",
+            f"CI and scripts that pass --workspace {ws} or refer to name@{ws}",
+            f"bookmarked UI links that contain @{ws}",
+        ]
+    current = await store.get_by_name(body.name or ws)
+    return {**(await _one(store, current)), "update": update}
+
+
+def _may_archive(actor: Actor, ws: WorkspaceRef) -> bool:
+    return actor.is_admin if ws.is_shared else ws.owner_user_id == actor.id
+
+
+@router.post("/{ws}/archive")
+async def archive_workspace(ws: str, request: Request, actor: ActorDep) -> dict:
+    store = _store(request)
+    ref = await store.get_by_name(ws)
+    require(_may_archive(actor, ref), "only admins archive shared workspaces, owners their own")
+    return await _one(store, await store.set_archived(ws, datetime.now(UTC)))
+
+
+@router.post("/{ws}/restore")
+async def restore_workspace(ws: str, request: Request, actor: ActorDep) -> dict:
+    store = _store(request)
+    ref = await store.get_by_name(ws)
+    require(_may_archive(actor, ref), "only admins restore shared workspaces, owners their own")
+    return await _one(store, await store.set_archived(ws, None))
+
+
+@router.get("/{ws}/members")
+async def list_members(ws: str, request: Request, role: WorkspaceRole | None = None) -> list[dict]:
+    return [m.model_dump(mode="json") for m in await _store(request).members(ws, role)]
+
+
+class AddMember(BaseModel):
+    username: str
+    role: WorkspaceRole
+
+
+@router.post("/{ws}/members")
+async def add_member(ws: str, body: AddMember, request: Request, admin: AdminUser) -> dict:
+    await _store(request).add_member(ws, body.username, body.role, admin.id)
+    return {"workspace": ws, "username": body.username, "role": body.role}
+
+
+@router.delete("/{ws}/members/{username}/{role}")
+async def remove_member(
+    ws: str, username: str, role: WorkspaceRole, request: Request, admin: AdminUser
+) -> dict:
+    if not await _store(request).remove_member(ws, username, role):
+        raise HTTPException(status_code=404, detail=f"{username} is not a {role} of {ws}")
+    return {"removed": username, "role": role}

evalkit/bundled/contracts/atif-accepted.json

@@ -0,0 +1,5 @@
+{
+  "accepted_schema_versions": [
+    "ATIF-v1.8"
+  ]
+}