generic-ml-wrapper 0.1.0__py3-none-any.whl

generic_ml_wrapper/adapter/outbound/gateway/anthropic_sse.py

@@ -0,0 +1,148 @@
+# SPDX-FileCopyrightText: 2026 Daniel Slobozian
+# SPDX-License-Identifier: Apache-2.0
+"""Extract per-turn token usage from an Anthropic Messages SSE stream.
+
+Anthropic streams a turn as Server-Sent Events. Usage is split across two events:
+``message_start`` carries ``input_tokens`` and the model; the final
+``message_delta`` carries the cumulative ``output_tokens``. This reads both from
+the raw ``data:`` lines a relay tees off the response, so the relay can record a
+turn without buffering or re-encoding the body.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, cast
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+_DATA_PREFIX = "data:"
+
+
+@dataclass(frozen=True)
+class StreamUsage:
+    """The token usage read off one Anthropic turn (session-agnostic).
+
+    Attributes:
+        input_tokens: Fresh prompt tokens from ``message_start``.
+        output_tokens: Completion tokens from the final ``message_delta``.
+        model: The serving model, or ``None`` if the stream omitted it.
+        cache_creation_tokens: Prompt tokens written to the cache this turn.
+        cache_read_tokens: Prompt tokens served from the cache this turn.
+        turn_id: The provider's id for this turn, or ``None`` if unreported.
+    """
+
+    input_tokens: int
+    output_tokens: int
+    model: str | None
+    cache_creation_tokens: int = 0
+    cache_read_tokens: int = 0
+    turn_id: str | None = None
+
+
+def read_usage(text: str) -> StreamUsage | None:
+    """Read a turn's usage from a raw Anthropic response body (SSE or JSON).
+
+    Claude streams (SSE) interactively but may answer a non-interactive request
+    with a single JSON envelope; this tries the streaming shape first, then the
+    non-streaming one.
+
+    Args:
+        text: The decoded response body.
+
+    Returns:
+        The turn's usage, or ``None`` if none was found.
+    """
+    streaming = extract_usage(text.splitlines())
+    return streaming if streaming is not None else _from_json(text)
+
+
+def _from_json(text: str) -> StreamUsage | None:
+    try:
+        decoded: object = json.loads(text)
+    except (json.JSONDecodeError, ValueError):
+        return None
+    if not isinstance(decoded, dict):
+        return None
+    envelope = cast("dict[str, object]", decoded)
+    usage = envelope.get("usage")
+    input_tokens = _usage_token(usage, "input_tokens")
+    output_tokens = _usage_token(usage, "output_tokens")
+    if input_tokens is None and output_tokens is None:
+        return None
+    return StreamUsage(
+        input_tokens or 0,
+        output_tokens or 0,
+        _as_str(envelope.get("model")),
+        _usage_token(usage, "cache_creation_input_tokens") or 0,
+        _usage_token(usage, "cache_read_input_tokens") or 0,
+        _as_str(envelope.get("id")),
+    )
+
+
+def extract_usage(lines: Iterable[str]) -> StreamUsage | None:
+    """Read token usage from the ``data:`` lines of an Anthropic SSE stream.
+
+    Args:
+        lines: The response's lines (``event:``/blank lines are ignored).
+
+    Returns:
+        The turn's usage, or ``None`` if no usage was seen (e.g. a non-stream
+        error response).
+    """
+    input_tokens: int | None = None
+    output_tokens: int | None = None
+    cache_creation: int | None = None
+    cache_read: int | None = None
+    model: str | None = None
+    turn_id: str | None = None
+    for raw in lines:
+        event = _data_event(raw)
+        if event is None:
+            continue
+        kind = event.get("type")
+        if kind == "message_start":
+            message = event.get("message")
+            if isinstance(message, dict):
+                fields = cast("dict[str, object]", message)
+                model = _as_str(fields.get("model"))
+                turn_id = _as_str(fields.get("id"))
+                usage = fields.get("usage")
+                input_tokens = _usage_token(usage, "input_tokens")
+                cache_creation = _usage_token(usage, "cache_creation_input_tokens")
+                cache_read = _usage_token(usage, "cache_read_input_tokens")
+        elif kind == "message_delta":
+            seen = _usage_token(event.get("usage"), "output_tokens")
+            if seen is not None:
+                output_tokens = seen
+    if input_tokens is None and output_tokens is None:
+        return None
+    return StreamUsage(
+        input_tokens or 0, output_tokens or 0, model, cache_creation or 0, cache_read or 0, turn_id
+    )
+
+
+def _data_event(raw: str) -> dict[str, object] | None:
+    line = raw.strip()
+    if not line.startswith(_DATA_PREFIX):
+        return None
+    try:
+        decoded: object = json.loads(line[len(_DATA_PREFIX) :].strip())
+    except (json.JSONDecodeError, ValueError):
+        return None
+    return cast("dict[str, object]", decoded) if isinstance(decoded, dict) else None
+
+
+def _usage_token(usage: object, key: str) -> int | None:
+    if not isinstance(usage, dict):
+        return None
+    value = cast("dict[str, object]", usage).get(key)
+    if isinstance(value, bool) or not isinstance(value, int):
+        return None
+    return value
+
+
+def _as_str(value: object) -> str | None:
+    return value if isinstance(value, str) and value else None

generic_ml_wrapper/adapter/outbound/gateway/openai_chat.py

@@ -0,0 +1,112 @@
+# SPDX-FileCopyrightText: 2026 Daniel Slobozian
+# SPDX-License-Identifier: Apache-2.0
+"""Extract per-turn token usage from an OpenAI Chat Completions response (vibe/Mistral).
+
+Mistral's vibe CLI speaks the OpenAI Chat Completions API. Usage arrives in a
+standard ``usage`` block — in the final SSE chunk when streaming (that chunk
+carries ``usage`` with an empty ``choices``), or in the single JSON envelope of a
+non-streaming response. ``prompt_tokens`` counts the whole input including any
+cache reads, so fresh input is ``prompt_tokens - cached_tokens``.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import cast
+
+from generic_ml_wrapper.adapter.outbound.gateway.anthropic_sse import StreamUsage
+
+_DATA_PREFIX = "data:"
+
+
+def read_usage(text: str) -> StreamUsage | None:
+    """Read a turn's usage from an OpenAI Chat Completions response body.
+
+    Tries the streaming shape (SSE ``data:`` lines) first, then a single
+    non-streaming JSON envelope.
+
+    Args:
+        text: The decoded response body.
+
+    Returns:
+        The turn's usage, or ``None`` if none was found.
+    """
+    streaming = _from_sse(text)
+    return streaming if streaming is not None else _from_json(text)
+
+
+def _from_sse(text: str) -> StreamUsage | None:
+    model: str | None = None
+    turn_id: str | None = None
+    tokens: tuple[int, int, int] | None = None
+    for raw in text.splitlines():
+        chunk = _data_event(raw)
+        if chunk is None:
+            continue
+        model = _as_str(chunk.get("model")) or model
+        turn_id = _as_str(chunk.get("id")) or turn_id
+        seen = _tokens(chunk.get("usage"))
+        if seen is not None:
+            tokens = seen
+    if tokens is None:
+        return None
+    fresh, output, cached = tokens
+    return StreamUsage(fresh, output, model, 0, cached, turn_id)
+
+
+def _from_json(text: str) -> StreamUsage | None:
+    try:
+        decoded: object = json.loads(text)
+    except (json.JSONDecodeError, ValueError):
+        return None
+    if not isinstance(decoded, dict):
+        return None
+    envelope = cast("dict[str, object]", decoded)
+    tokens = _tokens(envelope.get("usage"))
+    if tokens is None:
+        return None
+    fresh, output, cached = tokens
+    return StreamUsage(
+        fresh, output, _as_str(envelope.get("model")), 0, cached, _as_str(envelope.get("id"))
+    )
+
+
+def _data_event(raw: str) -> dict[str, object] | None:
+    line = raw.strip()
+    if not line.startswith(_DATA_PREFIX):
+        return None
+    payload = line[len(_DATA_PREFIX) :].strip()
+    if payload == "[DONE]":
+        return None
+    try:
+        decoded: object = json.loads(payload)
+    except (json.JSONDecodeError, ValueError):
+        return None
+    return cast("dict[str, object]", decoded) if isinstance(decoded, dict) else None
+
+
+def _tokens(usage: object) -> tuple[int, int, int] | None:
+    """Return ``(fresh_input, output, cache_read)`` from a ``usage`` block, or ``None``."""
+    if not isinstance(usage, dict):
+        return None
+    fields = cast("dict[str, object]", usage)
+    prompt = _as_count(fields.get("prompt_tokens"))
+    completion = _as_count(fields.get("completion_tokens"))
+    cached = _as_count(_get(fields.get("prompt_tokens_details"), "cached_tokens"))
+    return (max(prompt - cached, 0), completion, cached)  # prompt_tokens includes cache reads
+
+
+def _get(mapping: object, key: str) -> object:
+    if not isinstance(mapping, dict):
+        return None
+    return cast("dict[str, object]", mapping).get(key)
+
+
+def _as_count(value: object) -> int:
+    if isinstance(value, bool) or not isinstance(value, int):
+        return 0
+    return value
+
+
+def _as_str(value: object) -> str | None:
+    return value if isinstance(value, str) and value else None

generic_ml_wrapper/adapter/outbound/gateway/openai_responses.py

@@ -0,0 +1,85 @@
+# SPDX-FileCopyrightText: 2026 Daniel Slobozian
+# SPDX-License-Identifier: Apache-2.0
+"""Extract per-turn token usage from an OpenAI Responses API SSE stream (Codex).
+
+Codex streams a turn as Server-Sent Events and reports usage once, in the final
+``response.completed`` event's ``response.usage``. Unlike Anthropic, its
+``input_tokens`` is the TOTAL input including cache reads, so fresh input is
+``input_tokens - cached_tokens``.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import cast
+
+from generic_ml_wrapper.adapter.outbound.gateway.anthropic_sse import StreamUsage
+
+_DATA_PREFIX = "data:"
+
+
+def read_usage(text: str) -> StreamUsage | None:
+    """Read a turn's usage from a Codex Responses API SSE stream.
+
+    Args:
+        text: The decoded response body.
+
+    Returns:
+        The turn's usage, or ``None`` if no ``response.completed`` usage was seen.
+    """
+    for raw in text.splitlines():
+        line = raw.strip()
+        if not line.startswith(_DATA_PREFIX):
+            continue
+        try:
+            decoded: object = json.loads(line[len(_DATA_PREFIX) :].strip())
+        except (json.JSONDecodeError, ValueError):
+            continue
+        if not isinstance(decoded, dict):
+            continue
+        event = cast("dict[str, object]", decoded)
+        if event.get("type") != "response.completed":
+            continue
+        usage = _usage(event.get("response"))
+        if usage is not None:
+            return usage
+    return None
+
+
+def _usage(response: object) -> StreamUsage | None:
+    if not isinstance(response, dict):
+        return None
+    fields = cast("dict[str, object]", response)
+    usage = fields.get("usage")
+    if not isinstance(usage, dict):
+        return None
+    tokens = cast("dict[str, object]", usage)
+    total_input = _as_count(tokens.get("input_tokens"))
+    output = _as_count(tokens.get("output_tokens"))
+    details = tokens.get("input_tokens_details")
+    cached = _as_count(_get(details, "cached_tokens"))
+    cache_write = _as_count(_get(details, "cache_write_tokens"))
+    return StreamUsage(
+        input_tokens=max(total_input - cached, 0),  # codex input_tokens includes cache reads
+        output_tokens=output,
+        model=_as_str(fields.get("model")),
+        cache_creation_tokens=cache_write,
+        cache_read_tokens=cached,
+        turn_id=_as_str(fields.get("id")),
+    )
+
+
+def _get(mapping: object, key: str) -> object:
+    if not isinstance(mapping, dict):
+        return None
+    return cast("dict[str, object]", mapping).get(key)
+
+
+def _as_count(value: object) -> int:
+    if isinstance(value, bool) or not isinstance(value, int):
+        return 0
+    return value
+
+
+def _as_str(value: object) -> str | None:
+    return value if isinstance(value, str) and value else None