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

wmo/cli/route_app.py

@@ -0,0 +1,274 @@
+"""`wmo optimize route`: fit, tune, and report learned inference policies from outcome matrices.
+
+The routing optimizer's CLI face, sitting beside `wmo optimize harness` in the optimizer
+family. Consumes a persisted `OutcomeMatrix` (produced by closed-loop pool evaluation or a
+research adapter such as RouterBench) and emits the policy artifact serving loads, plus the
+improvement report the endpoint cites. `tune` is the one post-fit control: it moves a fitted
+policy's cost/quality dial without refitting. Vocabulary note: "route" is developer-facing CLI
+only; customer copy never says router.
+"""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from pathlib import Path
+
+import typer
+from rich.console import Console
+
+from wmo.optimize.knn import (
+    COST_QUALITY_ANCHORS,
+    apply_cost_quality,
+    cost_quality_knobs,
+    cost_quality_named_point,
+    fit_knn_policy,
+)
+from wmo.optimize.outcomes import OutcomeMatrix
+from wmo.optimize.policy import (
+    KNN_BANK_FILENAME,
+    POLICY_FILENAME,
+    EmbedderSpec,
+    RoutingPolicy,
+)
+from wmo.optimize.report import build_report
+from wmo.optimize.routing import evaluate_policy, fit_rank_policy, rerank_policy
+
+route_app = typer.Typer(
+    help="Fit and report learned inference policies from closed-loop outcome matrices.",
+    no_args_is_help=True,
+)
+
+_console = Console()
+
+
+@route_app.command("fit")
+def fit(
+    matrix_file: str = typer.Argument(..., help="OutcomeMatrix JSON (closed-loop eval output)."),
+    kind: str = typer.Option(
+        "rank",
+        "--kind",
+        help="knn (guarded nearest-neighbor evidence, the validated champion) | rank "
+        "(Avengers cluster ranks).",
+    ),
+    out: str = typer.Option(
+        POLICY_FILENAME, "--out", help="Where to write the fitted policy JSON."
+    ),
+    fallback: str = typer.Option(
+        None,
+        "--fallback",
+        help="(knn) Baseline model every request uses unless the evidence says otherwise. "
+        "Default: the best single model on the fit set.",
+    ),
+    z: float = typer.Option(
+        0.5,
+        "--z",
+        min=0.0,
+        help="(knn) Confidence knob: standard errors of paired evidence a pick must clear to "
+        "leave the fallback (doubled when it is also pricier). Higher = stricter = more "
+        "requests stay on the fallback; 0 routes on any positive difference.",
+    ),
+    rag_num: int = typer.Option(50, "--rag-num", min=1, help="(knn) Neighbor budget."),
+    rag_thres: float = typer.Option(
+        0.95,
+        "--rag-thres",
+        min=0.0,
+        max=1.0,
+        help="(knn) Keep neighbors above this fraction of the rag-num-th best similarity.",
+    ),
+    min_pairs: int = typer.Option(
+        8, "--min-pairs", min=0, help="(knn) Neighbors scored on both sides before routing away."
+    ),
+    floor_q: float = typer.Option(
+        0.05,
+        "--floor-q",
+        min=0.0,
+        max=1.0,
+        help="Novelty floor quantile: abstain to the fallback when a query's best bank "
+        "similarity is below this quantile of the bank's own nearest-neighbor sims "
+        "(coverage/robustness knob for task drift; 0 = off, the exact validated champion).",
+    ),
+    se_floor: bool = typer.Option(
+        True,
+        "--se-floor/--no-se-floor",
+        help="(knn) Floor the guard's standard error on thin neighborhoods (small-bank safety).",
+    ),
+    clusters: int = typer.Option(64, "--clusters", min=1, help="k-means cluster count."),
+    seed: int = typer.Option(42, "--seed", help="Clustering seed."),
+    top_k_clusters: int = typer.Option(2, "--top-k-clusters", min=1),
+    beta: float = typer.Option(6.0, "--beta", help="Cluster softmax sharpness."),
+    cost_weight: float = typer.Option(
+        0.0,
+        "--cost-weight",
+        min=0.0,
+        help="Quality/cost knob: reward points paid per average-call-cost unit (0 = pure "
+        "accuracy ranking, the Avengers reference behavior).",
+    ),
+    embedder: str = typer.Option("hashing", "--embedder", help="hashing | azure"),
+    dim: int = typer.Option(512, "--dim", help="Embedding dimension."),
+    deployment: str = typer.Option(None, "--deployment", help="(azure) embedding deployment."),
+    endpoint: str = typer.Option(None, "--endpoint", help="(azure) resource endpoint."),
+    api_key_env: str = typer.Option(
+        None, "--api-key-env", help="(azure) env var holding the account key."
+    ),
+) -> None:
+    """Fit a routing policy on an outcome matrix (kNN evidence or Avengers cluster ranks)."""
+    if kind not in ("rank", "knn"):
+        raise typer.BadParameter(f"unknown kind '{kind}'; use knn or rank")
+    matrix = OutcomeMatrix.load(Path(matrix_file))
+    if embedder not in ("hashing", "azure"):
+        raise typer.BadParameter(f"unknown embedder '{embedder}'; use hashing or azure")
+    spec = (
+        EmbedderSpec(dim=dim)
+        if embedder == "hashing"
+        else EmbedderSpec(
+            kind="azure",
+            dim=dim,
+            deployment=deployment,
+            endpoint=endpoint,
+            api_key_env=api_key_env,
+        )
+    )
+    out_path = Path(out)
+    if rag_thres <= 0.0:
+        # typer's min is inclusive but the artifact field requires > 0; fail before the fit
+        # writes a sidecar it will then abandon.
+        raise typer.BadParameter("--rag-thres must be greater than 0")
+    built = spec.build()  # ONE embedder for fit and evaluation; azure would otherwise embed twice
+    if kind == "knn":
+        if cost_weight > 0.0:
+            raise typer.BadParameter(
+                "--cost-weight re-ranks cluster evidence and applies to --kind rank only; a knn "
+                "policy trades cost through its dial instead: fit it, then "
+                "`wmo optimize route tune <policy.json> --cost-quality <0..1>`"
+            )
+        # The sidecar goes beside the policy file: that is where serving resolves it from.
+        policy = fit_knn_policy(
+            matrix,
+            bank_path=out_path.parent / KNN_BANK_FILENAME,
+            embedder=spec,
+            embed_with=built,
+            guard_model=fallback,
+            rag_num=rag_num,
+            rag_thres=rag_thres,
+            z=z,
+            min_pairs=min_pairs,
+            se_floor=se_floor,
+            floor_q=floor_q,
+            fitted_from=f"{matrix_file} knn z={z} k={rag_num} q={floor_q} {embedder}-{dim}",
+        )
+    else:
+        policy = fit_rank_policy(
+            matrix,
+            embedder=spec,
+            n_clusters=clusters,
+            seed=seed,
+            top_k_clusters=top_k_clusters,
+            beta=beta,
+            fitted_from=f"{matrix_file} seed={seed} k={clusters} {embedder}-{dim}",
+        )
+        if cost_weight > 0.0:
+            policy = rerank_policy(policy, cost_weight=cost_weight)
+    policy.save(out_path)
+    result = evaluate_policy(policy, matrix, matrix.scenario_ids(), embedder=built)
+    if kind == "knn":
+        routed = 1.0 - result.model_mix.get(policy.default_model, 0.0)
+        _console.print(
+            f"[green]✓[/green] fitted knn policy over {result.scenarios} scenarios -> {out}\n"
+            f"  bank {out_path.parent / KNN_BANK_FILENAME}, fallback {policy.default_model}, "
+            f"z={z}\n"
+            f"  routed away from the fallback {routed:.1%} of the time; cost/scenario "
+            f"${result.cost_per_scenario:.5f}\n"
+            f"  fit-set accuracy {result.accuracy:.4f} is IN-SAMPLE (every request retrieves its "
+            "own row); measure on held-out scenarios with `wmo optimize route report`"
+        )
+        return
+    _console.print(
+        f"[green]✓[/green] fitted {len(policy.clusters)} clusters over "
+        f"{result.scenarios} scenarios -> {out}\n"
+        f"  fit-set accuracy {result.accuracy:.4f}, cost/scenario ${result.cost_per_scenario:.5f}"
+    )
+
+
+@route_app.command("tune")
+def tune(
+    policy_file: str = typer.Argument(POLICY_FILENAME, help="Fitted knn policy JSON to re-tune."),
+    cost_quality: float = typer.Option(
+        ...,
+        "--cost-quality",
+        min=0.0,
+        max=1.0,
+        help="The endpoint's one dial: 0.0 = max quality, 1.0 = max savings. 0.25 is the "
+        "shipped default. See the anchor table this command prints for what each end measured.",
+    ),
+) -> None:
+    """Set a fitted policy's cost/quality dial in place, without refitting anything.
+
+    The dial maps to the policy's knobs along the measured frontier (see
+    `wmo.optimize.knn.apply_cost_quality`). The first run copies the un-tuned artifact to
+    `policy.base.json` and every later run re-reads THAT, so the dial is always applied to the
+    policy as fitted and sliding twice never compounds:
+
+        wmo optimize route tune models/support/policy.json --cost-quality 0.6
+
+    The evidence bank is untouched, so this is instant. A served endpoint can be dialed without
+    touching files at all: `PUT /v1/endpoints/{name}/config`.
+    """
+    path = Path(policy_file)
+    if not path.is_file():
+        raise typer.BadParameter(f"no policy file at {path}")
+    base_path = path.with_name(f"{path.stem}.base{path.suffix}")
+    if not base_path.is_file():
+        # Preserve the artifact as fitted the first time, so `tune` is always re-appliable from
+        # the fit and never from an already-slid copy of itself.
+        base_path.write_bytes(path.read_bytes())
+    base = RoutingPolicy.load(base_path)
+    try:
+        tuned = apply_cost_quality(base, cost_quality)
+    except ValueError as exc:
+        raise typer.BadParameter(str(exc)) from exc
+    tuned.save(path)
+    knobs = cost_quality_knobs(cost_quality)
+    _console.print(
+        f"[green]✓[/green] cost_quality={cost_quality:g} "
+        f"({cost_quality_named_point(cost_quality)}) -> {path}\n"
+        f"  knobs: floor_q={knobs.floor_q:g}, cost knob lam={knobs.pick_lam:g}, "
+        f"guard={knobs.guard_mode}, z={knobs.knn_z:g}\n"
+        f"  as fitted: {base_path}\n"
+        f"  measured on routerbench-ours9 (5 held-out splits, vs the best single model):"
+    )
+    for anchor in COST_QUALITY_ANCHORS:
+        marker = "->" if anchor.cost_quality == cost_quality else "  "
+        _console.print(
+            f"  {marker} {anchor.cost_quality:<5g} {anchor.quality_delta_points:+.2f}pt "
+            f"@ {anchor.cost_delta_percent:+.1f}% cost"
+            + (f"  [dim]{anchor.named_point}[/dim]" if anchor.named_point != "Custom" else "")
+        )
+
+
+@route_app.command("report")
+def report(
+    matrix_file: str = typer.Argument(..., help="OutcomeMatrix JSON with held-out scenarios."),
+    policy_file: str = typer.Argument(..., help="Fitted policy JSON."),
+    baseline: str = typer.Option(
+        ..., "--baseline", help="Frontier pool model the report compares against."
+    ),
+    endpoint: str = typer.Option("endpoint", "--endpoint", help="Endpoint id for the report."),
+    out: str = typer.Option("report.json", "--out", help="Where to write the report JSON."),
+) -> None:
+    """Build the improvement report for a fitted policy over a matrix."""
+    matrix = OutcomeMatrix.load(Path(matrix_file))
+    policy = RoutingPolicy.load(Path(policy_file))
+    improvement = build_report(
+        matrix,
+        policy,
+        baseline=baseline,
+        endpoint=endpoint,
+        generated_at=datetime.now(tz=UTC).isoformat(),
+    )
+    Path(out).write_text(improvement.model_dump_json(indent=2), encoding="utf-8")
+    headline = improvement.headline
+    _console.print(
+        f"[green]✓[/green] report -> {out}\n"
+        f"  routed acc {headline.accuracy:.4f} @ ${headline.cost_per_run_usd:.5f}/run vs "
+        f"{baseline} {headline.baseline_accuracy:.4f} @ ${headline.baseline_cost_per_run_usd:.5f}"
+    )

wmo/cli/session_state.py

@@ -0,0 +1,243 @@
+# Copyright (c) 2026 Experiential Labs. All rights reserved.
+
+"""Persisted local references to detached hosted agent sessions.
+
+``wmo run <agent-id> --detach`` leaves a platform-owned E2B session running
+with no local process attached. The CLI remembers how to address it again
+(platform URL, agent id, session id) plus the workspace-sync checkpoint in the
+user-global WMO state directory (``$WMO_HOME`` or ``~/.wmo``), never inside
+the directory being synchronized. Writes are atomic and owner-only; the state
+directory is injectable so tests never touch real user state.
+
+The workspace base archive is stored content-addressed next to the JSON state
+and referenced by digest: the archive lands first and the state referencing it
+lands last, so a crash between the two writes leaves the previous checkpoint
+intact instead of a dangling pointer.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import hashlib
+import os
+import re
+import tempfile
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, ValidationError
+
+from wmo.platform.credentials import wmo_home
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+SESSIONS_DIRNAME = "sessions"
+_CURRENT_FILENAME = "current"
+# Session ids come from the platform (UUIDs); refuse anything that could
+# escape the state directory when used as a file-name component.
+_SESSION_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
+
+
+class SessionStateError(RuntimeError):
+    """Detached-session state on disk is missing, unsafe, or corrupted."""
+
+
+class WorkspaceCheckpoint(BaseModel):
+    """Synchronization checkpoint between one local directory and the session."""
+
+    root: str
+    base_sha256: str = ""
+    conflicts: tuple[str, ...] = ()
+    # A patch applied locally whose acknowledgement did not reach the platform
+    # yet; the next command retries it so the server object never leaks.
+    pending_ack: str | None = None
+
+
+class DetachedSessionState(BaseModel):
+    """One hosted agent session the CLI can send to, attach to, or end later."""
+
+    api_url: str
+    web_url: str | None = None
+    agent_id: str
+    agent_name: str
+    session_id: str
+    created_at: str
+    cursor: int = 0
+    workspace: WorkspaceCheckpoint | None = None
+
+
+class SessionStateStore:
+    """Atomic, owner-only persistence for detached session references."""
+
+    def __init__(self, directory: Path | None = None) -> None:
+        """Store state under ``directory`` (default: the user-global WMO home)."""
+        self._directory = directory if directory is not None else wmo_home() / SESSIONS_DIRNAME
+
+    # -- session records ---------------------------------------------------------------------
+
+    def save(
+        self, state: DetachedSessionState, *, base_archive: bytes | None = None
+    ) -> DetachedSessionState:
+        """Persist one session's state, optionally with a new workspace base archive.
+
+        Returns:
+            The state as persisted (with ``base_sha256`` updated when an
+            archive was written).
+
+        Raises:
+            SessionStateError: If an archive is given without workspace
+                metadata, the id is unsafe, or a state path is a symlink.
+        """
+        self._ensure_directory()
+        session_id = self._validated(state.session_id)
+        if base_archive is not None:
+            if state.workspace is None:
+                msg = "a workspace base archive requires workspace checkpoint metadata"
+                raise SessionStateError(msg)
+            digest = hashlib.sha256(base_archive, usedforsecurity=False).hexdigest()
+            self._write_bytes(self._archive_path(session_id, digest), base_archive)
+            state = state.model_copy(
+                update={"workspace": state.workspace.model_copy(update={"base_sha256": digest})}
+            )
+        payload = state.model_dump_json(indent=2)
+        self._write_bytes(self._state_path(session_id), payload.encode("utf-8"))
+        self._prune_archives(state)
+        return state
+
+    def load(self, session_id: str) -> DetachedSessionState | None:
+        """Read one session's state; ``None`` when nothing is stored for the id."""
+        path = self._state_path(self._validated(session_id))
+        if not path.exists():
+            return None
+        try:
+            return DetachedSessionState.model_validate_json(path.read_text(encoding="utf-8"))
+        except (ValidationError, ValueError, OSError) as error:
+            msg = (
+                f"session state at {path} is unreadable; delete it and start a new "
+                "session with `wmo run <agent-id> --detach`"
+            )
+            raise SessionStateError(msg) from error
+
+    def load_base_archive(self, state: DetachedSessionState) -> bytes:
+        """Read and integrity-check the persisted workspace base archive."""
+        workspace = state.workspace
+        if workspace is None or not workspace.base_sha256:
+            msg = f"session {state.session_id} has no workspace checkpoint"
+            raise SessionStateError(msg)
+        path = self._archive_path(self._validated(state.session_id), workspace.base_sha256)
+        try:
+            content = path.read_bytes()
+        except OSError as error:
+            msg = (
+                f"workspace checkpoint archive is missing at {path}; end the session "
+                f"with `wmo run --session {state.session_id} --end` and recover from "
+                "the final workspace download"
+            )
+            raise SessionStateError(msg) from error
+        digest = hashlib.sha256(content, usedforsecurity=False).hexdigest()
+        if digest != workspace.base_sha256:
+            msg = (
+                f"workspace checkpoint archive at {path} failed its integrity check; "
+                f"delete it and end the session with `wmo run --session "
+                f"{state.session_id} --end`"
+            )
+            raise SessionStateError(msg)
+        return content
+
+    def write_recovery_archive(self, session_id: str, content: bytes) -> Path:
+        """Save a final workspace archive that could not be synchronized locally.
+
+        The file deliberately survives :meth:`delete`: it is the user's data,
+        not session state.
+        """
+        self._ensure_directory()
+        path = self._directory / f"{self._validated(session_id)}.recovered.tar.gz"
+        self._write_bytes(path, content)
+        return path
+
+    def delete(self, session_id: str) -> None:
+        """Remove the state, its checkpoint archives, and a matching current pointer."""
+        session_id = self._validated(session_id)
+        self._state_path(session_id).unlink(missing_ok=True)
+        for archive in self._directory.glob(f"{session_id}.workspace-*.tar.gz"):
+            archive.unlink(missing_ok=True)
+        if self.current_session_id() == session_id:
+            (self._directory / _CURRENT_FILENAME).unlink(missing_ok=True)
+
+    # -- current pointer ---------------------------------------------------------------------
+
+    def set_current(self, session_id: str) -> None:
+        """Publish ``session_id`` as the session bare send/attach/end commands use."""
+        self._ensure_directory()
+        content = f"{self._validated(session_id)}\n".encode()
+        self._write_bytes(self._directory / _CURRENT_FILENAME, content)
+
+    def current_session_id(self) -> str | None:
+        """The current session id, or ``None`` when no pointer is set."""
+        path = self._directory / _CURRENT_FILENAME
+        if not path.exists():
+            return None
+        value = path.read_text(encoding="utf-8").strip()
+        return value or None
+
+    # -- internals ---------------------------------------------------------------------------
+
+    def _ensure_directory(self) -> None:
+        """Create every missing directory level owner-only from the start.
+
+        ``mkdir(parents=True)`` would create intermediate levels with the
+        umask default, leaving a window (and, for ``~/.wmo``, a permanent
+        0755) around checkpoint archives that contain the user's source
+        tree. A umask can only narrow 0o700, so creating each level with
+        that mode closes the race; pre-existing directories are left alone.
+        """
+        missing: list[Path] = []
+        current = self._directory
+        while not current.exists():
+            missing.append(current)
+            parent = current.parent
+            if parent == current:
+                break
+            current = parent
+        for directory in reversed(missing):
+            with contextlib.suppress(FileExistsError):
+                directory.mkdir(mode=0o700)
+
+    def _state_path(self, session_id: str) -> Path:
+        return self._directory / f"{session_id}.json"
+
+    def _archive_path(self, session_id: str, digest: str) -> Path:
+        return self._directory / f"{session_id}.workspace-{digest[:16]}.tar.gz"
+
+    def _prune_archives(self, state: DetachedSessionState) -> None:
+        """Drop archives the just-written state no longer references."""
+        keep: str | None = None
+        if state.workspace is not None and state.workspace.base_sha256:
+            keep = self._archive_path(state.session_id, state.workspace.base_sha256).name
+        for candidate in self._directory.glob(f"{state.session_id}.workspace-*.tar.gz"):
+            if candidate.name != keep:
+                candidate.unlink(missing_ok=True)
+
+    def _validated(self, session_id: str) -> str:
+        """Reject ids that are unsafe as file-name components."""
+        if _SESSION_ID.fullmatch(session_id) is None:
+            msg = f"invalid session id: {session_id!r}"
+            raise SessionStateError(msg)
+        return session_id
+
+    def _write_bytes(self, path: Path, content: bytes) -> None:
+        """Write through a 0600 temporary file and swap into place atomically."""
+        if path.is_symlink():
+            msg = f"refusing to write session state through the symlink {path}; remove the link"
+            raise SessionStateError(msg)
+        fd, tmp_name = tempfile.mkstemp(dir=self._directory, prefix=f"{path.name}.")
+        try:
+            with os.fdopen(fd, "wb") as handle:
+                handle.write(content)
+            os.replace(tmp_name, path)
+        except BaseException:
+            # The replace may already have consumed the temp file; a missing
+            # file must not mask the original exception (e.g. an interrupt).
+            with contextlib.suppress(OSError):
+                os.unlink(tmp_name)
+            raise