world-model-optimizer 0.2.0__tar.gz

world_model_optimizer-0.2.0/.gitignore

@@ -0,0 +1,47 @@
+__pycache__/
+*.py[cod]
+.venv/
+.env
+.env.*
+!.env.example
+.wmo/
+dist/
+*.egg-info/
+.pytest_cache/
+.ruff_cache/
+.mypy_cache/
+.ty_cache/
+.DS_Store
+
+# Local-only design notes; not committed.
+DESIGN.md
+
+demos/
+
+# baseline-grid run scratch (not artifacts)
+.wmo-rebuild/
+logs/
+
+.wmo-fixed/
+
+# dev-run artifacts of the RL arms (official rows are committed deliberately under results/)
+examples/tau-bench/rl/*.results.jsonl
+examples/tau-bench/rl/icl_memory.jsonl
+
+# Local per-project settings (telemetry ids etc.) — never committed, wherever the root is.
+settings.toml
+
+# Claude Code local state (skills ARE committed; local settings/locks are not).
+.claude/settings.local.json
+.claude/scheduled_tasks.lock
+.claude/projects/
+
+# Raw experiment data artifacts of the distill/ablation program are NOT committed — they live
+# in the Notion experiments area (Research page) with a SHA-256 manifest. Analysis prose (.md)
+# stays in git, as do small result JSONs elsewhere under research/ (e.g. trace_scaling_results).
+.agents/docs/research/distill/*.json
+.agents/docs/research/distill/*.jsonl
+.agents/docs/research/distill/*.html
+
+# wandb run directories (local mirrors of what is already on the server)
+wandb/

world_model_optimizer-0.2.0/PKG-INFO

@@ -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`.

world_model_optimizer-0.2.0/README.md

@@ -0,0 +1,146 @@
+# 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`.

world_model_optimizer-0.2.0/conftest.py

@@ -0,0 +1,11 @@
+"""Repo-wide pytest configuration."""
+
+from __future__ import annotations
+
+import os
+
+# Rich consoles snapshot color support when constructed, and `wmo.cli.app` builds its console at
+# import time — so color-forcing vars must go before any test module imports it, or a dev shell
+# exporting FORCE_COLOR would inject ANSI codes into CliRunner captures and fail assertions.
+os.environ.pop("FORCE_COLOR", None)
+os.environ.pop("CLICOLOR_FORCE", None)

world_model_optimizer-0.2.0/packages/llm-waterfall/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Experiential Labs
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

world_model_optimizer-0.2.0/packages/llm-waterfall/llm_waterfall/__init__.py

@@ -0,0 +1,53 @@
+"""llm-waterfall: stateless LLM failover across an ordered chain of backends.
+
+Capacity errors (throttling, transient 5xx, timeouts) spill to the next backend; real client
+errors raise immediately. Every call returns which backend served it, token usage, USD cost, and
+the full attempt trail.
+"""
+
+from llm_waterfall.classify import is_capacity_error, outcome_for
+from llm_waterfall.pricing import ModelPrice, cost_usd, price_for
+from llm_waterfall.types import (
+    Attempt,
+    Backend,
+    ChatMaxTokensField,
+    ChatRequest,
+    ChatResponse,
+    ChatResult,
+    CompletionResult,
+    EmbeddingResult,
+    EmbeddingsUnsupported,
+    Message,
+    RetryPolicy,
+    TokenUsage,
+    ToolCallingUnsupported,
+    VerifyResult,
+    WaterfallExhausted,
+)
+from llm_waterfall.waterfall import Waterfall
+
+__version__ = "0.1.4"
+
+__all__ = [
+    "Attempt",
+    "Backend",
+    "ChatMaxTokensField",
+    "ChatRequest",
+    "ChatResponse",
+    "ChatResult",
+    "CompletionResult",
+    "EmbeddingResult",
+    "EmbeddingsUnsupported",
+    "Message",
+    "ModelPrice",
+    "RetryPolicy",
+    "TokenUsage",
+    "ToolCallingUnsupported",
+    "VerifyResult",
+    "Waterfall",
+    "WaterfallExhausted",
+    "cost_usd",
+    "is_capacity_error",
+    "outcome_for",
+    "price_for",
+]

world_model_optimizer-0.2.0/packages/llm-waterfall/llm_waterfall/adapters/__init__.py

@@ -0,0 +1,36 @@
+"""Adapter registry: map a Backend's provider name to its adapter class.
+
+Adapter modules import at module scope (they gate their SDK imports internally, so this package
+imports with zero SDKs installed); only the SDKs themselves are lazy.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from llm_waterfall.adapters.anthropic import AnthropicAdapter
+from llm_waterfall.adapters.aws_mantle import AwsMantleAdapter
+from llm_waterfall.adapters.azure_openai import AzureOpenAIAdapter
+from llm_waterfall.adapters.base import Adapter
+from llm_waterfall.adapters.bedrock import BedrockAdapter
+from llm_waterfall.adapters.openai import OpenAIAdapter
+from llm_waterfall.types import PROVIDERS, Backend
+
+_ADAPTERS: dict[str, Callable[[Backend], Adapter]] = {
+    "bedrock": BedrockAdapter,
+    "openai": OpenAIAdapter,
+    "anthropic": AnthropicAdapter,
+    "azure_openai": AzureOpenAIAdapter,  # real; construction validates endpoint/api_version
+    "aws_mantle": AwsMantleAdapter,  # stub: raises NotImplementedError at construction
+}
+
+
+def build_adapter(backend: Backend) -> Adapter:
+    """Construct the adapter for `backend`. Unimplemented providers fail here, fast."""
+    try:
+        factory = _ADAPTERS[backend.provider]
+    except KeyError:
+        raise ValueError(
+            f"unknown provider {backend.provider!r}; expected one of {', '.join(PROVIDERS)}"
+        ) from None
+    return factory(backend)

world_model_optimizer-0.2.0/packages/llm-waterfall/llm_waterfall/adapters/anthropic.py

@@ -0,0 +1,105 @@
+"""Anthropic direct-API adapter. Reads ANTHROPIC_API_KEY from the environment."""
+
+from __future__ import annotations
+
+import threading
+from typing import TYPE_CHECKING, cast
+
+from llm_waterfall.adapters.base import missing_sdk_error
+from llm_waterfall.types import (
+    Backend,
+    ChatRequest,
+    ChatResponse,
+    EmbeddingsUnsupported,
+    Message,
+    TokenUsage,
+    ToolCallingUnsupported,
+)
+
+if TYPE_CHECKING:
+    from anthropic import Anthropic
+    from anthropic.types import MessageParam
+
+
+class AnthropicAdapter:
+    """Claude via the direct Anthropic Messages API."""
+
+    def __init__(self, backend: Backend) -> None:
+        self.backend = backend
+        self._client: Anthropic | None = None
+        self._lock = threading.Lock()
+
+    def _get_client(self) -> Anthropic:
+        if self._client is None:
+            with self._lock:
+                if self._client is None:
+                    try:
+                        import httpx
+                        from anthropic import Anthropic
+                    except ModuleNotFoundError as exc:
+                        raise missing_sdk_error("anthropic", "anthropic") from exc
+
+                    # max_retries=0: the waterfall owns retry policy. Granular httpx.Timeout so
+                    # a dead endpoint fails over after connect_timeout_s, not read_timeout_s.
+                    self._client = Anthropic(
+                        max_retries=0,
+                        timeout=httpx.Timeout(
+                            self.backend.read_timeout_s,
+                            connect=self.backend.connect_timeout_s,
+                        ),
+                    )
+        return self._client
+
+    def complete(
+        self,
+        system: str,
+        messages: list[Message],
+        *,
+        temperature: float | None,
+        max_tokens: int,
+    ) -> tuple[str, TokenUsage]:
+        """One Messages API call; system is a top-level arg (Anthropic-native)."""
+        api_messages = [
+            cast("MessageParam", {"role": m.role, "content": m.content}) for m in messages
+        ]
+        client_messages = self._get_client().messages
+        # Claude 4.7+ rejects sampling params; only forward temperature when explicitly set.
+        if temperature is None:
+            response = client_messages.create(
+                model=self.backend.model,
+                system=system,
+                messages=api_messages,
+                max_tokens=max_tokens,
+            )
+        else:
+            response = client_messages.create(
+                model=self.backend.model,
+                system=system,
+                messages=api_messages,
+                max_tokens=max_tokens,
+                temperature=temperature,
+            )
+        text = "".join(block.text for block in response.content if block.type == "text")
+        usage = TokenUsage(
+            input_tokens=response.usage.input_tokens,
+            output_tokens=response.usage.output_tokens,
+        )
+        return text, usage
+
+    def complete_chat(self, request: ChatRequest) -> ChatResponse:
+        """Direct Anthropic structured mapping is not implemented yet."""
+        del request
+        raise ToolCallingUnsupported(
+            "the anthropic adapter does not yet map OpenAI-compatible tool calls; "
+            "put an openai, azure_openai, or bedrock backend in the chain"
+        )
+
+    def embed_model_id(self) -> str | None:
+        """Anthropic has no embeddings API."""
+        return None
+
+    def embed(self, texts: list[str]) -> tuple[list[list[float]], TokenUsage]:
+        raise EmbeddingsUnsupported(
+            "Anthropic has no embeddings API; put a bedrock or openai backend in the chain "
+            "for embeddings (the waterfall skips this backend for embed calls)."
+        )