world-model-optimizer 0.2.0__py3-none-any.whl

wmo/serving/traces_source.py

@@ -0,0 +1,206 @@
+"""Serve-side trace access: read a model's local traces, or fetch them from the Hugging Face Hub.
+
+The raw trace corpus (`traces.otel.jsonl`) is large and need not be committed. When it is present
+locally it is used directly; otherwise, if the model's card declares a `traces_hf` source, the
+backend streams it from the Hub's public resolve URL (no auth, no client-side Hub API) into the
+model directory, reporting byte progress the website can poll. A local copy always supersedes the
+Hub. Recorded traces are grouped by task into replayable scenarios for the Explore-traces tab.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Callable
+from enum import StrEnum
+from pathlib import Path
+
+import httpx
+from pydantic import BaseModel, Field
+
+from wmo.config.card import TracesSource
+from wmo.core.types import Action
+from wmo.ingest import get_adapter
+
+TRACES_FILENAME = "traces.otel.jsonl"
+
+
+def resolve_url(source: TracesSource) -> str:
+    """The public Hub resolve URL for a traces source (works unauthenticated for public repos)."""
+    prefix = "datasets/" if source.kind == "dataset" else ""
+    return f"https://huggingface.co/{prefix}{source.repo}/resolve/{source.revision}/{source.path}"
+
+
+def local_traces_path(model_dir: Path) -> Path | None:
+    """The traces file for a model, if present: a downloaded copy, else the example sibling."""
+    downloaded = model_dir / TRACES_FILENAME
+    if downloaded.is_file():
+        return downloaded
+    # examples/<task>/traces.otel.jsonl sits two levels above examples/<task>/models/<name>/.
+    sibling = model_dir.parent.parent / TRACES_FILENAME
+    if sibling.is_file():
+        return sibling
+    return None
+
+
+def _action_label(action: Action) -> str:
+    """Format an action in the wmo-play grammar (matches the web index generator)."""
+    if action.kind.value == "tool_call":
+        return f"{action.name} {action.arguments}" if action.arguments else (action.name or "")
+    return f"say {action.content or ''}"
+
+
+class ScenarioStep(BaseModel):
+    action: Action
+    action_label: str
+    observation: str
+    is_error: bool
+
+
+class TraceScenario(BaseModel):
+    id: str
+    label: str
+    task: str | None
+    steps: list[ScenarioStep]
+
+
+def _clip(text: str, limit: int) -> str:
+    flat = " ".join(text.split())
+    return flat if len(flat) <= limit else flat[: limit - 1] + "…"
+
+
+def scenarios_from_traces(
+    path: Path, *, max_scenarios: int = 6, max_steps: int = 10
+) -> list[TraceScenario]:
+    """Normalize a traces.otel.jsonl into a bounded list of replayable scenarios (one per trace)."""
+    traces = get_adapter("otel-genai").from_file(str(path))
+    out: list[TraceScenario] = []
+    for i, trace in enumerate(traces):
+        steps: list[ScenarioStep] = []
+        for step in trace.steps:
+            if step.observation.content is None:
+                continue
+            steps.append(
+                ScenarioStep(
+                    action=step.action,
+                    action_label=_clip(_action_label(step.action), 100),
+                    observation=step.observation.content,
+                    is_error=step.observation.is_error,
+                )
+            )
+            if len(steps) >= max_steps:
+                break
+        if not steps:
+            continue
+        task = trace.steps[0].task if trace.steps else None
+        out.append(
+            TraceScenario(id=f"t{i}", label=_scenario_label(task, i), task=task, steps=steps)
+        )
+        if len(out) >= max_scenarios:
+            break
+    return out
+
+
+def _scenario_label(task: str | None, index: int) -> str:
+    if not task:
+        return f"Scenario {index + 1}"
+    import json
+
+    try:
+        parsed = json.loads(task)
+        text = parsed.get("reason_for_call") or parsed.get("task_instructions") or task
+    except (json.JSONDecodeError, AttributeError):
+        text = task
+    return _clip(text, 68)
+
+
+class DownloadStatus(StrEnum):
+    RUNNING = "running"
+    DONE = "done"
+    FAILED = "failed"
+
+
+class DownloadProgress(BaseModel):
+    status: DownloadStatus
+    downloaded: int = 0
+    total: int | None = None  # bytes, when the server reports Content-Length
+    error: str | None = None
+
+
+class _DownloadState:
+    def __init__(self) -> None:
+        self.progress = DownloadProgress(status=DownloadStatus.RUNNING)
+        self.lock = threading.Lock()
+
+
+class TracesDownloader:
+    """Runs Hub trace downloads on background threads, one in flight per model name."""
+
+    def __init__(
+        self, *, fetch: Callable[[str, Path, Callable[[int, int | None], None]], None] | None = None
+    ) -> None:
+        self._fetch = fetch or _stream_to_file
+        self._states: dict[str, _DownloadState] = {}
+        self._lock = threading.Lock()
+
+    def start(self, name: str, url: str, dest: Path) -> None:
+        with self._lock:
+            existing = self._states.get(name)
+            if existing and existing.progress.status is DownloadStatus.RUNNING:
+                return  # already downloading; the client just polls
+            self._states[name] = _DownloadState()
+        thread = threading.Thread(
+            target=self._run, args=(name, url, dest), name=f"wmo-traces-{name}", daemon=True
+        )
+        thread.start()
+
+    def _run(self, name: str, url: str, dest: Path) -> None:
+        state = self._states[name]
+
+        def on_progress(downloaded: int, total: int | None) -> None:
+            with state.lock:
+                state.progress.downloaded = downloaded
+                state.progress.total = total
+
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        tmp = dest.with_suffix(dest.suffix + ".part")
+        try:
+            self._fetch(url, tmp, on_progress)
+            tmp.replace(dest)  # atomic: a partial download never looks complete
+        except Exception as exc:  # noqa: BLE001 - report any failure to the client
+            tmp.unlink(missing_ok=True)
+            with state.lock:
+                state.progress.status = DownloadStatus.FAILED
+                state.progress.error = str(exc)
+            return
+        with state.lock:
+            state.progress.status = DownloadStatus.DONE
+
+    def progress(self, name: str) -> DownloadProgress | None:
+        state = self._states.get(name)
+        if state is None:
+            return None
+        with state.lock:
+            return state.progress.model_copy()
+
+
+def _stream_to_file(url: str, dest: Path, on_progress: Callable[[int, int | None], None]) -> None:
+    """Stream a public Hub file to disk in chunks, following the CDN redirect."""
+    with httpx.stream("GET", url, follow_redirects=True, timeout=60.0) as resp:
+        resp.raise_for_status()
+        raw = resp.headers.get("content-length")
+        total = int(raw) if raw and raw.isdigit() else None
+        downloaded = 0
+        with dest.open("wb") as fh:
+            for chunk in resp.iter_bytes(1024 * 1024):
+                fh.write(chunk)
+                downloaded += len(chunk)
+                on_progress(downloaded, total)
+
+
+class TracesResponse(BaseModel):
+    """What the Explore-traces tab needs: local scenarios if present, else a Hub download offer."""
+
+    source: str  # "local" | "hub" | "none"
+    downloadable: bool
+    scenarios: list[TraceScenario] = Field(default_factory=list)
+    download: DownloadProgress | None = None

wmo/telemetry.py

@@ -0,0 +1,213 @@
+"""Best-effort anonymous usage telemetry."""
+
+from __future__ import annotations
+
+import os
+import sys
+from atexit import register
+from dataclasses import dataclass
+from importlib.metadata import PackageNotFoundError, version
+from pathlib import Path
+
+from posthog import Posthog
+
+from wmo.config import ARTIFACT_DIR
+from wmo.config.settings import ensure_telemetry_anonymous_id, load_settings
+from wmo.engine.reporting import BuildReporter
+from wmo.tracking import RunRecord
+
+POSTHOG_PROJECT_API_KEY = "phc_rPFfCufWpxyctR7duEZTTXovP4k5kbHqSqzd4Z4MQJdL"
+POSTHOG_HOST = "https://us.i.posthog.com"
+
+TelemetryValue = str | int | float | bool | None
+TelemetryProperties = dict[str, TelemetryValue]
+
+_FALSE_VALUES = {"0", "false", "off", "no"}
+_TRUE_VALUES = {"1", "true", "on", "yes"}
+_CLIENTS: dict[tuple[str, str], Posthog] = {}
+
+
+@dataclass
+class BuildTelemetryStats:
+    input_trace_count: int = 0
+    input_step_count: int = 0
+    train_trace_count: int = 0
+    val_trace_count: int = 0
+    heldout_trace_count: int = 0
+    indexed_step_count: int = 0
+
+
+class TelemetryBuildReporter:
+    def __init__(self, inner: BuildReporter, stats: BuildTelemetryStats) -> None:
+        self._inner = inner
+        self._stats = stats
+
+    def ingest_done(self, traces: int, steps: int) -> None:
+        self._stats.input_trace_count = traces
+        self._stats.input_step_count = steps
+        self._inner.ingest_done(traces, steps)
+
+    def split_done(self, train: int, val: int, test: int) -> None:
+        self._stats.train_trace_count = train
+        self._stats.val_trace_count = val
+        self._stats.heldout_trace_count = test
+        self._inner.split_done(train, val, test)
+
+    def index_done(self, steps: int) -> None:
+        self._stats.indexed_step_count = steps
+        self._inner.index_done(steps)
+
+    def optimize_start(self, budget: int) -> None:
+        self._inner.optimize_start(budget)
+
+    def activity(self, line: str) -> None:
+        self._inner.activity(line)
+
+    def rollout(self, done: int, budget: int, score: float | None) -> None:
+        self._inner.rollout(done, budget, score)
+
+    def optimize_done(self, held_out_accuracy: float, frontier_size: int, rollouts: int) -> None:
+        self._inner.optimize_done(held_out_accuracy, frontier_size, rollouts)
+
+
+def capture(
+    event: str,
+    properties: TelemetryProperties | None = None,
+    *,
+    root: str | Path = ARTIFACT_DIR,
+) -> bool:
+    """Send one anonymous metadata-only event. Returns False when skipped or failed."""
+    if not _enabled(root):
+        return False
+    api_key = os.getenv("WMO_POSTHOG_PROJECT_API_KEY", POSTHOG_PROJECT_API_KEY).strip()
+    if not api_key:
+        return False
+    host = os.getenv("WMO_POSTHOG_HOST", POSTHOG_HOST).rstrip("/")
+    try:
+        distinct_id = ensure_telemetry_anonymous_id(root)
+        event_properties: TelemetryProperties = {
+            "$process_person_profile": False,
+            "wmo_version": _wmo_version(),
+            "python_version": f"{sys.version_info.major}.{sys.version_info.minor}",
+            **(properties or {}),
+        }
+        # Never log prompts, traces, actions, observations, paths, models, credentials, or text.
+        message_id = _posthog_client(api_key, host).capture(
+            event,
+            distinct_id=distinct_id,
+            properties=event_properties,
+        )
+        return message_id is not None
+    except (OSError, ValueError):
+        return False
+
+
+def capture_build_completed(
+    *,
+    stats: BuildTelemetryStats,
+    gepa_budget: int,
+    rollouts_used: int,
+    frontier_size: int,
+    record: RunRecord,
+    root: str | Path,
+) -> None:
+    capture(
+        "wmo build completed",
+        {
+            "success": True,
+            "input_trace_count": stats.input_trace_count,
+            "input_step_count": stats.input_step_count,
+            "train_trace_count": stats.train_trace_count,
+            "val_trace_count": stats.val_trace_count,
+            "heldout_trace_count": stats.heldout_trace_count,
+            "indexed_step_count": stats.indexed_step_count,
+            "gepa_budget": gepa_budget,
+            "rollouts_used": rollouts_used,
+            "frontier_size": frontier_size,
+            "duration_seconds": round(record.duration_seconds, 3),
+            "llm_call_count": record.total.calls,
+            "input_tokens": record.total.input_tokens,
+            "output_tokens": record.total.output_tokens,
+            "cost_usd": round(record.total.cost_usd, 6),
+        },
+        root=root,
+    )
+
+
+def capture_eval_completed(
+    *,
+    mode: str,
+    file_count: int,
+    scored_step_count: int,
+    rag_enabled: bool,
+    sample_turns: str,
+    train_split: float,
+    top_k: int,
+    root: str | Path,
+) -> None:
+    capture(
+        "wmo eval completed",
+        {
+            "success": True,
+            "eval_mode": mode,
+            "file_count": file_count,
+            "scored_step_count": scored_step_count,
+            "rag_enabled": rag_enabled,
+            "sample_turns": sample_turns,
+            "train_split": train_split,
+            "top_k": top_k,
+        },
+        root=root,
+    )
+
+
+def settings_root_from_results_root(results_root: str) -> Path:
+    path = Path(results_root)
+    return path.parent if path.name == "evals" else Path(ARTIFACT_DIR)
+
+
+def _enabled(root: str | Path) -> bool:
+    if _env_truthy("DO_NOT_TRACK"):
+        return False
+    env = os.getenv("WMO_TELEMETRY")
+    if env is not None:
+        return env.strip().lower() in _TRUE_VALUES
+    if os.getenv("PYTEST_CURRENT_TEST"):
+        return False
+    try:
+        return load_settings(root).telemetry.enabled
+    except (OSError, ValueError):
+        return False
+
+
+def _env_truthy(name: str) -> bool:
+    value = os.getenv(name)
+    if value is None:
+        return False
+    normalized = value.strip().lower()
+    if normalized in _FALSE_VALUES:
+        return False
+    return bool(normalized)
+
+
+def _wmo_version() -> str:
+    try:
+        return version("world-model-optimizer")
+    except PackageNotFoundError:
+        return "unknown"
+
+
+def _posthog_client(api_key: str, host: str) -> Posthog:
+    key = (api_key, host)
+    client = _CLIENTS.get(key)
+    if client is None:
+        client = Posthog(
+            api_key,
+            host=host,
+            flush_interval=1.0,
+            max_retries=1,
+            timeout=0.5,
+        )
+        _CLIENTS[key] = client
+        register(client.shutdown)
+    return client

wmo/tracking/__init__.py

@@ -0,0 +1,36 @@
+"""Run tracking: time + cost + tokens across the harness lifecycle.
+
+Instrument at the provider boundary (`MeteredProvider`) so the world model, GEPA, and the judge are
+all metered without changes to those modules. `RunTracker` aggregates `UsageEvent`s into
+`UsageTotals` (priced via `wmo.tracking.pricing`) plus a wall-clock duration from an injectable
+`Clock`; `RunRecord`s persist under `.wmo/runs/`.
+"""
+
+from wmo.tracking.clock import Clock, SystemClock
+from wmo.tracking.metered import MeteredProvider, classify_build_call
+from wmo.tracking.pricing import ModelPrice, cost_usd, price_for
+from wmo.tracking.store import load_runs, save_run
+from wmo.tracking.tracker import (
+    Phase,
+    RunRecord,
+    RunTracker,
+    UsageEvent,
+    UsageTotals,
+)
+
+__all__ = [
+    "Clock",
+    "SystemClock",
+    "MeteredProvider",
+    "classify_build_call",
+    "ModelPrice",
+    "cost_usd",
+    "price_for",
+    "load_runs",
+    "save_run",
+    "Phase",
+    "RunRecord",
+    "RunTracker",
+    "UsageEvent",
+    "UsageTotals",
+]

wmo/tracking/clock.py

@@ -0,0 +1,24 @@
+"""Injectable clock so run-tracking durations are deterministic in tests.
+
+Production uses `SystemClock` (real `time.monotonic`). Tests pass a `FakeClock` with scripted ticks
+so cost/time assertions don't depend on wall-clock timing.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Protocol, runtime_checkable
+
+
+@runtime_checkable
+class Clock(Protocol):
+    def monotonic(self) -> float:
+        """Seconds from an arbitrary epoch; only differences are meaningful (for durations)."""
+        ...
+
+
+class SystemClock:
+    """Default clock backed by `time.monotonic` (monotonic, unaffected by wall-clock changes)."""
+
+    def monotonic(self) -> float:
+        return time.monotonic()

wmo/tracking/metered.py

@@ -0,0 +1,125 @@
+"""`MeteredProvider`: a Provider wrapper that records every call onto a RunTracker.
+
+Instrumenting at the provider boundary means the optimizer, the judge, and the world model are all
+metered without any of them knowing about tracking — we don't edit `gepa.py` or the judge. The
+wrapper forwards `complete`/`embed`/`verify`/`config` to the wrapped provider unchanged and records
+a `UsageEvent` per call.
+
+Phase attribution: build and serve share one provider, and within build the *same* provider serves
+GEPA rollouts, GEPA reflection, and the judge. We tell them apart by the system prompt each path
+uses (a stable, boundary-visible signal), defaulting `complete` to whatever base phase the wrapper
+was constructed with. Callers that want exact control can pass their own `classify`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator
+
+from wmo.optimize.judge import JUDGE_MARKER
+from wmo.providers.base import (
+    DEFAULT_MAX_TOKENS,
+    Completion,
+    Message,
+    Provider,
+    ProviderConfig,
+    StreamChunk,
+    StreamingProvider,
+    TokenUsage,
+    VerifyResult,
+)
+from wmo.tracking.tracker import Phase, RunTracker
+
+
+def classify_build_call(system: str) -> Phase:
+    """Default phase classifier for a build-time `complete`: judge vs GEPA (rollout/reflection).
+
+    The judge is recognized by `JUDGE_MARKER` (owned by judge.py, so the prompt and this classifier
+    can't drift). Everything else during build — GEPA rollouts (env-sim) and reflection — is GEPA.
+    """
+    if JUDGE_MARKER in system:
+        return Phase.JUDGE
+    return Phase.GEPA
+
+
+class MeteredProvider:
+    """Wraps a `Provider`, recording token usage + cost per call onto a `RunTracker`.
+
+    `base_phase` is the phase for `complete` calls when no `classify` is given (e.g. `Phase.SERVE`
+    for the live world model). For build, pass `classify=classify_build_call` to split judge from
+    GEPA.
+    """
+
+    def __init__(
+        self,
+        provider: Provider,
+        tracker: RunTracker,
+        *,
+        base_phase: Phase = Phase.OTHER,
+        classify: Callable[[str], Phase] | None = None,
+    ) -> None:
+        self._provider = provider
+        self._tracker = tracker
+        self._base_phase = base_phase
+        self._classify = classify
+
+    @property
+    def config(self) -> ProviderConfig:
+        return self._provider.config
+
+    def complete(
+        self,
+        system: str,
+        messages: list[Message],
+        *,
+        temperature: float = 0.7,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+    ) -> Completion:
+        completion = self._provider.complete(
+            system, messages, temperature=temperature, max_tokens=max_tokens
+        )
+        phase = self._classify(system) if self._classify is not None else self._base_phase
+        # Prefer the model the completion says actually served (failover chains set it); the
+        # configured model otherwise (Completion.model validates non-empty, so `or` is exact).
+        # Pricing a failed-over call at the primary's rate would silently mis-report cost.
+        model = completion.model or self._provider.config.model
+        self._tracker.record(phase, model, completion.usage)
+        return completion
+
+    def stream(
+        self,
+        system: str,
+        messages: list[Message],
+        *,
+        temperature: float = 0.7,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+    ) -> Iterator[StreamChunk]:
+        """Forward a native stream, recording usage from the terminal chunk.
+
+        A stream abandoned before its terminal chunk records nothing even though the provider
+        billed the partial generation; serving owns closing streams it starts.
+        """
+        if not isinstance(self._provider, StreamingProvider):
+            raise TypeError(
+                f"{type(self._provider).__name__} does not implement StreamingProvider; "
+                "stream() is only available for native backends"
+            )
+        phase = self._classify(system) if self._classify is not None else self._base_phase
+        for chunk in self._provider.stream(
+            system, messages, temperature=temperature, max_tokens=max_tokens
+        ):
+            if chunk.done and chunk.usage is not None:
+                model = chunk.model or self._provider.config.model
+                self._tracker.record(phase, model, chunk.usage)
+            yield chunk
+
+    def embed(self, texts: list[str]) -> list[list[float]]:
+        # Embeddings carry no token usage from our providers; record a zero-usage event for the
+        # call count so EMBED shows up in the breakdown. Attribute it to the embeddings model
+        # (`embed_model`), not the completion model, so any future embed pricing is keyed right.
+        vectors = self._provider.embed(texts)
+        embed_model = self._provider.config.embed_model or self._provider.config.model
+        self._tracker.record(Phase.EMBED, embed_model, TokenUsage())
+        return vectors
+
+    def verify(self) -> VerifyResult:
+        return self._provider.verify()