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

wmo/tracking/pricing.py

@@ -0,0 +1,99 @@
+"""Per-model token pricing → USD cost.
+
+Provider-agnostic: prices are keyed by a normalized model id (provider prefixes like Bedrock's
+`us.anthropic.` are stripped before lookup), so the same Opus 4.8 row covers the direct API and
+Bedrock. Prices are USD per 1M tokens; an unknown model costs 0.0 and is flagged so callers can
+surface "cost unavailable" rather than silently under-reporting.
+"""
+
+from __future__ import annotations
+
+import re
+
+from pydantic import BaseModel
+
+from wmo.providers.base import TokenUsage
+
+# Bedrock appends a snapshot date and/or version to the model id, e.g.
+# `claude-haiku-4-5-20251001-v1:0` or `claude-opus-4-6-v1`. Strip them so the lookup key matches the
+# undated table rows (`claude-haiku-4-5`). Only applied to `claude-*` ids.
+_BEDROCK_SUFFIX = re.compile(r"(-\d{8})?(-v\d+)?(:\d+)?$")
+
+
+class ModelPrice(BaseModel):
+    """USD per 1,000,000 tokens, split by input/output."""
+
+    input_per_mtok: float
+    output_per_mtok: float
+
+
+# Keyed by normalized model id (see `_normalize`). USD per 1M tokens.
+#
+# Completion prices verified 2026-06-25 against the live vendor pricing pages:
+#   - Claude: platform.claude.com/docs/en/about-claude/models/overview
+#   - OpenAI GPT-5.x: developers.openai.com/api/docs/pricing (Standard tier, short context)
+# Embedding prices are long-stable list prices NOT re-fetched in that pass (the OpenAI pricing
+# page no longer surfaces them); treat as approximate and re-verify if embed cost matters.
+_PRICES: dict[str, ModelPrice] = {
+    # --- Anthropic / Bedrock (Claude) ---
+    "claude-fable-5": ModelPrice(input_per_mtok=10.0, output_per_mtok=50.0),
+    "claude-mythos-5": ModelPrice(input_per_mtok=10.0, output_per_mtok=50.0),
+    "claude-opus-4-8": ModelPrice(input_per_mtok=5.0, output_per_mtok=25.0),
+    "claude-opus-4-7": ModelPrice(input_per_mtok=5.0, output_per_mtok=25.0),
+    "claude-opus-4-6": ModelPrice(input_per_mtok=5.0, output_per_mtok=25.0),
+    "claude-opus-4-5": ModelPrice(input_per_mtok=5.0, output_per_mtok=25.0),
+    "claude-opus-4-1": ModelPrice(input_per_mtok=15.0, output_per_mtok=75.0),
+    "claude-sonnet-5": ModelPrice(input_per_mtok=3.0, output_per_mtok=15.0),
+    "claude-sonnet-4-6": ModelPrice(input_per_mtok=3.0, output_per_mtok=15.0),
+    "claude-haiku-4-5": ModelPrice(input_per_mtok=1.0, output_per_mtok=5.0),
+    # --- OpenAI / Azure OpenAI (GPT-5.x; Azure deployments reuse the base model's price) ---
+    "gpt-5.5": ModelPrice(input_per_mtok=5.0, output_per_mtok=30.0),
+    "gpt-5.5-pro": ModelPrice(input_per_mtok=30.0, output_per_mtok=180.0),
+    "gpt-5.4": ModelPrice(input_per_mtok=2.5, output_per_mtok=15.0),
+    "gpt-5.4-mini": ModelPrice(input_per_mtok=0.75, output_per_mtok=4.5),
+    "gpt-5.4-nano": ModelPrice(input_per_mtok=0.2, output_per_mtok=1.25),
+    # Self-hosted models (vLLM on our own GPUs) intentionally have NO row: their cost is amortized
+    # GPU time, not a per-token API price. `price_for` returns None for them, which the eval grid
+    # renders as "no cost"; a 0.0 ModelPrice would instead report a misleading $0.00.
+    # --- Embeddings (output tokens are always 0 for embed calls) ---
+    "text-embedding-3-small": ModelPrice(input_per_mtok=0.02, output_per_mtok=0.0),
+    "text-embedding-3-large": ModelPrice(input_per_mtok=0.13, output_per_mtok=0.0),
+    "amazon.titan-embed-text-v2:0": ModelPrice(input_per_mtok=0.02, output_per_mtok=0.0),
+}
+
+
+def _normalize(model: str) -> str:
+    """Strip provider/region routing prefixes so one row covers a model across providers.
+
+    Bedrock ids look like `us.anthropic.claude-opus-4-8`; the direct API uses `claude-opus-4-8`.
+    We drop a leading region segment (`us.`/`eu.`/...) and an `anthropic.` vendor segment, but keep
+    `amazon.titan-...` (its `amazon.` is part of the canonical model id, not a routing prefix).
+    """
+    normalized = model.strip()
+    region_prefixes = ("us.", "eu.", "apac.", "us-gov.", "global.", "jp.", "au.", "ca.")
+    for prefix in region_prefixes:
+        if normalized.startswith(prefix):
+            normalized = normalized[len(prefix) :]
+            break
+    if normalized.startswith("anthropic."):
+        normalized = normalized[len("anthropic.") :]
+    if normalized.startswith("claude-"):
+        # Drop a trailing Bedrock snapshot date / version (`-20251001-v1:0`, `-v1`) so dated
+        # inference-profile ids match the undated table rows.
+        normalized = _BEDROCK_SUFFIX.sub("", normalized)
+    return normalized
+
+
+def price_for(model: str) -> ModelPrice | None:
+    """Return the price row for `model` (after normalization), or None if unknown."""
+    return _PRICES.get(_normalize(model))
+
+
+def cost_usd(model: str, usage: TokenUsage) -> float:
+    """USD cost of `usage` on `model`. Unknown models cost 0.0 (see `price_for` to detect that)."""
+    price = price_for(model)
+    if price is None:
+        return 0.0
+    return (
+        usage.input_tokens * price.input_per_mtok + usage.output_tokens * price.output_per_mtok
+    ) / 1_000_000

wmo/tracking/store.py

@@ -0,0 +1,31 @@
+"""Persist + list run records under `.wmo/runs/`.
+
+One JSON file per run (`<run_id>.json`). Kept tiny and dependency-free so `wmo build`/`serve` can
+write a record without pulling in the rest of the harness.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from wmo.tracking.tracker import RunRecord
+
+
+def save_run(record: RunRecord, runs_dir: str | Path) -> Path:
+    """Write `record` to `<runs_dir>/<run_id>.json`, creating the directory if needed."""
+    path = Path(runs_dir)
+    path.mkdir(parents=True, exist_ok=True)
+    out = path / f"{record.run_id}.json"
+    out.write_text(record.model_dump_json(indent=2), encoding="utf-8")
+    return out
+
+
+def load_runs(runs_dir: str | Path) -> list[RunRecord]:
+    """Load all run records from `runs_dir` (empty list if the directory doesn't exist)."""
+    path = Path(runs_dir)
+    if not path.exists():
+        return []
+    return [
+        RunRecord.model_validate_json(p.read_text(encoding="utf-8"))
+        for p in sorted(path.glob("*.json"))
+    ]

wmo/tracking/tracker.py

@@ -0,0 +1,149 @@
+"""Run tracking: aggregate tokens → cost + wall-clock across the harness lifecycle.
+
+A `RunTracker` collects `UsageEvent`s (one per LLM call, tagged by phase) and rolls them up into
+`UsageTotals` (tokens, USD, calls) plus a wall-clock duration measured off an injectable `Clock`.
+`RunRecord` is the persisted artifact (`.wmo/runs/<run_id>.json`).
+
+The tracker is provider-agnostic: it records `(model, TokenUsage)` and prices via
+`wmo.tracking.pricing`. It's fed at the provider boundary by `MeteredProvider` (so GEPA, the judge,
+and the world model are all captured without touching the optimizer), and directly by the world
+model's serve `step`.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections import defaultdict
+from collections.abc import Iterator
+from contextlib import contextmanager
+from enum import StrEnum
+
+from pydantic import BaseModel, Field
+
+from wmo.providers.base import TokenUsage
+from wmo.tracking.clock import Clock, SystemClock
+from wmo.tracking.pricing import cost_usd
+
+
+class Phase(StrEnum):
+    """Lifecycle phase a usage event is attributed to."""
+
+    BUILD = "build"  # world-model build (overall)
+    GEPA = "gepa"  # GEPA rollouts + reflection during optimization
+    JUDGE = "judge"  # LLM-judge scoring
+    SERVE = "serve"  # live world-model step calls
+    EMBED = "embed"  # embedding calls (phi)
+    OTHER = "other"
+
+
+class UsageEvent(BaseModel):
+    """One metered LLM call."""
+
+    phase: Phase
+    model: str
+    usage: TokenUsage = Field(default_factory=TokenUsage)
+    cost_usd: float = 0.0
+
+
+class UsageTotals(BaseModel):
+    """Rolled-up usage: tokens, cost, and call count."""
+
+    calls: int = 0
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cost_usd: float = 0.0
+
+    @property
+    def total_tokens(self) -> int:
+        return self.input_tokens + self.output_tokens
+
+    def _add(self, event: UsageEvent) -> None:
+        self.calls += 1
+        self.input_tokens += event.usage.input_tokens
+        self.output_tokens += event.usage.output_tokens
+        self.cost_usd += event.cost_usd
+
+
+class RunRecord(BaseModel):
+    """Persisted summary of one run (build or serve session)."""
+
+    run_id: str
+    kind: str  # "build" | "serve" | ... (free-form label for the run)
+    duration_seconds: float = 0.0
+    total: UsageTotals = Field(default_factory=UsageTotals)
+    by_phase: dict[Phase, UsageTotals] = Field(default_factory=dict)
+
+
+class RunTracker:
+    """Accumulates usage events and the wall-clock of a run.
+
+    Duration is measured between `start()` and `stop()` off the injected `Clock`, so tests can pass
+    a `FakeClock` and assert exact seconds. `record` is the single entry point for a metered call.
+    """
+
+    def __init__(self, run_id: str, kind: str, clock: Clock | None = None) -> None:
+        self._run_id = run_id
+        self._kind = kind
+        self._clock = clock or SystemClock()
+        self._events: list[UsageEvent] = []
+        self._lock = threading.Lock()
+        self._started_at: float | None = None
+        self._elapsed: float = 0.0
+
+    def start(self) -> None:
+        self._started_at = self._clock.monotonic()
+
+    def stop(self) -> None:
+        if self._started_at is not None:
+            self._elapsed = self._clock.monotonic() - self._started_at
+            self._started_at = None
+
+    @contextmanager
+    def timed(self) -> Iterator[RunTracker]:
+        """Time a run: `with tracker.timed(): ...` brackets start()/stop() even on error."""
+        self.start()
+        try:
+            yield self
+        finally:
+            self.stop()
+
+    def record(self, phase: Phase, model: str, usage: TokenUsage) -> UsageEvent:
+        """Record one metered LLM call, pricing it via the model pricing table.
+
+        Thread-safe: GEPA evaluates batches concurrently, so metered calls land in parallel.
+        """
+        event = UsageEvent(phase=phase, model=model, usage=usage, cost_usd=cost_usd(model, usage))
+        with self._lock:
+            self._events.append(event)
+        return event
+
+    @property
+    def events(self) -> list[UsageEvent]:
+        return list(self._events)
+
+    def totals(self) -> UsageTotals:
+        total = UsageTotals()
+        for event in self._events:
+            total._add(event)
+        return total
+
+    def by_phase(self) -> dict[Phase, UsageTotals]:
+        buckets: dict[Phase, UsageTotals] = defaultdict(UsageTotals)
+        for event in self._events:
+            buckets[event.phase]._add(event)
+        return dict(buckets)
+
+    def duration_seconds(self) -> float:
+        """Elapsed seconds; live (since start) if still running, else the frozen final span."""
+        if self._started_at is not None:
+            return self._clock.monotonic() - self._started_at
+        return self._elapsed
+
+    def record_summary(self) -> RunRecord:
+        return RunRecord(
+            run_id=self._run_id,
+            kind=self._kind,
+            duration_seconds=self.duration_seconds(),
+            total=self.totals(),
+            by_phase=self.by_phase(),
+        )

world_model_optimizer-0.2.0.dist-info/METADATA

@@ -0,0 +1,203 @@
+Metadata-Version: 2.4
+Name: world-model-optimizer
+Version: 0.2.0
+Summary: Run agents, build world models from traces, and optimize agent harnesses.
+Requires-Python: >=3.12
+Requires-Dist: anthropic>=0.39
+Requires-Dist: boto3>=1.34
+Requires-Dist: click>=8.2
+Requires-Dist: environment-capture
+Requires-Dist: fastapi>=0.128
+Requires-Dist: gepa>=0.1.1
+Requires-Dist: httpx>=0.27
+Requires-Dist: numpy>=1.26
+Requires-Dist: openai>=1.40
+Requires-Dist: posthog>=7.0
+Requires-Dist: pydantic>=2.6
+Requires-Dist: python-multipart>=0.0.9
+Requires-Dist: rich>=14.1
+Requires-Dist: scikit-learn>=1.4
+Requires-Dist: tomli-w>=1.0
+Requires-Dist: typer>=0.16
+Requires-Dist: uvicorn>=0.38
+Provides-Extra: connectors
+Requires-Dist: mcp>=1.27; extra == 'connectors'
+Provides-Extra: dev
+Requires-Dist: e2b==2.31.0; extra == 'dev'
+Requires-Dist: harbor==0.20.0; extra == 'dev'
+Requires-Dist: matplotlib>=3.8; extra == 'dev'
+Requires-Dist: mcp>=1.27; extra == 'dev'
+Requires-Dist: opentelemetry-proto>=1.24; extra == 'dev'
+Requires-Dist: pandas>=2.0; extra == 'dev'
+Requires-Dist: psycopg[binary]>=3.1; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Requires-Dist: ruff>=0.5; extra == 'dev'
+Requires-Dist: seaborn>=0.13; extra == 'dev'
+Requires-Dist: tinker-cookbook<0.5,>=0.4.3; extra == 'dev'
+Requires-Dist: tinker<0.24,>=0.23; extra == 'dev'
+Requires-Dist: ty>=0.0.1a1; extra == 'dev'
+Requires-Dist: wandb>=0.17; extra == 'dev'
+Provides-Extra: distill
+Requires-Dist: tinker-cookbook<0.5,>=0.4.3; extra == 'distill'
+Requires-Dist: tinker<0.24,>=0.23; extra == 'distill'
+Requires-Dist: wandb>=0.17; extra == 'distill'
+Provides-Extra: e2b
+Requires-Dist: e2b==2.31.0; extra == 'e2b'
+Provides-Extra: harbor
+Requires-Dist: harbor==0.20.0; extra == 'harbor'
+Provides-Extra: otel
+Requires-Dist: opentelemetry-proto>=1.24; extra == 'otel'
+Provides-Extra: postgres
+Requires-Dist: psycopg[binary]>=3.1; extra == 'postgres'
+Provides-Extra: viz
+Requires-Dist: matplotlib>=3.8; extra == 'viz'
+Requires-Dist: pandas>=2.0; extra == 'viz'
+Requires-Dist: seaborn>=0.13; extra == 'viz'
+Description-Content-Type: text/markdown
+
+# World Model Optimizer
+
+`wmo` is an open-source project for running and building continuously improving agents. It
+includes a flexible agent runtime, a world model that simulates tool calls, and an optimizer that
+builds task-specific harnesses for stronger performance at lower cost.
+
+![World model, runtime agent, and optimizer connected in a continuous improvement loop](assets/world-model-agent-loop.svg)
+
+<p align="center">
+  🌐 <a href="https://platform.experientiallabs.ai">Platform</a> |
+  📚 <a href="https://github.com/experientiallabs/world-model-optimizer/tree/main/docs">Docs</a> |
+  <a href="https://discord.gg/QwjJpEyHd"><img src="https://cdn.simpleicons.org/discord/5865F2" alt="" width="16" height="16"> Discord</a>
+</p>
+
+## Getting started
+
+### Local setup
+
+Install WMO, choose the model provider for the built-in runtime agent, and start a local run:
+
+```bash
+pip install world-model-optimizer
+wmo providers set
+wmo run --task "Inspect this repository and explain it"
+```
+
+Build a named world model from collected traces:
+
+```bash
+wmo build --file traces.jsonl --name my-environment
+```
+
+Then optimize an agent harness against that model and a set of tasks:
+
+```bash
+wmo optimize harness my-agent my-environment --tasks tasks.jsonl
+```
+
+### Hosted platform
+
+Create an account at [platform.experientiallabs.ai](https://platform.experientiallabs.ai), then
+authenticate the CLI:
+
+```bash
+wmo login
+```
+
+Copy an agent ID from the platform and run its current champion harness:
+
+```bash
+wmo run <agent-id>
+```
+
+### E2B backend
+
+Hosted agents already run in platform-managed E2B sandboxes. To evaluate a local optimization in
+E2B, install the extra and provide an E2B key:
+
+```bash
+pip install "world-model-optimizer[e2b]"
+export E2B_API_KEY=...
+wmo optimize harness my-agent my-environment --tasks tasks.jsonl --backend e2b
+```
+
+## Use a world model as an API
+
+```python
+from wmo import Action, ActionKind
+from wmo.config.store import WorldModelStore
+from wmo.engine.loader import load_world_model
+
+model_dir = WorldModelStore(".wmo").resolve("airline")
+wm, _provider = load_world_model(model_dir)
+
+session = wm.new_session(task="check out the cart")
+obs = wm.step(session.id, Action(kind=ActionKind.TOOL_CALL, name="add_to_cart",
+                                 arguments={"sku": "A1"}))
+print(obs.content)
+```
+
+Or over HTTP (same code path), namespaced by model name: `GET /world_models`, then `POST /world_models/{name}/sessions` and `POST /world_models/{name}/sessions/{id}/step`.
+
+## Run after platform login
+
+After `wmo login`, the same `wmo run` command can open a hosted world model or run an agent's
+current champion harness in E2B. The platform manages model and sandbox credentials, so hosted
+runs do not need local API keys.
+
+```bash
+wmo login
+wmo run <world-model-or-agent-id>
+wmo run <agent-id> -u . --task "fix the failing tests"
+```
+
+Workspace upload is opt-in with `-u`: WMO live-syncs changes and preserves concurrent local edits.
+Long-running agents can detach, continue in the platform, and be messaged or reattached later.
+
+```bash
+wmo run <agent-id> -u . --detach
+wmo run --send "Now run the full test suite"
+wmo run --attach
+wmo run --end
+```
+
+## Runtime agents and optimizers in E2B sandboxes
+
+WMO can run the real [pi](https://github.com/earendil-works/pi) worker inside isolated
+[E2B](https://e2b.dev) sandboxes while the world model supplies the environment. Optimization and
+evaluation rollouts run in parallel, and model credentials stay outside the sandbox.
+
+```bash
+wmo optimize harness my-agent my-environment --tasks tasks.jsonl --backend e2b
+wmo eval tasks.jsonl --mode closed-loop --harness my-agent --harness-backend e2b
+```
+
+The optimizer can change prompts, tools, policies, skills, and runtime code. Every candidate is
+measured against the same simulated tasks, and only changes that pass the evaluation gates become
+the new versioned champion harness.
+
+## Development
+
+Managed with [uv](https://docs.astral.sh/uv/); linting/formatting with [ruff](https://docs.astral.sh/ruff/); type checking with [ty](https://github.com/astral-sh/ty). Conventions live in [AGENTS.md](./AGENTS.md).
+
+```bash
+uv sync --extra dev      # env + dev tools
+uv run ruff check .      # lint
+uv run ruff format .     # format
+uv run ty check          # type check
+uv run pytest -q         # tests
+```
+
+## Usage telemetry
+
+`wmo` uses anonymous usage telemetry to track the volume of usage.
+Telemetry is strictly metadata. It never includes prompts, traces, actions, observations, file paths,
+model names, provider credentials, or raw user content.
+
+Telemetry is enabled by default. To opt out for a project:
+
+```bash
+uv run wmo config telemetry disable
+```
+
+This writes `.wmo/settings.toml`. You can re-enable it with `uv run wmo config telemetry enable`,
+check the current setting with `uv run wmo config telemetry status`, or disable it for a process
+with `DO_NOT_TRACK=1` or `WMO_TELEMETRY=0`.