frugal-sdk-python 0.1.0__py3-none-any.whl

frugal_metrics/wrappers/bedrock_models.py

@@ -0,0 +1,430 @@
+"""Per-provider response-shape dispatch for AWS Bedrock.
+
+Every Bedrock customer can hit any provider through `invoke_model`,
+`invoke_model_with_response_stream`, `converse`, or `converse_stream`.
+The shape of the request body and the response usage block depends on
+the underlying provider — so `wrap_bedrock` looks the modelId up in the
+DISPATCH table to find the right extractor.
+
+Coverage in v1: per-provider extractors for Anthropic Claude, OpenAI
+gpt-oss, Amazon Nova, Meta Llama, Mistral. Everything else falls
+through to the UNIVERSAL_FALLBACK_HANDLER, which uses the
+`amazon-bedrock-invocationMetrics` chunk / `x-amzn-bedrock-*-token-count`
+headers that Bedrock attaches to every response regardless of provider.
+"""
+
+from __future__ import annotations
+
+import io
+import json
+import re
+from dataclasses import dataclass
+from typing import Any, Callable
+
+from frugal_metrics.instrument import UsageReport
+from frugal_metrics.prompt_analyzer import RequestPayload
+
+
+# --- Cross-region inference profile prefixes ---------------------------------
+#
+# Newer Claude / Nova / Llama models on Bedrock are typically invoked through
+# cross-region inference profile IDs that prefix the underlying modelId with a
+# region scope, e.g. `us.anthropic.claude-haiku-4-5-20251001-v1:0`.
+
+CROSS_REGION_PREFIX_RE = re.compile(r"^(us|eu|jp|apac|au|global)\.")
+
+
+def strip_cross_region_prefix(model_id: str) -> str:
+    """Drop a leading region prefix so dispatch can match the base model ID.
+
+    The wrapper still emits the original (un-stripped) modelId in
+    `gen_ai.request.model` so customers can tell `us.` vs `eu.` apart.
+    """
+    return CROSS_REGION_PREFIX_RE.sub("", model_id)
+
+
+# --- Dispatch entry shape ----------------------------------------------------
+
+
+@dataclass(frozen=True)
+class BedrockModelHandler:
+    name: str
+    extract_invoke_model_usage: Callable[[dict[str, Any]], UsageReport]
+    extract_converse_usage: Callable[[dict[str, Any]], UsageReport]
+    extract_invoke_model_stream_event: Callable[[dict[str, Any]], "UsageReport | None"]
+    capture_invoke_model_payload: Callable[[dict[str, Any]], "RequestPayload | None"]
+
+
+# --- Per-provider InvokeModel response extractors ----------------------------
+
+
+def _anthropic_invoke_usage(body: dict[str, Any]) -> UsageReport:
+    usage = body.get("usage") or {}
+    thinking = usage.get("thinking_tokens")
+    return UsageReport(
+        response_model=body.get("model"),
+        input_tokens=usage.get("input_tokens"),
+        output_tokens=usage.get("output_tokens"),
+        cache_read_tokens=usage.get("cache_read_input_tokens"),
+        cache_write_tokens=usage.get("cache_creation_input_tokens"),
+        reasoning_tokens=thinking if isinstance(thinking, int) else None,
+    )
+
+
+def _openai_gpt_oss_invoke_usage(body: dict[str, Any]) -> UsageReport:
+    usage = body.get("usage") or {}
+    return UsageReport(
+        response_model=body.get("model"),
+        input_tokens=usage.get("prompt_tokens"),
+        output_tokens=usage.get("completion_tokens"),
+    )
+
+
+def _amazon_nova_invoke_usage(body: dict[str, Any]) -> UsageReport:
+    usage = body.get("usage") or {}
+    return UsageReport(
+        response_model=body.get("model") or body.get("modelId"),
+        input_tokens=usage.get("inputTokens"),
+        output_tokens=usage.get("outputTokens"),
+    )
+
+
+def _meta_llama_invoke_usage(body: dict[str, Any]) -> UsageReport:
+    return UsageReport(
+        input_tokens=body.get("prompt_token_count"),
+        output_tokens=body.get("generation_token_count"),
+    )
+
+
+def _mistral_invoke_usage(body: dict[str, Any]) -> UsageReport:
+    # Mistral on Bedrock historically did not include usage in the body; rely
+    # on the universal x-amzn-bedrock-*-token-count headers that the wrapper
+    # also reads.
+    return UsageReport()
+
+
+def _empty_invoke_usage(body: dict[str, Any]) -> UsageReport:
+    return UsageReport()
+
+
+# --- Universal Converse extractor (every provider returns the same shape) ----
+
+
+def extract_converse_usage_universal(response: dict[str, Any]) -> UsageReport:
+    usage = response.get("usage") or {}
+    return UsageReport(
+        input_tokens=usage.get("inputTokens"),
+        output_tokens=usage.get("outputTokens"),
+        cache_read_tokens=usage.get("cacheReadInputTokens"),
+        cache_write_tokens=usage.get("cacheWriteInputTokens"),
+    )
+
+
+def extract_converse_stream_event(event: dict[str, Any]) -> "UsageReport | None":
+    """Converse streams a terminal `metadata` event with the usage block."""
+    metadata = event.get("metadata")
+    if metadata is None:
+        return None
+    usage = metadata.get("usage") or {}
+    if not usage:
+        return None
+    return UsageReport(
+        input_tokens=usage.get("inputTokens"),
+        output_tokens=usage.get("outputTokens"),
+        cache_read_tokens=usage.get("cacheReadInputTokens"),
+        cache_write_tokens=usage.get("cacheWriteInputTokens"),
+    )
+
+
+# --- Universal InvokeModel stream extractor (final-chunk metrics) ------------
+
+
+def extract_invoke_model_stream_event_universal(
+    event: dict[str, Any],
+) -> "UsageReport | None":
+    """Pull token counts from `amazon-bedrock-invocationMetrics` in the last chunk.
+
+    Bedrock attaches these to the terminal chunk of every InvokeModel stream
+    regardless of provider — gives us non-zero coverage for any model.
+    """
+    decoded = _decode_invoke_model_chunk(event)
+    if decoded is None:
+        return None
+    metrics = decoded.get("amazon-bedrock-invocationMetrics")
+    if not metrics:
+        return None
+    return UsageReport(
+        input_tokens=metrics.get("inputTokenCount"),
+        output_tokens=metrics.get("outputTokenCount"),
+    )
+
+
+def _invoke_model_stream_event_anthropic(
+    event: dict[str, Any],
+) -> "UsageReport | None":
+    """Anthropic-specific extractor for invoke_model_with_response_stream chunks.
+
+    Anthropic emits SSE-shaped events inside chunk bytes:
+      message_start  → initial usage including cache_*
+      message_delta  → running output_tokens total
+    """
+    decoded = _decode_invoke_model_chunk(event)
+    if decoded is None:
+        return None
+    event_type = decoded.get("type")
+    if event_type == "message_start":
+        message = decoded.get("message") or {}
+        usage = message.get("usage") or {}
+        return UsageReport(
+            response_model=message.get("model"),
+            input_tokens=usage.get("input_tokens"),
+            output_tokens=usage.get("output_tokens"),
+            cache_read_tokens=usage.get("cache_read_input_tokens"),
+            cache_write_tokens=usage.get("cache_creation_input_tokens"),
+        )
+    if event_type == "message_delta":
+        usage = decoded.get("usage") or {}
+        out = usage.get("output_tokens")
+        if out is not None:
+            return UsageReport(output_tokens=out)
+    return None
+
+
+def _decode_invoke_model_chunk(event: dict[str, Any]) -> dict[str, Any] | None:
+    chunk = event.get("chunk") or {}
+    raw = chunk.get("bytes")
+    if not raw:
+        return None
+    if isinstance(raw, (bytes, bytearray)):
+        raw = bytes(raw).decode("utf-8", errors="replace")
+    try:
+        return json.loads(raw)
+    except (ValueError, TypeError):
+        return None
+
+
+# --- Universal HTTP-header extractor (for non-streaming unrecognized models) -
+
+
+def extract_universal_headers_usage(response: dict[str, Any]) -> UsageReport:
+    """Pull token counts from boto3's response['ResponseMetadata']['HTTPHeaders'].
+
+    Bedrock returns `x-amzn-bedrock-input-token-count` and
+    `x-amzn-bedrock-output-token-count` on every InvokeModel response — same
+    shape across providers. Used as the fallback when no per-provider handler
+    matches.
+    """
+    headers = (response.get("ResponseMetadata") or {}).get("HTTPHeaders") or {}
+    return UsageReport(
+        input_tokens=_try_int(headers.get("x-amzn-bedrock-input-token-count")),
+        output_tokens=_try_int(headers.get("x-amzn-bedrock-output-token-count")),
+    )
+
+
+def _try_int(value: Any) -> int | None:
+    if value is None:
+        return None
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return None
+
+
+# --- Request-payload capture (drives the M5 prompt analyzer) -----------------
+
+
+def _capture_anthropic_payload(body: dict[str, Any]) -> RequestPayload | None:
+    max_tokens = body.get("max_tokens")
+    if not isinstance(max_tokens, int):
+        max_tokens = None
+    return RequestPayload(
+        messages=body.get("messages"),
+        tools=body.get("tools"),
+        response_format=None,
+        max_tokens=max_tokens,
+        system=body.get("system"),
+    )
+
+
+def _capture_openai_payload(body: dict[str, Any]) -> RequestPayload | None:
+    max_tokens = body.get("max_tokens")
+    if not isinstance(max_tokens, int):
+        max_tokens = body.get("max_completion_tokens")
+        if not isinstance(max_tokens, int):
+            max_tokens = None
+    return RequestPayload(
+        messages=body.get("messages"),
+        tools=body.get("tools"),
+        response_format=body.get("response_format"),
+        max_tokens=max_tokens,
+        system=None,
+    )
+
+
+def _no_payload(_body: dict[str, Any]) -> RequestPayload | None:
+    return None
+
+
+# --- Body decode helpers (shared by per-provider extractors) ------------------
+
+
+def decode_invoke_model_body(raw_body: Any) -> dict[str, Any]:
+    """Decode the `body` kwarg of an InvokeModel call into a dict.
+
+    Customers usually pass `body=json.dumps({...})` or `body=b'{...}'`. We
+    always parse into a dict so handlers can reason about request kwargs (for
+    the analyzer) and so the universal stream parser shares the same shape.
+    """
+    if raw_body is None:
+        return {}
+    if isinstance(raw_body, (bytes, bytearray)):
+        raw_body = bytes(raw_body).decode("utf-8", errors="replace")
+    if isinstance(raw_body, str):
+        try:
+            decoded = json.loads(raw_body)
+        except (ValueError, TypeError):
+            return {}
+        return decoded if isinstance(decoded, dict) else {}
+    if isinstance(raw_body, dict):
+        return raw_body
+    return {}
+
+
+def consume_invoke_model_response_body(response: dict[str, Any]) -> dict[str, Any]:
+    """Read response['body'].read() once, replace it with a re-readable BytesIO,
+    and return the parsed JSON. Customers can still call .read() afterwards.
+    """
+    body_stream = response.get("body")
+    if body_stream is None:
+        return {}
+    reader = getattr(body_stream, "read", None)
+    if not callable(reader):
+        return {}
+    raw = reader()
+    response["body"] = _ReReadableStreamingBody(raw)
+    if isinstance(raw, (bytes, bytearray)):
+        raw = bytes(raw).decode("utf-8", errors="replace")
+    try:
+        decoded = json.loads(raw)
+    except (ValueError, TypeError):
+        return {}
+    return decoded if isinstance(decoded, dict) else {}
+
+
+class _ReReadableStreamingBody:
+    """Tiny stand-in for botocore.response.StreamingBody after we've consumed it.
+
+    We only need to satisfy the customer if they call `.read()` again.
+    Anything fancier (close, iter chunks) is best-effort.
+    """
+
+    def __init__(self, data: Any) -> None:
+        if isinstance(data, str):
+            data = data.encode("utf-8")
+        if not isinstance(data, (bytes, bytearray)):
+            data = b""
+        self._data = bytes(data)
+        self._buf = io.BytesIO(self._data)
+
+    def read(self, amt: int | None = None) -> bytes:
+        if amt is None:
+            return self._buf.read()
+        return self._buf.read(amt)
+
+    def close(self) -> None:
+        self._buf.close()
+
+    def iter_chunks(self) -> Any:
+        yield self._data
+
+
+# --- Handler registry --------------------------------------------------------
+
+
+ANTHROPIC_CLAUDE_HANDLER = BedrockModelHandler(
+    name="anthropic.claude",
+    extract_invoke_model_usage=_anthropic_invoke_usage,
+    extract_converse_usage=extract_converse_usage_universal,
+    extract_invoke_model_stream_event=_invoke_model_stream_event_anthropic,
+    capture_invoke_model_payload=_capture_anthropic_payload,
+)
+
+OPENAI_GPT_OSS_HANDLER = BedrockModelHandler(
+    name="openai.gpt-oss",
+    extract_invoke_model_usage=_openai_gpt_oss_invoke_usage,
+    extract_converse_usage=extract_converse_usage_universal,
+    extract_invoke_model_stream_event=extract_invoke_model_stream_event_universal,
+    capture_invoke_model_payload=_capture_openai_payload,
+)
+
+AMAZON_NOVA_HANDLER = BedrockModelHandler(
+    name="amazon.nova",
+    extract_invoke_model_usage=_amazon_nova_invoke_usage,
+    extract_converse_usage=extract_converse_usage_universal,
+    extract_invoke_model_stream_event=extract_invoke_model_stream_event_universal,
+    capture_invoke_model_payload=_no_payload,
+)
+
+META_LLAMA_HANDLER = BedrockModelHandler(
+    name="meta.llama",
+    extract_invoke_model_usage=_meta_llama_invoke_usage,
+    extract_converse_usage=extract_converse_usage_universal,
+    extract_invoke_model_stream_event=extract_invoke_model_stream_event_universal,
+    capture_invoke_model_payload=_no_payload,
+)
+
+MISTRAL_HANDLER = BedrockModelHandler(
+    name="mistral",
+    extract_invoke_model_usage=_mistral_invoke_usage,
+    extract_converse_usage=extract_converse_usage_universal,
+    extract_invoke_model_stream_event=extract_invoke_model_stream_event_universal,
+    capture_invoke_model_payload=_no_payload,
+)
+
+UNIVERSAL_FALLBACK_HANDLER = BedrockModelHandler(
+    name="<universal>",
+    extract_invoke_model_usage=_empty_invoke_usage,  # body parsing skipped — header path used instead
+    extract_converse_usage=extract_converse_usage_universal,
+    extract_invoke_model_stream_event=extract_invoke_model_stream_event_universal,
+    capture_invoke_model_payload=_no_payload,
+)
+
+
+# Order matters: longest provider+family prefix first so e.g.
+# `openai.gpt-oss-...` matches before any future shorter `openai.` entry.
+DISPATCH: list[tuple[str, BedrockModelHandler]] = [
+    ("anthropic.claude-", ANTHROPIC_CLAUDE_HANDLER),
+    ("openai.gpt-oss-", OPENAI_GPT_OSS_HANDLER),
+    ("openai.gpt-", OPENAI_GPT_OSS_HANDLER),
+    ("amazon.nova-", AMAZON_NOVA_HANDLER),
+    ("meta.llama", META_LLAMA_HANDLER),
+    ("mistral.", MISTRAL_HANDLER),
+]
+
+
+def lookup_handler(model_id: str) -> BedrockModelHandler:
+    """Strip cross-region prefix, find the longest matching provider+family
+    prefix, fall back to the universal handler.
+    """
+    stripped = strip_cross_region_prefix(model_id)
+    for prefix, handler in DISPATCH:
+        if stripped.startswith(prefix):
+            return handler
+    return UNIVERSAL_FALLBACK_HANDLER
+
+
+def merge_usage(into: UsageReport, partial: "UsageReport | None") -> None:
+    """Update `into` in place with non-None fields from `partial`."""
+    if partial is None:
+        return
+    for field_name in (
+        "response_model",
+        "input_tokens",
+        "output_tokens",
+        "cache_read_tokens",
+        "cache_write_tokens",
+        "reasoning_tokens",
+    ):
+        value = getattr(partial, field_name)
+        if value is not None:
+            setattr(into, field_name, value)