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

wmo/platform/transfer.py

@@ -0,0 +1,158 @@
+"""Deterministic model-bundle packing and safe unpacking for push/pull.
+
+A pushed bundle is byte-compatible with the bundles the platform's own build
+pipeline produces: a gzipped tarball of the model directory with
+archive-relative member paths. Packing is an include-list — the model's
+`config.toml`, `metrics.json`, `card.json`, `prompts/`, and `index/` — so
+local `runs/` cost records and raw `traces/` (customer data) never leave the
+machine.
+
+Bundles can reach the platform's 1GB cap, so packing and unpacking are
+file-based: bytes stream between disk and the network without ever being held
+in memory whole.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import shutil
+import tarfile
+import tomllib
+import uuid
+from pathlib import Path
+
+from pydantic import BaseModel
+
+from wmo.config.config import HarnessConfig
+from wmo.core.types import JsonValue
+
+_INCLUDED_FILES = ("config.toml", "metrics.json", "card.json")
+_INCLUDED_DIRS = ("prompts", "index")
+
+_HASH_CHUNK_BYTES = 1024 * 1024
+
+
+class PackedModelBundle(BaseModel):
+    """A packed model bundle on disk, ready for upload."""
+
+    path: Path
+    sha256: str
+    byte_size: int
+
+
+class BundleFormatError(ValueError):
+    """The directory or bytes are not a valid world-model bundle."""
+
+
+def sha256_file(path: Path) -> str:
+    """Digest a file's contents without loading it whole."""
+    digest = hashlib.sha256()
+    with path.open("rb") as fh:
+        while chunk := fh.read(_HASH_CHUNK_BYTES):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def pack_model_dir(directory: Path, dest: Path) -> PackedModelBundle:
+    """Pack a model directory into the platform's bundle format at ``dest``.
+
+    Args:
+        directory: A built model directory (must contain `config.toml`).
+        dest: Where to write the gzipped tarball (parent must exist).
+
+    Returns:
+        The bundle file plus integrity metadata; member order is sorted so
+        identical inputs produce identical archives.
+
+    Raises:
+        BundleFormatError: If the directory is missing or has no config.toml.
+    """
+    if not directory.is_dir():
+        msg = f"model directory does not exist: {directory}"
+        raise BundleFormatError(msg)
+    if not (directory / "config.toml").is_file():
+        msg = f"{directory} has no config.toml; only built world models can be pushed"
+        raise BundleFormatError(msg)
+
+    members: list[Path] = []
+    for name in _INCLUDED_FILES:
+        path = directory / name
+        if path.is_file():
+            members.append(path)
+    for name in _INCLUDED_DIRS:
+        root = directory / name
+        if root.is_dir():
+            members.extend(sorted(path for path in root.rglob("*")))
+            members.append(root)
+
+    with tarfile.open(dest, mode="w:gz") as tar:
+        for path in sorted(set(members)):
+            tar.add(path, arcname=str(path.relative_to(directory)), recursive=False)
+    return PackedModelBundle(
+        path=dest,
+        sha256=sha256_file(dest),
+        byte_size=dest.stat().st_size,
+    )
+
+
+def unpack_model_bundle(source: Path, dest_dir: Path, *, force: bool = False) -> None:
+    """Unpack a pulled bundle file into a local model directory.
+
+    Extraction happens in a temporary sibling renamed into place, so a crashed
+    unpack never leaves a half-written model that later loads as real.
+
+    Args:
+        source: Downloaded gzipped tarball.
+        dest_dir: Target model directory (`.wmo/models/<name>`).
+        force: Replace an existing directory instead of refusing.
+
+    Raises:
+        BundleFormatError: If the file is not a readable bundle or a member
+            would escape the destination.
+        FileExistsError: If ``dest_dir`` exists and ``force`` is not set.
+    """
+    if dest_dir.exists() and not force:
+        msg = f"{dest_dir} already exists; pass --force to replace it"
+        raise FileExistsError(msg)
+    staging_dir = dest_dir.with_name(f"{dest_dir.name}.pull-{uuid.uuid4().hex}")
+    staging_dir.mkdir(parents=True)
+    try:
+        with tarfile.open(source, mode="r:gz") as tar:
+            # The "data" filter rejects absolute paths, traversal, and special
+            # members instead of writing them.
+            tar.extractall(staging_dir, filter="data")
+    except (tarfile.TarError, OSError) as error:
+        shutil.rmtree(staging_dir, ignore_errors=True)
+        msg = f"bundle could not be unpacked: {error}"
+        raise BundleFormatError(msg) from error
+    if dest_dir.exists():
+        shutil.rmtree(dest_dir, ignore_errors=True)
+    staging_dir.rename(dest_dir)
+
+
+def extract_push_meta(directory: Path) -> dict[str, JsonValue]:
+    """Derive the push metadata the platform stores alongside a bundle.
+
+    Parses the model's own `config.toml` (and `metrics.json` when present)
+    through wmo's typed config so the platform never reads files out of the
+    tarball.
+    """
+    config = HarnessConfig.model_validate(
+        tomllib.loads((directory / "config.toml").read_text(encoding="utf-8"))
+    )
+    meta: dict[str, JsonValue] = {
+        "serve_provider": config.serve_provider.value,
+        "embed_provider": config.embed_provider.value,
+        "embed_dim": config.embed_dim,
+        "gepa_budget": config.gepa_budget,
+    }
+    try:
+        meta["serve_model"] = config.serve_provider_config().model
+    except ValueError:
+        # No provider block for the serve provider; the platform column stays unset.
+        pass
+    metrics_path = directory / "metrics.json"
+    if metrics_path.is_file():
+        meta["metrics"] = json.loads(metrics_path.read_text(encoding="utf-8"))
+    return meta

wmo/providers/__init__.py

@@ -0,0 +1,40 @@
+"""Unified LLM provider layer.
+
+One interface (`Provider`), multiple backends, one entry point (`get_provider` — or
+`provider_or_chain`, which upgrades to the local `.wmo/fallback.toml` failover chain when present).
+All can be verified on startup with a cheap ping. Built fresh for this repo; no external client
+framework.
+"""
+
+from wmo.providers.base import (
+    DEFAULT_MAX_TOKENS,
+    Completion,
+    EmbedderKind,
+    Message,
+    Provider,
+    ProviderConfig,
+    ProviderKind,
+    VerifyResult,
+)
+from wmo.providers.models import ProviderModel, model_types_for_provider, resolve_provider_model
+from wmo.providers.registry import get_provider, verify_all, verify_embedder
+from wmo.providers.waterfall import WaterfallProvider, provider_or_chain
+
+__all__ = [
+    "Provider",
+    "ProviderConfig",
+    "ProviderKind",
+    "ProviderModel",
+    "EmbedderKind",
+    "DEFAULT_MAX_TOKENS",
+    "Completion",
+    "Message",
+    "VerifyResult",
+    "get_provider",
+    "provider_or_chain",
+    "WaterfallProvider",
+    "verify_all",
+    "verify_embedder",
+    "model_types_for_provider",
+    "resolve_provider_model",
+]

wmo/providers/_bedrock_chat.py

@@ -0,0 +1,155 @@
+"""Structured tool-calling translation for Bedrock Converse providers."""
+
+from __future__ import annotations
+
+import json
+from typing import cast
+
+from llm_waterfall import ChatRequest, ChatResponse
+from pydantic import JsonValue
+
+
+def converse_request(request: ChatRequest, model: str) -> dict[str, object]:
+    """Translate the provider-neutral structured contract to Bedrock Converse."""
+    system: list[dict[str, str]] = []
+    messages: list[dict[str, object]] = []
+
+    def push(role: str, content: list[dict[str, object]]) -> None:
+        if messages and messages[-1]["role"] == role:
+            existing = cast("list[dict[str, object]]", messages[-1]["content"])
+            existing.extend(content)
+        else:
+            messages.append({"role": role, "content": content})
+
+    for message in request.messages:
+        if message.role in ("system", "developer"):
+            text = _chat_text(message.content)
+            if text:
+                system.append({"text": text})
+            continue
+        if message.role == "tool":
+            push(
+                "user",
+                [
+                    {
+                        "toolResult": {
+                            "toolUseId": message.tool_call_id or "",
+                            "content": [{"text": _chat_text(message.content)}],
+                        }
+                    }
+                ],
+            )
+            continue
+        blocks: list[dict[str, object]] = []
+        text = _chat_text(message.content)
+        if text:
+            blocks.append({"text": text})
+        for tool_call in message.tool_calls or []:
+            try:
+                arguments = json.loads(tool_call.function.arguments)
+            except ValueError:
+                arguments = {}
+            blocks.append(
+                {
+                    "toolUse": {
+                        "toolUseId": tool_call.id,
+                        "name": tool_call.function.name,
+                        "input": arguments,
+                    }
+                }
+            )
+        if blocks:
+            push("assistant" if message.role == "assistant" else "user", blocks)
+
+    max_tokens = request.max_tokens or request.max_completion_tokens or 4096
+    inference: dict[str, float | int] = {"maxTokens": max_tokens}
+    if request.temperature is not None:
+        inference["temperature"] = request.temperature
+    result: dict[str, object] = {
+        "modelId": model,
+        "messages": messages,
+        "inferenceConfig": inference,
+    }
+    if system:
+        result["system"] = system
+    if request.tools:
+        tools = [
+            {
+                "toolSpec": {
+                    "name": tool.function.name,
+                    "description": tool.function.description,
+                    "inputSchema": {"json": tool.function.parameters},
+                }
+            }
+            for tool in request.tools
+        ]
+        tool_config: dict[str, object] = {"tools": tools}
+        choice = request.tool_choice
+        if choice == "required":
+            tool_config["toolChoice"] = {"any": {}}
+        elif isinstance(choice, dict):
+            function = choice.get("function")
+            if isinstance(function, dict) and isinstance(function.get("name"), str):
+                tool_config["toolChoice"] = {"tool": {"name": function["name"]}}
+        if choice != "none":
+            result["toolConfig"] = tool_config
+    return result
+
+
+def converse_response(raw: object, model: str) -> ChatResponse:
+    """Translate a Bedrock Converse response to the structured provider contract."""
+    response = cast("dict[str, object]", raw)
+    output = cast("dict[str, object]", response["output"])
+    message_data = cast("dict[str, object]", output["message"])
+    blocks = cast("list[dict[str, object]]", message_data["content"])
+    text = "".join(str(block["text"]) for block in blocks if "text" in block)
+    tool_calls: list[dict[str, object]] = []
+    for block in blocks:
+        use = block.get("toolUse")
+        if not isinstance(use, dict):
+            continue
+        tool_calls.append(
+            {
+                "id": str(use.get("toolUseId", "")),
+                "type": "function",
+                "function": {
+                    "name": str(use.get("name", "")),
+                    "arguments": json.dumps(use.get("input", {})),
+                },
+            }
+        )
+    stop_reason = str(response.get("stopReason", "end_turn"))
+    finish_reason = {
+        "tool_use": "tool_calls",
+        "max_tokens": "length",
+        "content_filtered": "content_filter",
+        "guardrail_intervened": "content_filter",
+    }.get(stop_reason, "stop")
+    message: dict[str, object] = {"role": "assistant", "content": text}
+    if tool_calls:
+        message["tool_calls"] = tool_calls
+    usage = response.get("usage")
+    usage_data = usage if isinstance(usage, dict) else {}
+    return ChatResponse.model_validate(
+        {
+            "model": model,
+            "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}],
+            "usage": {
+                "prompt_tokens": usage_data.get("inputTokens", 0),
+                "completion_tokens": usage_data.get("outputTokens", 0),
+            },
+        }
+    )
+
+
+def _chat_text(content: JsonValue) -> str:
+    """Flatten the text-bearing forms used by OpenAI-compatible chat messages."""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        parts: list[str] = []
+        for item in content:
+            if isinstance(item, dict) and isinstance(item.get("text"), str):
+                parts.append(item["text"])
+        return "".join(parts)
+    return "" if content is None else str(content)

wmo/providers/_openai_common.py

@@ -0,0 +1,182 @@
+"""Shared request mapping / response parsing for the two OpenAI-shaped backends.
+
+`OpenAIProvider` and `AzureOpenAIProvider` differ only in how their client is constructed; the
+chat-completion and embedding wire formats are identical, so that logic lives here.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Protocol, cast
+
+from llm_waterfall import ChatMaxTokensField, ChatRequest, ChatResponse
+from openai import BadRequestError
+
+from wmo.providers.base import Completion, Message, StreamChunk, TokenUsage
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from openai.types import CreateEmbeddingResponse
+    from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
+
+
+class _ChatCompletions(Protocol):
+    def create(
+        self,
+        *,
+        model: str,
+        messages: list[ChatCompletionMessageParam],
+        max_completion_tokens: int,
+        temperature: float = ...,
+    ) -> ChatCompletion: ...
+
+
+class _Embeddings(Protocol):
+    def create(
+        self, *, model: str, input: list[str], dimensions: int = ...
+    ) -> CreateEmbeddingResponse: ...
+
+
+def to_messages(system: str, messages: list[Message]) -> list[ChatCompletionMessageParam]:
+    """Fold the system prompt into the message list as OpenAI's leading `system` turn."""
+    out: list[dict[str, str]] = []
+    if system:
+        out.append({"role": "system", "content": system})
+    out.extend({"role": m.role, "content": m.content} for m in messages)
+    return cast("list[ChatCompletionMessageParam]", out)
+
+
+def complete(
+    chat_completions: _ChatCompletions,
+    model: str,
+    system: str,
+    messages: list[Message],
+    max_tokens: int,
+    temperature: float | None = None,
+    max_tokens_field: str = "max_completion_tokens",
+) -> Completion:
+    """Run one chat completion and map it onto our `Completion`.
+
+    `max_tokens_field` names the output-budget parameter the deployment accepts: GPT-5.x wants
+    `max_completion_tokens`, while Azure MaaS open models (DeepSeek, Kimi) still take the
+    classic `max_tokens` (see `ProviderConfig.resolved_chat_max_tokens_field`). `temperature`
+    is sent ONLY when given: GPT 5.5's reasoning models reject non-default sampling params
+    (callers pass None), while OpenAI-compatible servers (vLLM policies) need it.
+    """
+    # The output-budget param name is dynamic, so this call crosses the SDK boundary through
+    # the same one-line cast `stream` uses below.
+    resource = cast("Any", chat_completions)
+    base_kwargs: dict[str, Any] = {
+        "model": model,
+        "messages": to_messages(system, messages),
+        max_tokens_field: max_tokens,
+    }
+    if temperature is None:
+        response: ChatCompletion = resource.create(**base_kwargs)
+    else:
+        try:
+            response = resource.create(**base_kwargs, temperature=temperature)
+        except BadRequestError as exc:
+            # Reasoning-model deployments (GPT-5.x behind Azure/custom endpoints) reject any
+            # non-default temperature with a 400 unsupported_value. The caller can't know which
+            # models sample; degrade to the model's default rather than failing the request.
+            if "temperature" not in str(exc):
+                raise
+            response = resource.create(**base_kwargs)
+    if not response.choices:
+        # Content filtering (and some error modes) can return zero choices; surface it clearly
+        # rather than letting choices[0] raise a bare IndexError.
+        raise ValueError(f"{model} returned no choices")
+    text = response.choices[0].message.content or ""
+    usage = response.usage
+    token_usage = _chat_usage(usage) if usage is not None else TokenUsage()
+    return Completion(text=text, usage=token_usage)
+
+
+def stream(
+    chat_completions: object,
+    model: str,
+    system: str,
+    messages: list[Message],
+    max_tokens: int,
+    temperature: float | None = None,
+    max_tokens_field: str = "max_completion_tokens",
+) -> Iterator[StreamChunk]:
+    """Stream one chat completion as `StreamChunk`s (deltas, then a terminal chunk with usage).
+
+    `stream_options.include_usage` makes the wire stream end with a usage-bearing chunk, so the
+    terminal `StreamChunk` carries real token counts instead of estimates. `temperature` and
+    `max_tokens_field` follow the same rules as `complete`.
+    """
+    resource = cast("Any", chat_completions)
+    kwargs: dict[str, Any] = {
+        "model": model,
+        "messages": to_messages(system, messages),
+        max_tokens_field: max_tokens,
+        "stream": True,
+        "stream_options": {"include_usage": True},
+    }
+    if temperature is not None:
+        kwargs["temperature"] = temperature
+    usage = TokenUsage()
+    upstream = resource.create(**kwargs)
+    try:
+        for chunk in upstream:
+            choices = getattr(chunk, "choices", None) or []
+            if choices:
+                delta = getattr(choices[0], "delta", None)
+                text = getattr(delta, "content", None) if delta is not None else None
+                if text:
+                    yield StreamChunk(delta=text)
+            chunk_usage = getattr(chunk, "usage", None)
+            if chunk_usage is not None:
+                usage = _chat_usage(chunk_usage)
+    finally:
+        # The SDK Stream holds an httpx response; without an explicit close an abandoned
+        # stream releases the connection only when the object is garbage collected.
+        close = getattr(upstream, "close", None)
+        if callable(close):
+            close()
+    yield StreamChunk(done=True, usage=usage)
+
+
+def _chat_usage(usage: object) -> TokenUsage:
+    """Chat-completions usage -> TokenUsage, including the cached-prompt split when reported."""
+    details = getattr(usage, "prompt_tokens_details", None)
+    return TokenUsage(
+        input_tokens=getattr(usage, "prompt_tokens", 0) or 0,
+        output_tokens=getattr(usage, "completion_tokens", 0) or 0,
+        cached_input_tokens=(getattr(details, "cached_tokens", None) or 0) if details else 0,
+    )
+
+
+def complete_chat(
+    chat_completions: object,
+    model: str,
+    request: ChatRequest,
+    *,
+    max_tokens_field: ChatMaxTokensField,
+) -> ChatResponse:
+    """Run a validated structured request against an OpenAI-compatible SDK resource."""
+    # ChatRequest validates the stable tool-calling core before this SDK boundary. The OpenAI
+    # package models its evolving request surface as a large TypedDict union, so the narrow cast
+    # preserves forward-compatible extra fields without leaking Any into the public contract.
+    resource = cast("Any", chat_completions)
+    response = resource.create(**request.provider_payload(model, max_tokens_field=max_tokens_field))
+    return ChatResponse.model_validate(response.model_dump(mode="json"))
+
+
+def embed(
+    embeddings: _Embeddings, model: str, texts: list[str], dim: int | None = None
+) -> list[list[float]]:
+    """Embed `texts` against `model` (an OpenAI model id, or an Azure embedding deployment).
+
+    `dim`, when set, requests a specific output dimension via the `dimensions` param (supported by
+    text-embedding-3-* and their Azure deployments) so the index and query vectors match.
+    """
+    response = (
+        embeddings.create(model=model, input=texts, dimensions=dim)
+        if dim is not None
+        else embeddings.create(model=model, input=texts)
+    )
+    return [item.embedding for item in response.data]