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

evalkit/core/traces/sources/phoenix.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from datetime import datetime
+from typing import Any
+
+import httpx
+
+from evalkit.core.traces.raw import RawSpan
+from evalkit.core.traces.sources.base import TracePage
+
+
+def flatten_attributes(nested: dict[str, Any], prefix: str = "") -> dict[str, Any]:
+    flat: dict[str, Any] = {}
+    for key, value in nested.items():
+        name = f"{prefix}{key}"
+        if isinstance(value, dict):
+            flat.update(flatten_attributes(value, f"{name}."))
+        elif isinstance(value, list) and value and all(isinstance(v, dict) for v in value):
+            for i, item in enumerate(value):
+                flat.update(flatten_attributes(item, f"{name}.{i}."))
+        else:
+            flat[name] = value
+    return flat
+
+
+def _to_raw(span: dict[str, Any]) -> RawSpan:
+    context = span.get("context") or {}
+    end = span.get("end_time")
+    return RawSpan(
+        trace_id=str(context.get("trace_id") or span.get("trace_id")),
+        span_id=str(context.get("span_id") or span.get("span_id")),
+        parent_span_id=span.get("parent_id") or None,
+        name=str(span.get("name") or ""),
+        kind=str(span.get("span_kind") or "INTERNAL"),
+        start_time=datetime.fromisoformat(span["start_time"]),
+        end_time=datetime.fromisoformat(end) if end else None,
+        status_code=str(span.get("status_code") or "UNSET"),
+        status_message=span.get("status_message") or None,
+        attributes=flatten_attributes(span.get("attributes") or {}),
+        events=[
+            {
+                "name": e.get("name"),
+                "time": e.get("timestamp"),
+                "attributes": e.get("attributes") or {},
+            }
+            for e in span.get("events") or []
+        ],
+        resource={},
+    )
+
+
+SPAN_PAGE = 1000
+
+
+class PhoenixClient:
+    def __init__(
+        self,
+        url: str,
+        api_key: str | None,
+        page_traces: int,
+        *,
+        transport: httpx.AsyncBaseTransport | None = None,
+    ):
+        self._url = url.rstrip("/")
+        self._api_key = api_key
+        self._page_traces = page_traces
+        self._transport = transport
+
+    async def _pages(
+        self, path: str, params: dict[str, str | list[str]]
+    ) -> AsyncIterator[list[dict[str, Any]]]:
+        cursor: str | None = None
+        async with httpx.AsyncClient(
+            base_url=self._url,
+            headers={"Authorization": f"Bearer {self._api_key}"} if self._api_key else {},
+            timeout=120,
+            transport=self._transport,
+        ) as client:
+            while True:
+                resp = await client.get(
+                    path, params={**params, **({"cursor": cursor} if cursor else {})}
+                )
+                resp.raise_for_status()
+                body = resp.json()
+                rows = body.get("data") or []
+                if rows:
+                    yield rows
+                cursor = body.get("next_cursor")
+                if not cursor or not rows:
+                    return
+
+    def _traces(self, project: str, params: dict[str, str]) -> AsyncIterator[list[dict[str, Any]]]:
+        return self._pages(
+            f"/v1/projects/{project}/traces",
+            {"sort": "start_time", "order": "asc", "limit": str(self._page_traces), **params},
+        )
+
+    async def trace_spans(self, project: str, trace_ids: list[str]) -> AsyncIterator[list[RawSpan]]:
+        params: dict[str, str | list[str]] = {"trace_id": trace_ids, "limit": str(SPAN_PAGE)}
+        async for spans in self._pages(f"/v1/projects/{project}/spans", params):
+            yield [_to_raw(span) for span in spans]
+
+    async def session_spans(self, project: str, session_id: str) -> list[RawSpan]:
+        found: list[RawSpan] = []
+        async for traces in self._traces(project, {"session_identifier": session_id}):
+            async for spans in self.trace_spans(project, [t["trace_id"] for t in traces]):
+                found.extend(spans)
+        return found
+
+    async def trace_pages(
+        self, project: str, since: datetime, until: datetime
+    ) -> AsyncIterator[TracePage]:
+        window = {"start_time": since.isoformat(), "end_time": until.isoformat()}
+        async for traces in self._traces(project, window):
+            yield TracePage(
+                traces=[
+                    (str(trace["trace_id"]), datetime.fromisoformat(trace["start_time"]))
+                    for trace in traces
+                ]
+            )

evalkit/core/trajectory.py

@@ -0,0 +1,248 @@
+import json
+
+from pydantic import BaseModel, Field
+
+from evalkit.contracts.atif import ContentPart, ObservationResult, Step, ToolCall, Trajectory
+from evalkit.core.costmap import cost_usd, resolve_rates
+
+
+def _message_text(message: str | list[ContentPart] | None) -> str:
+    if message is None:
+        return ""
+    if isinstance(message, str):
+        return message
+    return " ".join(part.text for part in message if part.type == "text" and part.text)
+
+
+def _result_text(result: ObservationResult) -> str:
+    if result.content is None:
+        return ""
+    if isinstance(result.content, str):
+        return result.content
+    return " ".join(part.text for part in result.content if part.type == "text" and part.text)
+
+
+class TokenTotals(BaseModel):
+    input: int | None = None
+    output: int | None = None
+    cached: int | None = None
+
+
+class TrajectoryMetrics(BaseModel):
+    agent: str
+    input_tokens: int | None = None
+    output_tokens: int | None = None
+    cached_tokens: int | None = None
+    cost_usd: float | None = None
+    total_tool_calls: int = 0
+    errored_tool_calls: int = 0
+    turns: int = 0
+
+
+class RunMetrics(BaseModel):
+    duration_ms: int | None = None
+    input_tokens: int | None = None
+    output_tokens: int | None = None
+    cached_tokens: int | None = None
+    cost_usd: float | None = None
+    total_tool_calls: int = 0
+    errored_tool_calls: int = 0
+    subagents: int = 0
+    turns: int = 0
+    by_trajectory: dict[str, TrajectoryMetrics] = Field(default_factory=dict)
+
+
+def _add_reported[Num: (int, float)](total: Num | None, value: Num | None) -> Num | None:
+    if total is None:
+        return value
+    if value is None:
+        return total
+    return total + value
+
+
+def _errored_results(step: Step) -> int:
+    if step.observation is None:
+        return 0
+    return sum(1 for r in step.observation.results if (r.extra or {}).get("is_error") is True)
+
+
+def _own_metrics(trajectory: Trajectory, agent: str) -> TrajectoryMetrics:
+    metrics = TrajectoryMetrics(agent=agent)
+    for step in trajectory.steps:
+        if step.source == "agent":
+            metrics.turns += 1
+        metrics.total_tool_calls += len(step.tool_calls or [])
+        metrics.errored_tool_calls += _errored_results(step)
+        if step.metrics is not None:
+            metrics.input_tokens = _add_reported(metrics.input_tokens, step.metrics.prompt_tokens)
+            metrics.output_tokens = _add_reported(
+                metrics.output_tokens, step.metrics.completion_tokens
+            )
+            metrics.cached_tokens = _add_reported(metrics.cached_tokens, step.metrics.cached_tokens)
+            metrics.cost_usd = _add_reported(metrics.cost_usd, step.metrics.cost_usd)
+    if metrics.cost_usd is not None:
+        metrics.cost_usd = round(metrics.cost_usd, 10)
+    return metrics
+
+
+class TrajectoryView:
+    def __init__(self, trajectory: Trajectory):
+        self.trajectory = trajectory
+
+    def _all_steps(self) -> list[Step]:
+        steps = list(self.trajectory.steps)
+        for sub in self.trajectory.subagent_trajectories or []:
+            steps.extend(sub.steps)
+        return steps
+
+    def tool_calls(self) -> list[ToolCall]:
+        calls: list[ToolCall] = []
+        for step in self._all_steps():
+            calls.extend(step.tool_calls or [])
+        return calls
+
+    def errored_tool_call_count(self) -> int:
+        return sum(_errored_results(step) for step in self._all_steps())
+
+    def token_totals(self) -> TokenTotals:
+        final = self.trajectory.final_metrics
+        if final is not None and any(
+            v is not None
+            for v in (
+                final.total_prompt_tokens,
+                final.total_completion_tokens,
+                final.total_cached_tokens,
+            )
+        ):
+            return TokenTotals(
+                input=final.total_prompt_tokens,
+                output=final.total_completion_tokens,
+                cached=final.total_cached_tokens,
+            )
+        totals = {"input": None, "output": None, "cached": None}
+        keys = {"input": "prompt_tokens", "output": "completion_tokens", "cached": "cached_tokens"}
+        for step in self._all_steps():
+            if step.metrics is None:
+                continue
+            for name, attr in keys.items():
+                value = getattr(step.metrics, attr)
+                if value is not None:
+                    totals[name] = (totals[name] or 0) + value
+        return TokenTotals(**totals)
+
+    def turn_count(self) -> int:
+        return sum(1 for step in self.trajectory.steps if step.source == "agent")
+
+    def subagent_count(self) -> int:
+        return len(self.trajectory.subagent_trajectories or [])
+
+    def metrics_by_trajectory(self) -> dict[str, TrajectoryMetrics]:
+        main = self.trajectory
+        split = {"main": _own_metrics(main, main.agent.name if main.agent else "main")}
+        for i, sub in enumerate(self.trajectory.subagent_trajectories or []):
+            key = sub.trajectory_id or f"subagent_trajectories[{i}]"
+            split[key] = _own_metrics(sub, (sub.agent.name if sub.agent else None) or "subagent")
+        return split
+
+    def render_steps(self) -> list[str]:
+        lines: list[str] = []
+        for step in self._all_steps():
+            if step.reasoning_content:
+                lines.append(f"[reasoning] {step.reasoning_content}")
+            text = _message_text(step.message)
+            if text:
+                lines.append(f"[{step.source}] {text}")
+            results_by_call: dict[str, ObservationResult] = {}
+            if step.observation is not None:
+                for result in step.observation.results:
+                    if result.source_call_id:
+                        results_by_call[result.source_call_id] = result
+            for call in step.tool_calls or []:
+                args = json.dumps(call.arguments, ensure_ascii=False, default=str)
+                line = f"[tool] {call.function_name}  args: {args}"
+                matched = results_by_call.get(call.tool_call_id)
+                if matched is not None:
+                    line += f"  result: {_result_text(matched)}"
+                    if (matched.extra or {}).get("is_error") is True:
+                        line += "  [ERROR]"
+                lines.append(line)
+        return [f"{i}. {line}" for i, line in enumerate(lines, 1)]
+
+    def render(self) -> str:
+        numbered = "\n".join(self.render_steps()) or "(none)"
+        return f"Recorded steps (in order):\n{numbered}"
+
+
+def trajectory_model(trajectory: Trajectory) -> str | None:
+    if trajectory.agent and trajectory.agent.model_name:
+        return trajectory.agent.model_name
+    for step in trajectory.steps:
+        if step.model_name:
+            return step.model_name
+    return None
+
+
+def _price_steps(steps: list[Step], cost_map: dict, default_model: str | None) -> None:
+    for step in steps:
+        metrics = step.metrics
+        # Only fill a gap: a step with token counts but no price. A zero from
+        # the harness counts as a gap too: harnesses that lack rates for a
+        # model (pi defaults them to 0) report 0.0 rather than omitting it.
+        if metrics is None or metrics.cost_usd or metrics.prompt_tokens is None:
+            continue
+        model = step.model_name or default_model
+        if model is None:
+            continue
+        rates = resolve_rates(cost_map, model)
+        if rates is None:
+            continue
+        metrics.cost_usd = cost_usd(
+            rates,
+            metrics.prompt_tokens,
+            metrics.completion_tokens or 0,
+            metrics.cached_tokens or 0,
+        )
+
+
+def backfill_step_costs(trajectory: Trajectory, cost_map: dict, model: str | None = None) -> None:
+    # The run's model override is the fallback when the harness didn't record a
+    # model on the trajectory; a step's own model_name still wins.
+    default_model = model or trajectory_model(trajectory)
+    _price_steps(trajectory.steps, cost_map, default_model)
+    for sub in trajectory.subagent_trajectories or []:
+        _price_steps(sub.steps, cost_map, trajectory_model(sub) or default_model)
+
+
+def _sum_step_costs(view: TrajectoryView) -> float | None:
+    costs = [
+        step.metrics.cost_usd
+        for step in view._all_steps()
+        if step.metrics is not None and step.metrics.cost_usd is not None
+    ]
+    return round(sum(costs), 10) if costs else None
+
+
+def derive_metrics(view: TrajectoryView, wall_ms: int | None) -> RunMetrics:
+    totals = view.token_totals()
+    final = view.trajectory.final_metrics
+    # Prefer a harness-reported total; otherwise sum the per-step costs.
+    # backfill_step_costs prices any step the harness left unset before the
+    # trajectory is stored, so the gauge and the UI cost view read the same
+    # per-step numbers.
+    cost = (
+        final.total_cost_usd
+        if (final is not None and final.total_cost_usd)
+        else _sum_step_costs(view)
+    )
+    return RunMetrics(
+        duration_ms=wall_ms,
+        input_tokens=totals.input,
+        output_tokens=totals.output,
+        cached_tokens=totals.cached,
+        cost_usd=cost,
+        total_tool_calls=len(view.tool_calls()),
+        errored_tool_calls=view.errored_tool_call_count(),
+        subagents=view.subagent_count(),
+        turns=view.turn_count(),
+        by_trajectory=view.metrics_by_trajectory(),
+    )

evalkit/core/workspace.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from typing import NamedTuple
+
+
+class Reference(NamedTuple):
+    name: str
+    workspace: str
+
+
+class BareReference(ValueError):
+    def __init__(self, text: str):
+        super().__init__(f"use name@workspace, not {text!r}")
+
+
+def parse_reference(text: str, default_workspace: str) -> Reference:
+    name, sep, workspace = text.partition("@")
+    if not name or (sep and (not workspace or "@" in workspace)):
+        raise ValueError(f"malformed reference: {text!r}")
+    return Reference(name, workspace if sep else default_workspace)
+
+
+def parse_slug(value: str) -> Reference:
+    if "@" not in value:
+        raise BareReference(value)
+    return parse_reference(value, "")
+
+
+def slug(ref: Reference) -> str:
+    return f"{ref.name}@{ref.workspace}"

evalkit/credentials.py

@@ -0,0 +1,119 @@
+import os
+import re
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Literal, Protocol
+
+from dotenv import dotenv_values
+from pydantic import BaseModel, ConfigDict, field_validator, model_validator
+
+Kind = Literal["agent", "evaluator", "llm_provider"]
+
+
+class MissingCredential(ValueError):
+    pass
+
+
+ENV_REF = re.compile(r"^env:[A-Za-z_][A-Za-z0-9_]*$")
+
+
+class CredentialsConfig(BaseModel):
+    model_config = ConfigDict(extra="forbid", hide_input_in_errors=True)
+
+    agents: dict[str, str] = {}
+    evaluators: dict[str, str] = {}
+    llm_providers: dict[str, str] = {}
+
+    # TODO: remove in 0.5.0, two minor versions after 0.3.0 renamed credentials.judges
+    # to credentials.llm_providers.
+    @model_validator(mode="before")
+    @classmethod
+    def _judges_renamed(cls, data: object) -> object:
+        if isinstance(data, dict) and "judges" in data:
+            raise ValueError("credentials.judges is now credentials.llm_providers")
+        return data
+
+    @field_validator("agents", "evaluators", "llm_providers")
+    @classmethod
+    def _env_refs_only(cls, value: dict[str, str]) -> dict[str, str]:
+        bad = [k for k, v in value.items() if not ENV_REF.match(v)]
+        if bad:
+            raise ValueError(f"credentials must be env:NAME references (no defaults): {bad}")
+        return value
+
+    @field_validator("llm_providers")
+    @classmethod
+    def _provider_keys_are_urls_or_env_refs(cls, value: dict[str, str]) -> dict[str, str]:
+        bad = [k for k in value if k.startswith("env:") and not ENV_REF.match(k)]
+        if bad:
+            raise ValueError(f"a credentials.llm_providers key is a base URL or env:NAME: {bad}")
+        return value
+
+
+class KeySource(Protocol):
+    def key_for(self, kind: Kind, ref: str, url: str | None = None) -> str | None: ...
+
+    def require_judge_key(self, base_url: str) -> str | None: ...
+
+
+class CredentialRegistry:
+    def __init__(
+        self, cfg: CredentialsConfig, environ: Mapping[str, str], *, home: str | None = None
+    ):
+        self._cfg = cfg
+        self._env = environ
+        self._home = home
+
+    def _table(self, kind: Kind) -> dict[str, str]:
+        if kind == "llm_provider":
+            return {
+                url: v
+                for k, v in self._cfg.llm_providers.items()
+                if (url := self._judge_url(k)) is not None
+            }
+        return {"agent": self._cfg.agents, "evaluator": self._cfg.evaluators}[kind]
+
+    def _judge_url(self, key: str) -> str | None:
+        if not key.startswith("env:"):
+            return key
+        return self._env.get(key.removeprefix("env:"))
+
+    def _unset_judge_keys(self) -> list[str]:
+        return [
+            k
+            for k in self._cfg.llm_providers
+            if k.startswith("env:") and self._judge_url(k) is None
+        ]
+
+    def _entry(self, kind: Kind, ref: str) -> str | None:
+        table = self._table(kind)
+        if ref in table:
+            return table[ref]
+        name, _, workspace = ref.partition("@")
+        if kind != "llm_provider" and self._home is not None and workspace == self._home:
+            return table.get(name)
+        return None
+
+    def key_for(self, kind: Kind, ref: str, url: str | None = None) -> str | None:
+        entry = self._entry(kind, ref)
+        if entry is None:
+            return None
+        name = entry.removeprefix("env:")
+        if name not in self._env:
+            raise MissingCredential(f"credentials.{kind}s[{ref!r}] needs env var {name}")
+        return self._env[name]
+
+    def require_judge_key(self, base_url: str) -> str:
+        key = self.key_for("llm_provider", base_url)
+        if key is None:
+            unset = self._unset_judge_keys()
+            hint = f" (unset: {', '.join(k.removeprefix('env:') for k in unset)})" if unset else ""
+            raise MissingCredential(f"no credentials.llm_providers entry for {base_url}{hint}")
+        return key
+
+
+def load_dotenv_beside(config_path: Path) -> dict[str, str]:
+    file_values = {
+        k: v for k, v in dotenv_values(config_path.parent / ".env").items() if v is not None
+    }
+    return {**file_values, **os.environ}