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

wmo/cli/ingest_cmd.py

@@ -0,0 +1,165 @@
+"""`wmo ingest`: normalize traces from any source into OTel-GenAI JSONL, with live progress.
+
+The standalone half of what `wmo build` does first: pick a source (auto-detected for files),
+normalize to the harness's OTel span JSONL, and report progress. The output feeds any downstream
+command (`wmo build --source otel-genai --file <out>`), and the same event stream drives the
+platform's SSE trace-ingest endpoint (`wmo.ingest.stream.ingest_events`), so the CLI and the
+hosted flow can't drift.
+
+Output modes: a rich progress bar on a TTY, and `--json` for one D-INGEST event object per line
+(machine-readable; what a wrapping process pipes to a wire).
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterator
+from pathlib import Path
+
+import typer
+from rich.console import Console
+from rich.progress import BarColumn, Progress, TextColumn
+
+from wmo.ingest import VendorPull, list_adapters
+from wmo.ingest.stream import (
+    DetectedEvent,
+    DoneEvent,
+    ErrorEvent,
+    IngestEvent,
+    ProgressEvent,
+    event_json,
+    ingest_events,
+)
+
+_console = Console()
+
+
+def _default_out(file: str | None, source: str | None) -> Path:
+    if file is not None:
+        return Path(f"{Path(file).stem}.otel.jsonl")
+    return Path(f"{source or 'traces'}.otel.jsonl")
+
+
+def _render_rich(stream: Iterator[IngestEvent]) -> bool:
+    """Consume the event stream behind a progress bar; returns True on success."""
+    ok = False
+    with Progress(
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        TextColumn("{task.completed}/{task.total} traces"),
+        console=_console,
+    ) as progress:
+        task_id = None
+        for event in stream:
+            if isinstance(event, DetectedEvent):
+                _console.print(f"detected [bold]{event.format}[/bold] · {event.traces} traces")
+                task_id = progress.add_task("normalizing", total=event.traces)
+            elif isinstance(event, ProgressEvent) and task_id is not None:
+                progress.update(task_id, completed=event.normalized)
+                if event.note:
+                    progress.update(task_id, description=event.note)
+            elif isinstance(event, DoneEvent):
+                ok = True
+                progress.stop()
+                _console.print(
+                    f"[green]done[/green] · {event.traces} traces / {event.steps} steps "
+                    f"→ [bold]{event.otel_object}[/bold]"
+                )
+                _console.print(
+                    f"build from it: wmo build --name <name> --source otel-genai "
+                    f"--file {event.otel_object}"
+                )
+            elif isinstance(event, ErrorEvent):
+                progress.stop()
+                code = f" [{event.code}]" if event.code else ""
+                _console.print(f"[red]error{code}[/red] {event.message}")
+    return ok
+
+
+def ingest(
+    file: str = typer.Option(
+        None, "--file", help="Path to an exported traces file (format auto-detected)."
+    ),
+    source: str = typer.Option(
+        None,
+        "--source",
+        help="Trace source adapter (default: auto-detect from the file). One of: "
+        "otel-genai, chat-json, braintrust, phoenix, langfuse, langsmith, posthog, mastra, "
+        "postgres.",
+    ),
+    pull: bool = typer.Option(
+        False, "--pull", help="Pull traces live from the source's API (instead of --file)."
+    ),
+    project: str = typer.Option(None, "--project", help="Vendor project/workspace id (--pull)."),
+    api_key: str = typer.Option(None, "--api-key", help="Vendor API key (else env var)."),
+    since: str = typer.Option(None, "--since", help="Only pull traces since this ISO timestamp."),
+    limit: int = typer.Option(None, "--limit", help="Max number of traces to ingest."),
+    dsn: str = typer.Option(
+        None, "--dsn", help="Postgres connection string (implies --source postgres)."
+    ),
+    table: str = typer.Option(None, "--table", help="Postgres table holding the trace rows."),
+    trace_id_column: str = typer.Option(
+        None, "--trace-id-column", help="Postgres column grouping rows into traces."
+    ),
+    payload_column: str = typer.Option(
+        None, "--payload-column", help="Postgres JSON column holding the trace payloads."
+    ),
+    order_column: str = typer.Option(
+        None, "--order-column", help="Postgres column ordering rows (default: created_at)."
+    ),
+    out: str = typer.Option(
+        None, "--out", help="Output OTel JSONL path (default: <input>.otel.jsonl)."
+    ),
+    json_events: bool = typer.Option(
+        False, "--json", help="Emit one JSON progress event per line instead of the rich UI."
+    ),
+) -> None:
+    """Normalize traces from a file, vendor API, or Postgres table into OTel JSONL.
+
+    No model is built: the output corpus feeds `wmo build --source otel-genai --file <out>` (or
+    any other command that reads a trace file). Progress events follow the D-INGEST vocabulary:
+    detected -> progress... -> done | error.
+    """
+    if (dsn is not None or table is not None) and source is None:
+        source = "postgres"
+    is_pull = pull or dsn is not None or table is not None
+    if file is None and not is_pull:
+        raise typer.BadParameter(
+            "provide --file <export>, --pull (with --source), or --dsn/--table for postgres"
+        )
+    if file is not None and is_pull:
+        raise typer.BadParameter("pass either --file or a pull (--pull/--dsn/--table), not both")
+    if source is not None and source not in list_adapters():
+        raise typer.BadParameter(
+            f"unknown --source {source!r}; choose one of: {', '.join(list_adapters())}"
+        )
+    out_path = Path(out) if out else _default_out(file, source)
+    vendor_pull = (
+        VendorPull(
+            api_key=api_key,
+            project=project,
+            since=since,
+            limit=limit,
+            dsn=dsn,
+            table=table,
+            trace_id_column=trace_id_column,
+            payload_column=payload_column,
+            order_column=order_column,
+        )
+        if is_pull
+        else None
+    )
+    stream = ingest_events(file=file, pull=vendor_pull, source=source, out=out_path, limit=limit)
+
+    if json_events:
+        ok = False
+        for event in stream:
+            # soft_wrap: machine-readable lines must never be broken at terminal width.
+            _console.print(
+                json.dumps(event_json(event)), markup=False, highlight=False, soft_wrap=True
+            )
+            ok = isinstance(event, DoneEvent) or (ok and not isinstance(event, ErrorEvent))
+    else:
+        ok = _render_rich(stream)
+    if not ok:
+        raise typer.Exit(code=1)

wmo/cli/model_roles.py

@@ -0,0 +1,82 @@
+"""Provider resolution for opt-in model roles shared by CLI workflows."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import typer
+
+from wmo.config.settings import ModelRole, load_settings
+from wmo.providers.base import Provider, ProviderConfig, ProviderKind
+from wmo.providers.registry import get_provider
+
+OptInModelRole = Literal["agent", "meta"]
+
+# Azure OpenAI chat completions need an API version on every call. When an opt-in role does not
+# pin one in settings, this shared default API version applies.
+_DEFAULT_AZURE_API_VERSION = "2024-05-01-preview"
+
+
+def resolve_opt_in_model_provider(
+    root: str,
+    role: OptInModelRole,
+    fallback: Provider,
+) -> tuple[Provider, str | None]:
+    """Resolve one opt-in model role, or return the caller's fallback provider.
+
+    Args:
+        root: Project artifact root containing ``settings.toml``.
+        role: The opt-in role to resolve.
+        fallback: Provider retained when the role is not configured.
+
+    Returns:
+        The resolved provider and configured model name, or the fallback and ``None``.
+
+    Raises:
+        typer.BadParameter: The configured provider kind is unknown.
+    """
+    configured = load_settings(root).models.resolve(role)
+    if configured is None:
+        return fallback, None
+    config = _model_config(configured, role=role)
+    return get_provider(config), configured.model
+
+
+def resolve_required_model_config(root: str, role: OptInModelRole) -> ProviderConfig:
+    """Resolve one opt-in role a workflow requires (no fallback provider exists for it).
+
+    The harbor optimize flow has no world model whose provider could stand in, so its
+    ``agent`` (worker) and ``meta`` (proposer) roles must be configured explicitly.
+    """
+    configured = load_settings(root).models.resolve(role)
+    if configured is None:
+        raise typer.BadParameter(
+            f"settings [models.{role}] must be configured in <root>/settings.toml for this "
+            f"workflow; add a [models.{role}] table with provider and model"
+        )
+    return _model_config(configured, role=role)
+
+
+def _model_config(configured: ModelRole, *, role: OptInModelRole) -> ProviderConfig:
+    """Turn one configured role into provider-neutral config with the Azure default."""
+    try:
+        kind = ProviderKind(configured.provider)
+    except ValueError:
+        kinds = ", ".join(kind.value for kind in ProviderKind)
+        raise typer.BadParameter(
+            f"settings [models.{role}] has unknown provider {configured.provider!r}; "
+            f"choose one of: {kinds}"
+        ) from None
+    api_version = configured.api_version
+    if api_version is None and kind is ProviderKind.AZURE_OPENAI:
+        api_version = _DEFAULT_AZURE_API_VERSION
+    return ProviderConfig(
+        kind=kind,
+        model=configured.model,
+        model_type=configured.model_type,
+        region=configured.region,
+        endpoint=configured.endpoint,
+        deployment=configured.deployment,
+        api_version=api_version,
+        reasoning_effort=configured.reasoning_effort,
+    )

wmo/cli/platform_cmds.py

@@ -0,0 +1,372 @@
+"""Platform commands: `wmo login`, `wmo logout`, `wmo status`, `wmo push`, `wmo pull`.
+
+`login` connects this machine to a platform account (browser flow by default,
+`--token` for headless); `push`/`pull` round-trip world models and harnesses
+against the platform registry, auto-detecting the artifact kind from what
+exists locally (or remotely, for pulls).
+"""
+
+from __future__ import annotations
+
+import socket
+import tempfile
+import webbrowser
+from pathlib import Path
+from typing import Annotated
+
+import typer
+from rich.console import Console
+from rich.table import Table
+
+from wmo.config.store import WorldModelStore
+from wmo.harness.doc import HarnessDoc
+from wmo.harness.store import CHAMPION_ALIAS, HarnessStore
+from wmo.platform.auth import BrowserLogin
+from wmo.platform.client import PlatformClient, PlatformError, WhoAmI, fetch_cli_config
+from wmo.platform.credentials import (
+    DEFAULT_WEB_URL,
+    PlatformCredentials,
+    clear_credentials,
+    credentials_path,
+    load_credentials,
+    save_credentials,
+)
+from wmo.platform.transfer import extract_push_meta, pack_model_dir, unpack_model_bundle
+
+_console = Console()
+_CHECK = "[green]✓[/green]"
+
+# Module-level singletons: typer.Option calls can't be defaults inline (ruff B008).
+# Annotated-style options carry no default here; the parameter's own `=` does.
+_LOGIN_URL = typer.Option(
+    "--url", help="Platform URL (defaults to the saved one, then the hosted platform)."
+)
+_LOGIN_API_URL = typer.Option(
+    "--api-url",
+    help="Platform API URL (skips web discovery; useful for protected previews).",
+)
+_LOGIN_TOKEN = typer.Option(
+    "--token", help="Paste an existing API key instead of using the browser."
+)
+_LOGIN_NO_BROWSER = typer.Option(
+    "--no-browser", help="Print the authorization URL instead of opening a browser."
+)
+_ORG = typer.Option("--org", help="Organization id (defaults to the login's default organization).")
+_KIND = typer.Option(
+    "--kind", help="Artifact kind: model or harness (auto-detected when unambiguous)."
+)
+_PUSH_AS = typer.Option(
+    "--as", help="Remote name to publish under (local names may not be slug-safe)."
+)
+_PUSH_REF = typer.Option(
+    "--ref", help="Harness version or alias to push (default: champion, else latest)."
+)
+_PULL_VERSION = typer.Option("--version", help="Harness version to pull (default: latest).")
+_PULL_FORCE = typer.Option("--force", help="Replace an existing local artifact.")
+_ROOT = typer.Option("--root", help="Artifact root directory.")
+
+
+def login(
+    url: Annotated[str | None, _LOGIN_URL] = None,
+    api_url: Annotated[str | None, _LOGIN_API_URL] = None,
+    token: Annotated[str | None, _LOGIN_TOKEN] = None,
+    no_browser: Annotated[bool, _LOGIN_NO_BROWSER] = False,
+) -> None:
+    """Connect this machine to a platform account."""
+    credentials = load_credentials()
+    web_url = (url or credentials.web_url or DEFAULT_WEB_URL).rstrip("/")
+
+    if api_url is None:
+        try:
+            api_url = fetch_cli_config(web_url)
+        except PlatformError as error:
+            raise typer.BadParameter(f"{web_url} does not look like a platform: {error}") from error
+        if api_url is None:
+            raise typer.BadParameter(f"{web_url} did not advertise a backend URL; is it deployed?")
+    else:
+        api_url = api_url.rstrip("/")
+
+    if token is None:
+        token = _browser_login(web_url, open_browser=not no_browser)
+    if token is None or not token.strip():
+        _console.print("[red]No key received; nothing saved.[/red]")
+        raise typer.Exit(code=1)
+    token = token.strip()
+
+    with PlatformClient(api_url, token) as client:
+        try:
+            identity = client.whoami()
+        except PlatformError as error:
+            _console.print(f"[red]The key was rejected:[/red] {error}")
+            raise typer.Exit(code=1) from error
+
+    # A relogin may land on a different account: keep the saved default
+    # organization only if the new identity can still see it.
+    visible_orgs = {org.id for org in identity.orgs}
+    default_org = credentials.default_org if credentials.default_org in visible_orgs else None
+    if default_org is None and len(identity.orgs) == 1:
+        default_org = identity.orgs[0].id
+    updated = credentials.model_copy(
+        update={
+            "web_url": web_url,
+            "api_url": api_url,
+            "token": token,
+            "default_org": default_org,
+        }
+    )
+    path = save_credentials(updated)
+    org_names = ", ".join(org.name for org in identity.orgs) or "no organizations"
+    _console.print(f"{_CHECK} Connected to [bold]{org_names}[/bold] ({path})")
+    _print_orgs(identity, updated.default_org)
+
+
+def logout() -> None:
+    """Disconnect: delete the saved credential."""
+    credentials = load_credentials()
+    removed = clear_credentials()
+    if not removed:
+        _console.print("Not logged in; nothing to remove.")
+        return
+    _console.print(f"{_CHECK} Logged out.")
+    if credentials.token:
+        _console.print("The key itself stays valid until revoked on the platform's API keys page.")
+
+
+def status() -> None:
+    """Show the platform connection: account and organizations."""
+    credentials = load_credentials()
+    if not credentials.is_complete():
+        _console.print(
+            f"Not connected (no credential at {credentials_path()}). Run [bold]wmo login[/bold]."
+        )
+        raise typer.Exit(code=1)
+    with _client(credentials) as client:
+        try:
+            identity = client.whoami()
+        except PlatformError as error:
+            _console.print(f"[red]Connection check failed:[/red] {error}")
+            raise typer.Exit(code=1) from error
+    _console.print(f"{_CHECK} Connected to [bold]{credentials.web_url}[/bold]")
+    _console.print(f"  acting as: {identity.actor.kind} {identity.actor.id}")
+    _print_orgs(identity, credentials.default_org)
+
+
+def push(
+    name: Annotated[str, typer.Argument(help="Local world model or harness name.")],
+    org: Annotated[str | None, _ORG] = None,
+    kind: Annotated[str | None, _KIND] = None,
+    push_as: Annotated[str | None, _PUSH_AS] = None,
+    ref: Annotated[str | None, _PUSH_REF] = None,
+    root: Annotated[str, _ROOT] = ".wmo",
+) -> None:
+    """Publish a local world model or harness to the platform registry."""
+    model_dir = WorldModelStore(root).dir_for(name)
+    harness_exists = HarnessStore(root).exists(name)
+    resolved_kind = _resolve_kind(kind, model=model_dir is not None, harness=harness_exists)
+    remote_name = push_as or name
+
+    credentials, org_id = _require_connection(org)
+    with _client(credentials) as client:
+        if resolved_kind == "model" and model_dir is not None:
+            _push_model(client, org_id, remote_name, model_dir)
+        else:
+            _push_harness(client, org_id, remote_name, name, ref, root)
+
+
+def pull(
+    name: Annotated[str, typer.Argument(help="Remote world model or harness name.")],
+    org: Annotated[str | None, _ORG] = None,
+    kind: Annotated[str | None, _KIND] = None,
+    version: Annotated[int | None, _PULL_VERSION] = None,
+    force: Annotated[bool, _PULL_FORCE] = False,
+    root: Annotated[str, _ROOT] = ".wmo",
+) -> None:
+    """Fetch a world model or harness from the platform registry."""
+    if kind is not None and kind not in ("model", "harness"):
+        raise typer.BadParameter("--kind must be 'model' or 'harness'")
+    credentials, org_id = _require_connection(org)
+    with _client(credentials) as client:
+        resolved_kind = kind or _detect_remote_kind(client, org_id, name)
+        if resolved_kind == "model":
+            _pull_model(client, org_id, name, root, force=force)
+        else:
+            _pull_harness(client, org_id, name, root, version=version)
+
+
+# -- helpers -------------------------------------------------------------------------------------
+
+
+def _browser_login(web_url: str, *, open_browser: bool) -> str | None:
+    """Run the loopback browser flow; fall back to a hidden paste prompt."""
+    login_attempt = BrowserLogin(web_url)
+    try:
+        login_attempt.start()
+        key_name = f"wmo on {socket.gethostname()}"
+        authorize_url = login_attempt.authorize_url(key_name=key_name)
+        _console.print(f"Approve the request in your browser:\n  [bold]{authorize_url}[/bold]")
+        if open_browser:
+            webbrowser.open(authorize_url)
+        token = login_attempt.wait()
+    finally:
+        login_attempt.close()
+    if token is None:
+        _console.print("Timed out waiting for the browser.")
+        return typer.prompt("Paste an API key instead", hide_input=True, default="") or None
+    return token
+
+
+def _client(credentials: PlatformCredentials) -> PlatformClient:
+    if credentials.api_url is None or credentials.token is None:
+        raise typer.BadParameter("not connected to a platform; run `wmo login` first")
+    return PlatformClient(credentials.api_url, credentials.token)
+
+
+def _require_connection(org: str | None) -> tuple[PlatformCredentials, str]:
+    credentials = load_credentials()
+    if not credentials.is_complete():
+        raise typer.BadParameter("not connected to a platform; run `wmo login` first")
+    org_id = org or credentials.default_org
+    if not org_id:
+        raise typer.BadParameter(
+            "no organization selected; pass --org <id> (see `wmo status` for your organizations)"
+        )
+    return credentials, org_id
+
+
+def _resolve_kind(kind: str | None, *, model: bool, harness: bool) -> str:
+    if kind is not None:
+        if kind not in ("model", "harness"):
+            raise typer.BadParameter("--kind must be 'model' or 'harness'")
+        if kind == "model" and not model:
+            raise typer.BadParameter("no local world model has this name")
+        if kind == "harness" and not harness:
+            raise typer.BadParameter("no local harness has this name")
+        return kind
+    if model and harness:
+        raise typer.BadParameter("both a model and a harness have this name locally; pass --kind")
+    if model:
+        return "model"
+    if harness:
+        return "harness"
+    raise typer.BadParameter("no local world model or harness has this name")
+
+
+def _detect_remote_kind(client: PlatformClient, org_id: str, name: str) -> str:
+    model_names = {model.name for model in client.list_world_models(org_id)}
+    harness_names = {harness.name for harness in client.list_harnesses(org_id)}
+    if name in model_names and name in harness_names:
+        raise typer.BadParameter("both a model and a harness have this name remotely; pass --kind")
+    if name in model_names:
+        return "model"
+    if name in harness_names:
+        return "harness"
+    raise typer.BadParameter(f"the organization has no world model or harness named {name!r}")
+
+
+def _push_model(client: PlatformClient, org_id: str, remote_name: str, model_dir: Path) -> None:
+    meta = extract_push_meta(model_dir)
+    with tempfile.TemporaryDirectory(prefix="wmo-push-") as staging:
+        bundle = pack_model_dir(model_dir, Path(staging) / f"{remote_name}.tar.gz")
+        try:
+            pushed = client.push_model_bundle(
+                org_id,
+                remote_name,
+                bundle.path,
+                bundle.sha256,
+                bundle.byte_size,
+                meta,
+            )
+        except PlatformError as error:
+            if error.status_code == 422 and "name" in str(error):
+                raise typer.BadParameter(
+                    f"{error} — publish under a slug-safe name with --as"
+                ) from error
+            raise
+    _console.print(
+        f"{_CHECK} Pushed world model [bold]{pushed.name}[/bold] "
+        f"({bundle.byte_size:,} bytes, sha256 {bundle.sha256[:12]}…)"
+    )
+
+
+def _push_harness(
+    client: PlatformClient,
+    org_id: str,
+    remote_name: str,
+    local_name: str,
+    ref: str | None,
+    root: str,
+) -> None:
+    doc = HarnessStore(root).load(local_name, ref)
+    if remote_name != local_name:
+        doc = doc.model_copy(update={"name": remote_name})
+    pushed = client.push_harness_version(
+        org_id, remote_name, doc.model_dump(mode="json"), doc.doc_hash
+    )
+    if pushed.created:
+        _console.print(
+            f"{_CHECK} Pushed harness [bold]{pushed.name}[/bold] as remote v{pushed.version}"
+        )
+    else:
+        _console.print(
+            f"Remote [bold]{pushed.name}[/bold] v{pushed.version} already has this exact doc; "
+            "nothing to push."
+        )
+
+
+def _pull_model(client: PlatformClient, org_id: str, name: str, root: str, *, force: bool) -> None:
+    dest_dir = WorldModelStore(root).model_dir(name)
+    with tempfile.TemporaryDirectory(prefix="wmo-pull-") as staging:
+        bundle_path = Path(staging) / f"{name}.tar.gz"
+        client.download_model_bundle(org_id, name, bundle_path)
+        try:
+            unpack_model_bundle(bundle_path, dest_dir, force=force)
+        except FileExistsError as error:
+            raise typer.BadParameter(str(error)) from error
+    _console.print(f"{_CHECK} Pulled world model [bold]{name}[/bold] into {dest_dir}")
+
+
+def _pull_harness(
+    client: PlatformClient, org_id: str, name: str, root: str, *, version: int | None
+) -> None:
+    if version is None:
+        harness, _versions = client.get_harness(org_id, name)
+        if harness.latest_version < 1:
+            raise typer.BadParameter(f"remote harness {name!r} has no versions yet")
+        version = harness.latest_version
+    payload = client.get_harness_version(org_id, name, version)
+    doc = HarnessDoc.model_validate(payload.doc)
+    # Versions the platform recorded before doc_hash covered materialized paths carry the legacy
+    # hash; accept either so pre-change pathful (pi-node) records stay pullable.
+    if payload.doc_hash not in (doc.doc_hash, doc.legacy_doc_hash):
+        _console.print("[red]Pulled doc failed its integrity check; not saving.[/red]")
+        raise typer.Exit(code=1)
+    saved = HarnessStore(root).save_version(
+        doc.model_copy(update={"name": name}), alias=CHAMPION_ALIAS
+    )
+    _console.print(
+        f"{_CHECK} Pulled harness [bold]{name}[/bold] remote v{version} → local v{saved.version} "
+        f"(champion)"
+    )
+
+
+def _print_orgs(identity: WhoAmI, default_org: str | None) -> None:
+    if not identity.orgs:
+        _console.print("  no organizations visible to this key")
+        return
+    table = Table(show_header=True, header_style="bold")
+    table.add_column("organization")
+    table.add_column("id")
+    table.add_column("")
+    for org in identity.orgs:
+        marker = "default" if org.id == default_org else ""
+        table.add_row(org.name, org.id, marker)
+    _console.print(table)
+
+
+def register(app: typer.Typer) -> None:
+    """Attach the platform commands to the root CLI."""
+    app.command("login")(login)
+    app.command("logout")(logout)
+    app.command("status")(status)
+    app.command("push")(push)
+    app.command("pull")(pull)