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

evalkit/sdk/tracing.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+import os
+import threading
+from collections import OrderedDict
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from importlib.util import find_spec
+from typing import Any
+
+from opentelemetry import trace
+from opentelemetry.context import Context
+
+from evalkit.client import connection
+
+OTEL_TOKEN_ENV = "EVALKIT_OTEL_TOKEN"
+SERVER_URL_ENV = "EVALKIT_SERVER_URL"
+AGENT_HEADER = "X-Evalkit-Agent"
+STAMPED = ("session.id", "evalkit.experiment_id", "evalkit.case_id", "evalkit.repetition")
+EVAL_MARKER = "evalkit.case_id"
+MAX_EVAL_TRACES = 10_000
+INSTALL_HINT = "evalkit tracing needs the otel extra: uv add 'evalkit-py[otel]'"
+
+CURRENT_RUN: ContextVar[dict[str, Any] | None] = ContextVar("evalkit_current_run", default=None)
+
+SpanFilter = Callable[[Any], bool]
+
+HAS_OTEL_SDK = bool(
+    find_spec("opentelemetry.sdk") and find_spec("opentelemetry.exporter.otlp.proto.http")
+)
+
+if HAS_OTEL_SDK:
+    from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+    from opentelemetry.sdk.resources import Resource
+    from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor, TracerProvider
+    from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter
+
+    class EvalkitSpanProcessor(SpanProcessor):
+        def __init__(self, exporter: SpanExporter, only: SpanFilter | None = None) -> None:
+            self.batch = BatchSpanProcessor(exporter)
+            self.only = only
+            self.max_traces = MAX_EVAL_TRACES
+            self.kept_traces: OrderedDict[int, None] = OrderedDict()
+            self.lock = threading.Lock()
+
+        def on_start(self, span: Any, parent_context: Context | None = None) -> None:
+            attributes = CURRENT_RUN.get()
+            if attributes:
+                for key in STAMPED:
+                    if key in attributes:
+                        span.set_attribute(key, attributes[key])
+            if self.only is not None and self.only(span):
+                self.keep_trace(trace_of(span))
+
+        def on_end(self, span: ReadableSpan) -> None:
+            if self.only is None:
+                self.batch.on_end(span)
+                return
+            trace_id = trace_of(span)
+            if self.only(span):
+                self.keep_trace(trace_id)
+            if self.is_kept(trace_id):
+                self.batch.on_end(span)
+
+        def is_kept(self, trace_id: int) -> bool:
+            with self.lock:
+                if trace_id not in self.kept_traces:
+                    return False
+                self.kept_traces.move_to_end(trace_id)
+                return True
+
+        def keep_trace(self, trace_id: int) -> None:
+            with self.lock:
+                self.kept_traces[trace_id] = None
+                self.kept_traces.move_to_end(trace_id)
+                while len(self.kept_traces) > self.max_traces:
+                    self.kept_traces.popitem(last=False)
+
+        def shutdown(self) -> None:
+            self.batch.shutdown()
+
+        def force_flush(self, timeout_millis: int = 30000) -> bool:
+            return self.batch.force_flush(timeout_millis)
+
+    def configure(
+        provider: TracerProvider, exporter: SpanExporter, only: SpanFilter | None = None
+    ) -> None:
+        provider.add_span_processor(EvalkitSpanProcessor(exporter, only))
+
+
+def trace_of(span: Any) -> int:
+    return span.get_span_context().trace_id
+
+
+def collector_url(server_url: str) -> str:
+    return f"{server_url.rstrip('/')}/api/otel/v1/traces"
+
+
+def resolve_destination() -> tuple[str, str | None]:
+    token = os.environ.get(OTEL_TOKEN_ENV)
+    url = os.environ.get(SERVER_URL_ENV)
+    if url and token:
+        return url, token
+    active = connection.active()
+    return url or active.url, token or active.token
+
+
+def headers(agent: str) -> dict[str, str]:
+    _, token = resolve_destination()
+    out = {AGENT_HEADER: agent}
+    if token:
+        out["Authorization"] = f"Bearer {token}"
+    return out
+
+
+def exporter(agent: str) -> OTLPSpanExporter:
+    if not HAS_OTEL_SDK:
+        raise RuntimeError(INSTALL_HINT)
+    url, _ = resolve_destination()
+    return OTLPSpanExporter(endpoint=collector_url(url), headers=headers(agent))
+
+
+def processor(agent: str, only: SpanFilter | None = None) -> EvalkitSpanProcessor:
+    return EvalkitSpanProcessor(exporter(agent), only)
+
+
+def is_eval_span(span: Any) -> bool:
+    return EVAL_MARKER in (span.attributes or {})
+
+
+@contextmanager
+def eval_run(metadata: dict[str, Any]) -> Iterator[None]:
+    attributes = {
+        "session.id": metadata.get("session_id"),
+        "evalkit.experiment_id": metadata.get("experiment_id"),
+        "evalkit.case_id": metadata.get("case_id"),
+        "evalkit.repetition": metadata.get("repetition"),
+    }
+    token = CURRENT_RUN.set({k: v for k, v in attributes.items() if v is not None})
+    try:
+        yield
+    finally:
+        CURRENT_RUN.reset(token)
+
+
+def register(agent: str, only: SpanFilter | None = None) -> None:
+    """Send this process's OpenTelemetry spans to the evalkit server.
+
+    Adds an exporter to the global tracer provider, or creates one when none is set.
+    Call it after your own tracing setup, so it joins your provider. Spans go to the
+    server you logged in to (or `EVALKIT_SERVER_URL`), with your token or
+    `EVALKIT_OTEL_TOKEN` (an agent token). Needs the `otel` extra.
+
+    Args:
+        agent: The agent the spans belong to, e.g. `myagent` or `myagent@dev`.
+        only: Export only the spans this function returns `True` for.
+    """
+    span_processor = processor(agent, only)
+    current = trace.get_tracer_provider()
+    if isinstance(current, TracerProvider):
+        provider = current
+    else:
+        provider = TracerProvider(resource=Resource.create({"evalkit.agent": agent}))
+        trace.set_tracer_provider(provider)
+    provider.add_span_processor(span_processor)

evalkit/store/base.py

@@ -0,0 +1,354 @@
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Literal, Protocol
+
+from pydantic import BaseModel, Field
+
+from evalkit.contracts.enums import (
+    EvaluationSource,
+    EvaluatorKind,
+    ExperimentStatus,
+    RunStatus,
+    TrajectorySource,
+    TrajectoryStatus,
+)
+from evalkit.contracts.settings import JudgeSettings, RemoteEvaluatorSettings
+
+
+@dataclass(frozen=True)
+class BlobStream:
+    size: int
+    chunks: AsyncIterator[bytes]
+
+
+class UnknownDataset(KeyError):
+    def __init__(self, ref: str):
+        super().__init__(ref)
+        self.ref = ref
+
+    def __str__(self) -> str:
+        return f"unknown dataset: {self.ref}"
+
+
+class CaseFile(BaseModel):
+    key: str
+    sha256: str
+    context: str
+
+
+class FixtureUse(BaseModel):
+    key: str
+    sha256: str
+    size: int
+    datasets: list[str]
+    uploaded_by: str | None
+
+
+class FixtureUsePage(BaseModel):
+    items: list[FixtureUse]
+    total: int
+
+
+class CaseView(BaseModel):
+    case_ref: int
+    case_id: str
+    content_hash: str
+    input: dict[str, Any]
+    expected: dict[str, Any] | None = None
+    evaluators: dict[str, Any] = Field(default_factory=dict)
+    effective_evaluators: dict[str, Any] = Field(default_factory=dict)
+    tags: list[str] = Field(default_factory=list)
+    split_tags: list[str] = Field(default_factory=list)
+    source: str | None = None
+    files: list[CaseFile] = Field(default_factory=list)
+    review_state: str | None = None
+
+
+class DatasetSyncView(BaseModel):
+    id: int
+    source: str
+    detail: dict[str, Any]
+    actor: str
+    created_at: datetime
+
+
+class DatasetRow(BaseModel):
+    id: int
+    workspace: str
+    workspace_is_shared: bool
+    name: str
+    kind: str
+    last_sync: DatasetSyncView | None = None
+
+    @property
+    def ref(self) -> str:
+        return f"{self.name}@{self.workspace}"
+
+
+class Selection(BaseModel):
+    all: bool = False
+    tag: str | None = None
+    case_ids: list[str] | None = None
+
+
+class IncomingCase(BaseModel):
+    case_id: str
+    content_hash: str
+    input: dict[str, Any]
+    expected: dict[str, Any] | None = None
+    evaluators: dict[str, Any] = Field(default_factory=dict)
+    tags: list[str] = Field(default_factory=list)
+    path_tags: list[str] = Field(default_factory=list)
+    source: str | None = None
+    files: list[CaseFile] = Field(default_factory=list)
+
+
+class IncomingOverlay(BaseModel):
+    target_case_id: str
+    agents: list[str]
+    tags: list[str]
+    evaluators: dict[str, Any]
+    source_file: str | None = None
+
+
+class CarriedReview(BaseModel):
+    case_id: str
+    reviewer: str
+
+
+class SyncResult(BaseModel):
+    sync_id: int
+    dataset_id: int
+    created: list[str] = Field(default_factory=list)
+    patched: list[str] = Field(default_factory=list)
+    removed: list[str] = Field(default_factory=list)
+    unchanged: int = 0
+    overlays: int = 0
+    carried: list[CarriedReview] = Field(default_factory=list)
+    now_unreviewed: list[str] = Field(default_factory=list)
+
+
+class GitRevision(BaseModel):
+    sha: str | None = None
+    branch: str | None = None
+    dirty: bool | None = None
+
+
+class ExperimentSnapshot(BaseModel):
+    capabilities: dict[str, Any]
+    trajectory_source: TrajectorySource
+    trajectory_wait_seconds: int
+    remote_evaluators: dict[str, RemoteEvaluatorSettings] = Field(default_factory=dict)
+    code_evaluators: list[str] = Field(default_factory=list)
+    evalkit_version: str | None = None
+
+
+class ExperimentConfig(BaseModel):
+    model: str | None = None
+    judge: JudgeSettings
+    selection: Selection
+    revision: GitRevision | None = None
+    snapshot: ExperimentSnapshot
+
+
+class Experiment(BaseModel):
+    id: int
+    number: int
+    workspace: str
+    dataset: str
+    agent: str
+    agent_id: int
+    created_by: str
+    status: ExperimentStatus
+    status_error: str | None
+    lease_holder: str | None
+    lease_expires_at: datetime | None
+    repetitions: int
+    config: ExperimentConfig
+    started_at: datetime
+    ended_at: datetime | None
+
+
+class PinnedVersion(BaseModel):
+    case_ref: int
+    content_hash: str
+
+
+class NewExperiment(BaseModel):
+    dataset: str
+    agent: str
+    runner: str
+    repetitions: int
+    config: ExperimentConfig
+    cases: list[PinnedVersion]
+    holder: str
+
+
+class AgentRunRecord(BaseModel):
+    harness_run_id: str
+    status: RunStatus
+    answer: str | None = None
+    output: dict[str, Any] | None = None
+    error: str | None = None
+    started_at: datetime
+    ended_at: datetime | None = None
+    trajectory_key: str | None = None
+    trajectory_status: TrajectoryStatus | None = None
+    metrics: dict[str, Any] = Field(default_factory=dict)
+    title: str | None = None
+
+
+class RunRef(BaseModel):
+    experiment_run_id: int
+    agent_run_id: int
+
+
+class CaseRep(BaseModel):
+    case_id: str
+    repetition: int
+
+
+class EvaluationRow(BaseModel):
+    producer: str
+    name: str
+    kind: EvaluatorKind
+    source: EvaluationSource
+    value: float | None = None
+    label: str | None = None
+    explanation: str | None = None
+    error: str | None = None
+
+
+class RunEvalContext(BaseModel):
+    experiment_id: int
+    repetition: int
+    case: CaseView
+    evaluations: list[EvaluationRow]
+
+
+class RunForScoring(BaseModel):
+    experiment_run_id: int
+    case: CaseView
+    repetition: int
+    agent_run: AgentRunRecord
+    agent_run_id: int
+    done: set[str]
+
+
+class RunDetail(BaseModel):
+    experiment_run_id: int
+    case_id: str
+    repetition: int
+    agent_run_id: int
+    agent_run: AgentRunRecord
+    evaluations: list[EvaluationRow] = Field(default_factory=list)
+
+
+class AggregateRows(BaseModel):
+    experiment: Experiment
+    case_scores: list[tuple[str, str, float, str | None, str | None]]
+    case_gauges: list[tuple[str, str, float]]
+    case_errors: list[tuple[str, str, str]]
+    run_errors: list[tuple[str, str]]
+    repetitions_seen: dict[str, int]
+    sources: dict[str, str]
+
+
+RunState = Literal["pending", "ran", "evaluated"]
+
+
+class Progress(BaseModel):
+    total: int
+    ran: int
+    evaluated: int
+
+
+class HistoryPoint(BaseModel):
+    experiment_id: int
+    number: int
+    agent: str
+    model: str | None = None
+    created_by: str
+    created_at: datetime
+    repetitions: int
+    scores: dict[str, float | None]
+    cost_usd: float | None = None
+
+
+class UnknownExperiment(KeyError):
+    def __init__(self, experiment_id: int):
+        super().__init__(f"unknown experiment: {experiment_id}")
+
+
+class Conflict(Exception):
+    pass
+
+
+class InvalidRun(ValueError):
+    pass
+
+
+class ScoreRow(BaseModel):
+    name: str
+    kind: Literal["evaluator", "gauge"] = "evaluator"
+    score: float | None = None
+    label: str | None = None
+    explanation: str | None = None
+    source: str = "builtin"
+    error: str | None = None
+    annotator_kind: Literal["CODE", "LLM", "HUMAN"] = "CODE"
+
+
+class Store(Protocol):
+    async def cases_for(
+        self, dataset: str, selection: Selection, *, agent: str | None
+    ) -> list[CaseView]: ...
+
+    async def get_experiment(self, exp_id: int) -> Experiment: ...
+
+    async def pinned_cases(self, exp_id: int) -> list[CaseView]: ...
+
+    async def transition(
+        self, exp_id: int, to: ExperimentStatus, holder: str | None, error: str | None = None
+    ) -> Experiment: ...
+
+    async def renew_lease(self, exp_id: int, holder: str) -> bool: ...
+
+    async def set_model(self, exp_id: int, model: str) -> None: ...
+
+    async def pending_runs(self, exp_id: int) -> list[CaseRep]: ...
+
+    async def harness_run_id_used(self, exp_id: int, harness_run_id: str) -> str | None: ...
+
+    async def record_run(
+        self,
+        exp_id: int,
+        case_id: str,
+        repetition: int,
+        run: AgentRunRecord,
+        *,
+        holder: str | None = None,
+    ) -> RunRef: ...
+
+    async def runs_to_score(self, exp_id: int, *, force: bool) -> list[RunForScoring]: ...
+
+    async def publish_evaluations(
+        self,
+        experiment_run_id: int,
+        rows: list[EvaluationRow],
+        *,
+        force: bool,
+        holder: str | None = None,
+    ) -> None: ...
+
+    async def aclose(self) -> None: ...
+
+
+def select_cases(cases: list[CaseView], selection: Selection) -> list[CaseView]:
+    if selection.all:
+        return list(cases)
+    if selection.tag:
+        return [case for case in cases if selection.tag in case.split_tags]
+    wanted = set(selection.case_ids or [])
+    return [case for case in cases if case.case_id in wanted]

evalkit/store/blob.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+
+import obstore
+from obstore.exceptions import NotFoundError
+from obstore.store import ObjectStore, from_url
+
+from evalkit.store.base import BlobStream
+
+CHUNK_BYTES = 1024 * 1024
+
+_stores: dict[str, ObjectStore] = {}
+
+
+def _store_for(url: str) -> ObjectStore:
+    if url not in _stores:
+        _stores[url] = from_url(url, mkdir=True) if url.startswith("file://") else from_url(url)
+    return _stores[url]
+
+
+class Obstore:
+    def __init__(self, url: str):
+        self.url = url
+
+    @property
+    def _store(self) -> ObjectStore:
+        return _store_for(self.url)
+
+    async def put_object(self, key: str, data: bytes) -> None:
+        await obstore.put_async(self._store, key, data)
+
+    async def put_stream(self, key: str, chunks: AsyncIterator[bytes]) -> None:
+        await obstore.put_async(self._store, key, chunks)
+
+    async def get_object(self, key: str) -> bytes:
+        resp = await obstore.get_async(self._store, key)
+        return bytes(await resp.bytes_async())
+
+    async def get_stream(self, key: str) -> BlobStream:
+        resp = await obstore.get_async(self._store, key)
+        return BlobStream(resp.meta["size"], _chunks(resp.stream(min_chunk_size=CHUNK_BYTES)))
+
+    async def object_exists(self, key: str) -> bool:
+        try:
+            await obstore.head_async(self._store, key)
+            return True
+        except (NotFoundError, FileNotFoundError):
+            return False
+
+    async def list_objects(self, prefix: str) -> list[tuple[str, int]]:
+        objects: list[tuple[str, int]] = []
+        async for chunk in obstore.list(self._store, prefix=prefix):
+            objects.extend((obj["path"], obj["size"]) for obj in chunk)
+        return objects
+
+    async def delete_object(self, key: str) -> None:
+        await obstore.delete_async(self._store, key)
+
+    async def copy_object(self, source: str, destination: str) -> None:
+        await obstore.copy_async(self._store, source, destination)
+
+    async def object_url(self, key: str) -> str:
+        return f"{self.url.rstrip('/')}/{key}"
+
+
+async def _chunks(stream: AsyncIterator) -> AsyncIterator[bytes]:
+    async for chunk in stream:
+        yield bytes(chunk)

evalkit/version.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+import importlib.metadata
+import re
+from enum import StrEnum
+from typing import NamedTuple
+
+
+class ParsedVersion(NamedTuple):
+    major: int
+    minor: int
+    patch: int
+    build_id: str | None
+
+
+class VersionComparison(StrEnum):
+    SAME = "same"
+    DIFFERENT = "different"
+    INCOMPATIBLE = "incompatible"
+
+
+DISTRIBUTION = "evalkit-py"
+_VERSION = re.compile(r"(\d+)\.(\d+)\.(\d+)(?:\.(dev\d+))?(?:\+(.+))?")
+
+
+def current() -> str:
+    return importlib.metadata.version(DISTRIBUTION)
+
+
+def parse(version: str) -> ParsedVersion:
+    match = _VERSION.fullmatch(version)
+    if match is None:
+        raise ValueError(f"not an evalkit version: {version}")
+    major, minor, patch, dev, local = match.groups()
+    return ParsedVersion(int(major), int(minor), int(patch), dev or local)
+
+
+def compare(local: str, remote: str) -> VersionComparison:
+    mine, theirs = parse(local), parse(remote)
+    if mine.major != theirs.major or (mine.major == 0 and mine.minor != theirs.minor):
+        return VersionComparison.INCOMPATIBLE
+    if local == remote:
+        return VersionComparison.SAME
+    return VersionComparison.DIFFERENT