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

llm_waterfall/adapters/bedrock.py

@@ -0,0 +1,309 @@
+"""AWS Bedrock adapter (Anthropic Messages schema via InvokeModel, Titan embeddings).
+
+Credentials come from the boto3 chain, or a named AWS profile when `Backend.profile` is set —
+`boto3.Session(profile_name=...)` — so one waterfall chain can span multiple AWS accounts.
+"""
+
+from __future__ import annotations
+
+import json
+import threading
+from typing import TYPE_CHECKING, TypedDict, cast
+
+from pydantic import JsonValue
+
+from llm_waterfall.adapters.base import missing_sdk_error
+from llm_waterfall.types import (
+    Backend,
+    ChatRequest,
+    ChatResponse,
+    Message,
+    TokenUsage,
+)
+
+if TYPE_CHECKING:
+    from botocore.client import BaseClient
+
+# Bedrock speaks the same Anthropic Messages schema as the direct API, pinned by this version tag.
+_ANTHROPIC_BEDROCK_VERSION = "bedrock-2023-05-31"
+
+# Default Titan text-embeddings model (v2 supports `dimensions` 256/512/1024).
+_DEFAULT_EMBED_MODEL = "amazon.titan-embed-text-v2:0"
+
+
+class _ContentBlock(TypedDict):
+    type: str
+    text: str
+
+
+class _Usage(TypedDict):
+    input_tokens: int
+    output_tokens: int
+
+
+class _MessagesResponse(TypedDict):
+    content: list[_ContentBlock]
+    usage: _Usage
+
+
+class _TitanEmbedResponse(TypedDict, total=False):
+    embedding: list[float]
+    inputTextTokenCount: int
+
+
+class _ConverseContentBlock(TypedDict, total=False):
+    text: str
+    toolUse: dict[str, object]
+
+
+class _ConverseMessage(TypedDict):
+    role: str
+    content: list[_ConverseContentBlock]
+
+
+class _ConverseOutput(TypedDict):
+    message: _ConverseMessage
+
+
+class _ConverseUsage(TypedDict):
+    inputTokens: int
+    outputTokens: int
+
+
+class _ConverseResponse(TypedDict):
+    output: _ConverseOutput
+    stopReason: str
+    usage: _ConverseUsage
+
+
+class BedrockAdapter:
+    """Claude (and Titan embeddings) via the Bedrock Runtime."""
+
+    def __init__(self, backend: Backend) -> None:
+        self.backend = backend
+        self._client: BaseClient | None = None
+        self._lock = threading.Lock()
+
+    def _get_client(self) -> BaseClient:
+        # Lazy + lock-guarded: boto3 is an optional extra, and boto3.Session construction is not
+        # thread-safe (the resulting client is — one Waterfall is shared across thread pools).
+        if self._client is None:
+            with self._lock:
+                if self._client is None:
+                    try:
+                        import boto3
+                        from botocore.config import Config
+                    except ModuleNotFoundError as exc:
+                        raise missing_sdk_error("boto3", "bedrock") from exc
+
+                    # Bound each request so a stalled connection RAISES instead of blocking
+                    # forever — the waterfall can only fail over on a raised error. read_timeout
+                    # is generous because reasoning models can legitimately generate for minutes;
+                    # a mid-generation cutoff wastes the whole call and silently substitutes a
+                    # different model into an eval.
+                    #
+                    # total_max_attempts=1 disables botocore's OWN retries on purpose (it counts
+                    # the initial request; botocore's `max_attempts` counts retries AFTER it, so
+                    # `max_attempts: 1` would still allow one hidden retry). Throttling/5xx/
+                    # timeouts must surface IMMEDIATELY to the waterfall, which owns retry policy
+                    # — SDK retries stack multiplicatively under the failover chain.
+                    config = Config(
+                        connect_timeout=self.backend.connect_timeout_s,
+                        read_timeout=self.backend.read_timeout_s,
+                        retries={"total_max_attempts": 1},
+                    )
+                    session = boto3.Session(
+                        profile_name=self.backend.profile, region_name=self.backend.region
+                    )
+                    self._client = session.client("bedrock-runtime", config=config)
+        return self._client
+
+    def complete(
+        self,
+        system: str,
+        messages: list[Message],
+        *,
+        temperature: float | None,
+        max_tokens: int,
+    ) -> tuple[str, TokenUsage]:
+        """One InvokeModel call with the Anthropic Messages body."""
+        body: dict[str, object] = {
+            "anthropic_version": _ANTHROPIC_BEDROCK_VERSION,
+            "max_tokens": max_tokens,
+            "system": system,
+            "messages": [{"role": m.role, "content": m.content} for m in messages],
+        }
+        # Claude 4.7+ rejects sampling params; only forward temperature when explicitly set.
+        if temperature is not None:
+            body["temperature"] = temperature
+        raw = self._get_client().invoke_model(modelId=self.backend.model, body=json.dumps(body))
+        data = cast("_MessagesResponse", json.loads(raw["body"].read()))
+        text = "".join(block["text"] for block in data["content"] if block["type"] == "text")
+        usage = TokenUsage(
+            input_tokens=data["usage"]["input_tokens"],
+            output_tokens=data["usage"]["output_tokens"],
+        )
+        return text, usage
+
+    def complete_chat(self, request: ChatRequest) -> ChatResponse:
+        """Run a structured tool-calling request through Bedrock Converse."""
+        converse_request = _converse_request(request, self.backend.model)
+        response = cast("_ConverseResponse", self._get_client().converse(**converse_request))
+        blocks = response["output"]["message"]["content"]
+        text = "".join(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 use is None:
+                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 = response["stopReason"]
+        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
+        return ChatResponse.model_validate(
+            {
+                "model": self.backend.model,
+                "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}],
+                "usage": {
+                    "prompt_tokens": response["usage"]["inputTokens"],
+                    "completion_tokens": response["usage"]["outputTokens"],
+                },
+            }
+        )
+
+    def embed_model_id(self) -> str | None:
+        """The model embed() resolves to — the single source of truth for embed attribution."""
+        return self.backend.embed_model or _DEFAULT_EMBED_MODEL
+
+    def embed(self, texts: list[str]) -> tuple[list[list[float]], TokenUsage]:
+        """Embed via Amazon Titan (one InvokeModel per text — Titan has no batch input)."""
+        model = self.backend.embed_model or _DEFAULT_EMBED_MODEL
+        client = self._get_client()
+        vectors: list[list[float]] = []
+        input_tokens = 0
+        for text in texts:
+            body: dict[str, object] = {"inputText": text}
+            if self.backend.embed_dim is not None:
+                body["dimensions"] = self.backend.embed_dim
+                body["normalize"] = True
+            raw = client.invoke_model(modelId=model, body=json.dumps(body))
+            data = cast("_TitanEmbedResponse", json.loads(raw["body"].read()))
+            vectors.append(data["embedding"])
+            input_tokens += data.get("inputTextTokenCount", 0)
+        return vectors, TokenUsage(input_tokens=input_tokens)
+
+
+def _converse_request(request: ChatRequest, model: str) -> dict[str, object]:
+    """Translate the structured OpenAI-compatible 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 _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)

llm_waterfall/adapters/openai.py

@@ -0,0 +1,130 @@
+"""OpenAI adapter (chat completions + embeddings). Reads OPENAI_API_KEY from the environment."""
+
+from __future__ import annotations
+
+import threading
+from typing import TYPE_CHECKING, Any, cast
+
+from llm_waterfall.adapters.base import missing_sdk_error
+from llm_waterfall.types import (
+    Backend,
+    ChatRequest,
+    ChatResponse,
+    EmbeddingsUnsupported,
+    Message,
+    TokenUsage,
+)
+
+if TYPE_CHECKING:
+    from openai import OpenAI
+    from openai.types.chat import ChatCompletionMessageParam
+
+_DEFAULT_EMBED_MODEL = "text-embedding-3-small"
+
+
+class OpenAIAdapter:
+    """GPT-5.x via chat completions; text-embedding-3-* for embeddings."""
+
+    def __init__(self, backend: Backend) -> None:
+        self.backend = backend
+        self._client: OpenAI | None = None
+        self._lock = threading.Lock()
+
+    def _get_client(self) -> OpenAI:
+        if self._client is None:
+            with self._lock:
+                if self._client is None:
+                    try:
+                        import httpx
+                        from openai import OpenAI
+                    except ModuleNotFoundError as exc:
+                        raise missing_sdk_error("openai", "openai") from exc
+
+                    # max_retries=0: the waterfall owns retry policy — SDK retries would stack
+                    # multiplicatively under the failover chain. Granular httpx.Timeout so a
+                    # dead endpoint fails over after connect_timeout_s, not read_timeout_s.
+                    self._client = OpenAI(
+                        base_url=self.backend.endpoint,
+                        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 chat completion.
+
+        The backend contract selects the output-token field. Omitting a default temperature keeps
+        this compatible with GPT-5.x reasoning models, which reject non-default sampling params.
+        """
+        wire: list[dict[str, str]] = []
+        if system:
+            wire.append({"role": "system", "content": system})
+        wire.extend({"role": m.role, "content": m.content} for m in messages)
+        api_messages = cast("list[ChatCompletionMessageParam]", wire)
+        chat = self._get_client().chat.completions
+        model = self._request_model()
+        payload: dict[str, object] = {
+            "model": model,
+            "messages": api_messages,
+            self.backend.chat_max_tokens_field: max_tokens,
+        }
+        if temperature is not None:
+            payload["temperature"] = temperature
+        response = chat.create(**cast("Any", payload))
+        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 = (
+            TokenUsage(input_tokens=usage.prompt_tokens, output_tokens=usage.completion_tokens)
+            if usage is not None
+            else TokenUsage()
+        )
+        return text, token_usage
+
+    def complete_chat(self, request: ChatRequest) -> ChatResponse:
+        """Run a full OpenAI-compatible tool-calling request through this backend."""
+        payload = request.provider_payload(
+            self._request_model(), max_tokens_field=self.backend.chat_max_tokens_field
+        )
+        # The OpenAI SDK's input TypedDict is intentionally not our public contract. This one
+        # narrow cast sits at the SDK boundary after ChatRequest validated the structured core;
+        # provider_payload preserves forward-compatible extra fields emitted by agent SDKs.
+        response = self._get_client().chat.completions.create(**cast("Any", payload))
+        return ChatResponse.model_validate(response.model_dump(mode="json"))
+
+    def _request_model(self) -> str:
+        """The id sent as `model` on the wire (Azure overrides this with the deployment)."""
+        return self.backend.model
+
+    def embed_model_id(self) -> str | None:
+        """The model embed() resolves to — the single source of truth for embed attribution."""
+        return self.backend.embed_model or _DEFAULT_EMBED_MODEL
+
+    def embed(self, texts: list[str]) -> tuple[list[list[float]], TokenUsage]:
+        """Embed against `embed_model_id()` (OpenAI default: text-embedding-3-small)."""
+        model = self.embed_model_id()
+        if model is None:  # pragma: no cover - only reachable via subclasses
+            raise EmbeddingsUnsupported("no embedding deployment configured for this backend")
+        embeddings = self._get_client().embeddings
+        if self.backend.embed_dim is None:
+            response = embeddings.create(model=model, input=texts)
+        else:
+            response = embeddings.create(
+                model=model, input=texts, dimensions=self.backend.embed_dim
+            )
+        usage = getattr(response, "usage", None)
+        input_tokens = getattr(usage, "prompt_tokens", 0) if usage is not None else 0
+        return [item.embedding for item in response.data], TokenUsage(input_tokens=input_tokens)

llm_waterfall/classify.py

@@ -0,0 +1,184 @@
+"""Capacity-vs-client error classification — the contract that decides failover vs propagate.
+
+Capacity errors (throttling, transient 5xx, model-not-ready, transport timeouts) mean "this backend
+can't serve right now; the next one might" — the waterfall spills. Client errors (bad request,
+auth, validation) mean "this request is wrong" — failing over would just mask a real bug behind a
+different model's answer, so they propagate immediately.
+
+Classification is pure duck-typing with zero SDK imports, so the package works with any subset of
+provider SDKs installed and the logic is testable with fake exceptions. Signals are checked in
+fidelity order: a structured error code always wins over anything derived from the message.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from llm_waterfall.types import WaterfallExhausted
+
+Outcome = Literal["capacity_error", "client_error"]
+
+# Botocore error CODES that mean "this model is capacity-constrained right now" — the reliable
+# signal (from `exc.response["Error"]["Code"]`), preferred over everything else.
+_CAPACITY_ERROR_CODES = frozenset(
+    {
+        "ThrottlingException",
+        "TooManyRequestsException",
+        "ServiceUnavailableException",
+        "ServiceQuotaExceededException",
+        "ModelNotReadyException",
+        "ModelTimeoutException",
+        "InternalServerException",  # transient 5xx, safe to fail over
+        "InternalFailure",  # botocore's generic 5xx code (killed a live full-slice eval run)
+    }
+)
+
+# Exception class NAMES from the OpenAI/Anthropic SDKs (and their httpx/httpcore transport) that
+# mean capacity/transport failure. Matched by name + module gate rather than isinstance so the
+# SDKs never need to be importable.
+_SDK_CAPACITY_TYPE_NAMES = frozenset(
+    {
+        "RateLimitError",
+        "APITimeoutError",
+        "APIConnectionError",
+        "InternalServerError",
+        "OverloadedError",  # anthropic 529
+        # httpx/httpcore transient transport failures (no status code on any of them):
+        "ConnectTimeout",
+        "ReadTimeout",
+        "WriteTimeout",
+        "PoolTimeout",
+        "ConnectError",
+        "ReadError",
+        "WriteError",
+        "RemoteProtocolError",  # server disconnected mid-response
+    }
+)
+
+# Top-level modules whose exceptions we trust for name/status-code classification. An exception
+# named `RateLimitError` from arbitrary application code proves nothing. tinker's exceptions
+# mirror openai's shape (same capacity type names; `status_code` on its APIStatusError family).
+_TRUSTED_SDK_MODULES = frozenset({"openai", "anthropic", "httpx", "httpcore", "tinker"})
+
+# HTTP statuses on a trusted SDK error that mean capacity/transient failure. 529 is Anthropic's
+# "overloaded". Everything else (400/401/403/404/422/...) is a client error.
+_CAPACITY_STATUS_CODES = frozenset({408, 429, 500, 502, 503, 504, 529})
+
+# Last-resort substrings for structureless transport errors (e.g. botocore's Read/ConnectTimeout,
+# raw ConnectionError), which carry no code, no trusted type, no status. Kept conservative: only
+# phrases that unambiguously mean capacity/transport failure, NOT generic tokens like "429"/"503"/
+# "capacity" that can appear inside a bad-request message and turn a real error into a silent
+# failover. (That was a real bug once: `ValidationException: request timeout too large`.)
+_TRANSPORT_MARKERS = (
+    "throttl",
+    "read timeout",
+    "connect timeout",
+    "connection reset",
+    "connection aborted",
+    "connection was closed",  # botocore ConnectionClosedError (dropped mid-response)
+    "could not connect to the endpoint",  # botocore EndpointConnectionError (DNS/unreachable)
+    "timed out",
+    "service unavailable",
+    "model not ready",
+)
+
+
+def _from_trusted_sdk(exc: Exception) -> bool:
+    """Whether `exc` was defined by a provider SDK or its HTTP transport."""
+    module = type(exc).__module__ or ""
+    return module.split(".")[0] in _TRUSTED_SDK_MODULES
+
+
+# Botocore transport-failure base classes, matched anywhere in the exception's MRO by
+# (module, name) so subclasses we have never seen still classify without importing botocore.
+# Message matching alone lost this game twice: ConnectionClosedError and EndpointConnectionError
+# each killed a multi-dollar run before their phrasing joined _TRANSPORT_MARKERS.
+_BOTOCORE_TRANSPORT_BASES = frozenset({"ConnectionError", "HTTPClientError"})
+
+
+def _is_botocore_transport_error(exc: Exception) -> bool:
+    """Whether `exc` is any member of botocore's ConnectionError/HTTPClientError family."""
+    for klass in type(exc).__mro__:
+        module = getattr(klass, "__module__", "") or ""
+        if module.split(".")[0] == "botocore" and klass.__name__ in _BOTOCORE_TRANSPORT_BASES:
+            return True
+    return False
+
+
+# tinker's sidecar family (the SDK's local transport subprocess dying, failing to start, or
+# losing IPC mid-request), matched anywhere in the MRO by (module, name) like the botocore
+# transport bases above: these carry no status_code, no openai-shaped capacity name, and no
+# reliable message phrasing, but a fresh attempt (which respawns the sidecar) can succeed.
+_TINKER_TRANSPORT_BASES = frozenset({"SidecarError"})
+
+
+def _is_tinker_transport_error(exc: Exception) -> bool:
+    """Whether `exc` is any member of tinker's SidecarError family."""
+    for klass in type(exc).__mro__:
+        module = getattr(klass, "__module__", "") or ""
+        if module.split(".")[0] == "tinker" and klass.__name__ in _TINKER_TRANSPORT_BASES:
+            return True
+    return False
+
+
+def outcome_for(exc: Exception) -> Outcome:
+    """Classify an exception raised by a backend attempt.
+
+    Checks, in fidelity order:
+      1. Our own `WaterfallExhausted` — a nested chain that exhausted was capacity-constrained
+         by definition (outer waterfalls/retry loops must treat it as transient).
+      2. Structured botocore error code — authoritative; a non-capacity code stops here so a
+         message substring can never overrule it.
+      3. Botocore's ConnectionError/HTTPClientError family, matched by MRO (module, name) —
+         transport failures whose subclasses carry no code and keep inventing new phrasings.
+      4. Tinker's non-HTTP failure families, same MRO matching: the SidecarError family is
+         transport-shaped (retry can respawn the sidecar), and RequestFailedError classifies
+         by its `category` field (a "user" request is wrong; "server"/"unknown" are transient).
+      5. SDK exception type name, gated on the defining module.
+      6. HTTP status, from a `status_code` attribute on the exception or on its `.response`
+         (httpx.HTTPStatusError carries it there), same module gate.
+      7. Conservative transport-phrase substrings for structureless errors.
+    """
+    if isinstance(exc, WaterfallExhausted):
+        return "capacity_error"
+
+    response = getattr(exc, "response", None)
+    if isinstance(response, dict):
+        error = response.get("Error")
+        # Guard the shape: a duck-typed `.response` dict may carry a non-dict "Error" value;
+        # classification must never raise over the original exception.
+        code = error.get("Code") if isinstance(error, dict) else None
+        if code is not None:
+            return "capacity_error" if code in _CAPACITY_ERROR_CODES else "client_error"
+
+    if _is_botocore_transport_error(exc):
+        return "capacity_error"
+
+    if _is_tinker_transport_error(exc):
+        return "capacity_error"
+
+    if _from_trusted_sdk(exc):
+        if type(exc).__name__ in _SDK_CAPACITY_TYPE_NAMES:
+            return "capacity_error"
+        if type(exc).__name__ == "RequestFailedError":
+            # tinker's async-request failure carries a StrEnum `category` instead of a
+            # status code: "user" means this request is wrong (propagate); "server" and
+            # "unknown" mean the service failed it and a fresh attempt can succeed.
+            category = str(getattr(exc, "category", "")).lower()
+            return "client_error" if category == "user" else "capacity_error"
+        status = getattr(exc, "status_code", None)
+        if not isinstance(status, int):
+            # httpx.HTTPStatusError keeps the status on the response object instead.
+            status = getattr(response, "status_code", None)
+        if isinstance(status, int):
+            return "capacity_error" if status in _CAPACITY_STATUS_CODES else "client_error"
+
+    message = str(exc).lower()
+    if any(marker in message for marker in _TRANSPORT_MARKERS):
+        return "capacity_error"
+    return "client_error"
+
+
+def is_capacity_error(exc: Exception) -> bool:
+    """True when the waterfall should spill to the next backend instead of raising."""
+    return outcome_for(exc) == "capacity_error"