feat(providers): optional token-usage capture via usage_sink

Add an opt-in `usage_sink` dict to every process_* provider and the shared
call_openai_compatible_api / call_http_api_with_retry core. When supplied it is
populated in place with {prompt_tokens, completion_tokens, total_tokens} (plus a
native USD cost for OpenRouter, which sends usage:{include:true} only when a
sink is passed). Default path is unchanged: omit the sink and the request body,
network behavior, and list[str] return are byte-identical.

Also expose public extractors extract_chat_completions_usage and
extract_gemini_usage. Motivation: recover per-call token usage / cost for
consensus runs without monkeypatching requests.post / generate_content.

- anthropic: bespoke usage normalizer written to the documented Messages-API
  shape but NOT validated against the live API (flagged in-code; needs live
  tests and confirmation before it is relied on).
- tests: offline behavioral coverage + format-drift resilience (extra/missing/
  partial/malformed usage); default-path byte-identical test.
- CHANGELOG: note under Unreleased.

Author: tony-zhelonkin

python/CHANGELOG.md

@@ -2,6 +2,21 @@
 
 All notable changes to the Python implementation of mLLMCelltype will be documented in this file.
 
+## [Unreleased]
+
+### Added
+- Optional token-usage capture for provider calls. Every `process_*` provider function
+  (and the shared `call_openai_compatible_api` / `call_http_api_with_retry` core) now
+  accepts an opt-in `usage_sink` dict that is populated in place with
+  `{prompt_tokens, completion_tokens, total_tokens}` — plus a native `cost` (USD) for
+  OpenRouter, which opts in to `usage: {include: true}` only when a sink is supplied.
+- Public usage extractors `extract_chat_completions_usage` and `extract_gemini_usage`
+  (exported from `mllmcelltype.providers`) for callers that parse responses directly.
+
+### Notes
+- Default behavior is unchanged: omitting `usage_sink` leaves request shape and return
+  values byte-identical to prior releases.
+
 ## [2.0.5] - 2026-05-11
 
 ### Fixed

python/mllmcelltype/providers/__init__.py

@@ -2,8 +2,9 @@
 This package contains modules for interacting with various LLM providers."""
 
 from .anthropic import process_anthropic
+from .common import UsageSink, extract_chat_completions_usage
 from .deepseek import process_deepseek
-from .gemini import process_gemini
+from .gemini import extract_gemini_usage, process_gemini
 from .grok import process_grok
 from .minimax import process_minimax
 from .openai import process_openai
@@ -13,6 +14,9 @@
 from .zhipu import process_zhipu
 
 __all__ = [
+    "UsageSink",
+    "extract_chat_completions_usage",
+    "extract_gemini_usage",
     "process_anthropic",
     "process_deepseek",
     "process_gemini",

python/mllmcelltype/providers/anthropic.py

@@ -8,6 +8,7 @@
 
 from ..logger import write_log
 from .common import (
+    UsageSink,
     call_http_api_with_retry,
     ensure_api_key,
     resolve_endpoint_url,
@@ -64,8 +65,43 @@ def _parse_anthropic_response(content: dict[str, Any]) -> list[str]:
     return [line.rstrip(",") for line in lines]
 
 
+def _extract_anthropic_usage(content: dict[str, Any]) -> dict[str, Any] | None:
+    """Normalize Anthropic's ``usage`` block to the shared usage schema.
+
+    NOTE (honesty): written to Anthropic's documented Messages-API response shape
+    (``usage.input_tokens`` / ``usage.output_tokens``; ``total_tokens`` is derived
+    since Anthropic doesn't return it). I did NOT validate this against the live
+    Anthropic API — I only have OpenRouter + Gemini keys. This path still needs
+    live testing and confirmation against the real Anthropic API before it should
+    be relied on; the offline tests only pin the documented shape, not real
+    payloads. Please scrutinize; happy to split this normalizer into a follow-up
+    if you'd rather not carry unverified code. ``total_tokens`` stays ``None``
+    unless both token counts are present ints (no fabrication from partial data).
+    """
+    if not isinstance(content, dict):
+        return None
+    usage = content.get("usage")
+    if not isinstance(usage, dict):
+        return None
+
+    prompt_tokens = usage.get("input_tokens")
+    completion_tokens = usage.get("output_tokens")
+    total_tokens = None
+    if isinstance(prompt_tokens, int) and isinstance(completion_tokens, int):
+        total_tokens = prompt_tokens + completion_tokens
+    return {
+        "prompt_tokens": prompt_tokens,
+        "completion_tokens": completion_tokens,
+        "total_tokens": total_tokens,
+    }
+
+
 def process_anthropic(
-    prompt: str, model: str, api_key: str, base_url: str | None = None
+    prompt: str,
+    model: str,
+    api_key: str,
+    base_url: str | None = None,
+    usage_sink: UsageSink | None = None,
 ) -> list[str]:
     """Process request using Anthropic Claude models.
 
@@ -74,6 +110,7 @@ def process_anthropic(
         model: The model name (e.g., 'claude-opus-4-7', 'claude-sonnet-4-6')
         api_key: Anthropic API key
         base_url: Optional custom base URL
+        usage_sink: Optional dict populated in place with token usage.
 
     Returns:
         List[str]: Processed responses, one per cluster
@@ -111,4 +148,6 @@ def process_anthropic(
         retry_delay=2,
         timeout=30,
         request_json=False,
+        usage_sink=usage_sink,
+        usage_parser=_extract_anthropic_usage,
     )

python/mllmcelltype/providers/common.py

@@ -4,7 +4,7 @@
 
 import json
 import time
-from collections.abc import Callable
+from collections.abc import Callable, MutableMapping
 from typing import Any
 
 import requests
@@ -14,6 +14,45 @@
 
 ResponseParser = Callable[[dict[str, Any]], list[str]]
 
+# A caller-supplied dict that, when passed, is populated in place with token
+# usage for the call: prompt_tokens / completion_tokens / total_tokens, plus an
+# optional native `cost` (USD) when the provider returns one. None = no capture.
+UsageSink = MutableMapping[str, Any]
+
+# Maps a raw response payload to the normalized usage schema (or None if absent).
+UsageParser = Callable[[dict[str, Any]], dict[str, Any] | None]
+
+
+def extract_chat_completions_usage(content: dict[str, Any]) -> dict[str, Any] | None:
+    """Extract token usage from an OpenAI-compatible chat completions payload.
+
+    Returns a normalized ``{prompt_tokens, completion_tokens, total_tokens}`` dict
+    (plus ``cost`` when the provider includes a native USD cost, e.g. OpenRouter
+    when ``usage: {include: true}`` was requested), or ``None`` when no usage block
+    is present. Never raises on a malformed/absent ``usage`` block.
+
+    Resilient to provider format drift by design: only the three canonical keys
+    (plus ``cost``) are read; any extra/unknown fields are ignored. Missing keys
+    are reported as ``None`` rather than fabricated, so callers can tell "not
+    reported" from a real zero. Validated live against OpenRouter; the other
+    OpenAI-compatible providers (OpenAI, DeepSeek, Qwen, Grok, StepFun, Zhipu,
+    MiniMax) share this exact path but were not run against their native endpoints.
+    """
+    if not isinstance(content, dict):
+        return None
+    usage = content.get("usage")
+    if not isinstance(usage, dict):
+        return None
+
+    normalized: dict[str, Any] = {
+        "prompt_tokens": usage.get("prompt_tokens"),
+        "completion_tokens": usage.get("completion_tokens"),
+        "total_tokens": usage.get("total_tokens"),
+    }
+    if usage.get("cost") is not None:
+        normalized["cost"] = usage["cost"]
+    return normalized
+
 
 class NonRetryableProviderError(ValueError):
     """Provider error that should fail fast without retry."""
@@ -138,8 +177,16 @@ def call_http_api_with_retry(
     timeout: int = 30,
     request_json: bool = False,
     non_retry_exceptions: tuple[type[Exception], ...] = (),
+    usage_sink: UsageSink | None = None,
+    usage_parser: UsageParser = extract_chat_completions_usage,
 ) -> list[str]:
-    """Execute an HTTP API request with retry and unified error handling."""
+    """Execute an HTTP API request with retry and unified error handling.
+
+    When ``usage_sink`` is provided, it is populated in place with token usage
+    parsed from the successful response via ``usage_parser`` (default: the
+    OpenAI-compatible extractor). Leaving it ``None`` is byte-identical to the
+    prior behavior.
+    """
     write_log("Sending API request...")
 
     for attempt in range(max_retries):
@@ -190,6 +237,11 @@ def call_http_api_with_retry(
                 raise NonRetryableProviderError(
                     f"{provider_name} response parser returned {type(res).__name__}, expected list"
                 )
+            if usage_sink is not None:
+                usage = usage_parser(content)
+                if usage is not None:
+                    usage_sink.update(usage)
+
             normalized_res = [str(line) for line in res]
             write_log(f"Got response with {len(normalized_res)} lines")
             write_log(f"Raw response from {provider_name}:\n{normalized_res}", level="debug")
@@ -235,8 +287,14 @@ def call_openai_compatible_api(
     timeout: int = 30,
     request_json: bool = False,
     non_retry_exceptions: tuple[type[Exception], ...] = (),
+    usage_sink: UsageSink | None = None,
 ) -> list[str]:
-    """Execute a request against an OpenAI-compatible endpoint with retries."""
+    """Execute a request against an OpenAI-compatible endpoint with retries.
+
+    When ``usage_sink`` is provided, it is populated in place with the call's
+    token usage (see :func:`extract_chat_completions_usage`). Default ``None``
+    preserves prior behavior exactly.
+    """
     headers = {
         "Content-Type": "application/json",
         "Authorization": f"Bearer {api_key}",
@@ -258,4 +316,5 @@ def call_openai_compatible_api(
         timeout=timeout,
         request_json=request_json,
         non_retry_exceptions=non_retry_exceptions,
+        usage_sink=usage_sink,
     )

python/mllmcelltype/providers/deepseek.py

@@ -6,6 +6,7 @@
 
 from ..logger import write_log
 from .common import (
+    UsageSink,
     build_chat_completions_body,
     call_openai_compatible_api,
     ensure_api_key,
@@ -14,7 +15,11 @@
 
 
 def process_deepseek(
-    prompt: str, model: str, api_key: str, base_url: str | None = None
+    prompt: str,
+    model: str,
+    api_key: str,
+    base_url: str | None = None,
+    usage_sink: UsageSink | None = None,
 ) -> list[str]:
     """Process request using DeepSeek models.
 
@@ -23,6 +28,7 @@ def process_deepseek(
         model: The model name (e.g., 'deepseek-v4-flash', 'deepseek-v4-pro')
         api_key: DeepSeek API key
         base_url: Optional custom base URL
+        usage_sink: Optional dict populated in place with token usage.
 
     Returns:
         List[str]: Processed responses, one per cluster
@@ -52,4 +58,5 @@ def process_deepseek(
         retry_delay=3,
         timeout=90,
         request_json=True,
+        usage_sink=usage_sink,
     )

python/mllmcelltype/providers/gemini.py

@@ -6,7 +6,24 @@
 from typing import Any
 
 from ..logger import write_log
-from .common import ensure_api_key
+from .common import UsageSink, ensure_api_key
+
+
+def extract_gemini_usage(response: Any) -> dict[str, Any] | None:
+    """Extract token usage from a Gemini SDK response's ``usage_metadata``.
+
+    Returns the shared ``{prompt_tokens, completion_tokens, total_tokens}`` schema,
+    or ``None`` when no usage metadata is present. Never raises on absent/odd
+    metadata.
+    """
+    metadata = getattr(response, "usage_metadata", None)
+    if metadata is None:
+        return None
+    return {
+        "prompt_tokens": getattr(metadata, "prompt_token_count", None),
+        "completion_tokens": getattr(metadata, "candidates_token_count", None),
+        "total_tokens": getattr(metadata, "total_token_count", None),
+    }
 
 
 def _parse_gemini_response(response: Any) -> list[str]:
@@ -26,7 +43,11 @@ def _parse_gemini_response(response: Any) -> list[str]:
 
 
 def process_gemini(
-    prompt: str, model: str, api_key: str, base_url: str | None = None
+    prompt: str,
+    model: str,
+    api_key: str,
+    base_url: str | None = None,
+    usage_sink: UsageSink | None = None,
 ) -> list[str]:
     """Process request using Google Gemini models.
 
@@ -35,6 +56,7 @@ def process_gemini(
         model: The model name (e.g., 'gemini-3.1-pro-preview', 'gemini-3-flash-preview', 'gemini-3.1-flash-lite')
         api_key: Google API key
         base_url: Optional custom base URL (Note: Gemini uses SDK, base_url may not be applicable)
+        usage_sink: Optional dict populated in place with token usage.
 
     Returns:
         List[str]: Processed responses, one per cluster
@@ -85,6 +107,10 @@ def process_gemini(
             )
 
             result = _parse_gemini_response(response)
+            if usage_sink is not None:
+                usage = extract_gemini_usage(response)
+                if usage is not None:
+                    usage_sink.update(usage)
             write_log(f"Got response with {len(result)} lines")
             write_log(f"Raw response from Gemini:\n{result}", level="debug")
             return result

python/mllmcelltype/providers/grok.py

@@ -6,6 +6,7 @@
 
 from ..logger import write_log
 from .common import (
+    UsageSink,
     build_chat_completions_body,
     call_openai_compatible_api,
     ensure_api_key,
@@ -14,7 +15,11 @@
 
 
 def process_grok(
-    prompt: str, model: str, api_key: str, base_url: str | None = None
+    prompt: str,
+    model: str,
+    api_key: str,
+    base_url: str | None = None,
+    usage_sink: UsageSink | None = None,
 ) -> list[str]:
     """Process request using Grok models from xAI.
 
@@ -23,6 +28,7 @@ def process_grok(
         model: The model name (e.g., 'grok-4.3', 'grok-4.3-latest')
         api_key: xAI API key
         base_url: Optional custom base URL
+        usage_sink: Optional dict populated in place with token usage.
 
     Returns:
         List[str]: Processed responses, one per cluster
@@ -42,4 +48,5 @@ def process_grok(
         url=url,
         body=body,
         post_func=requests.post,
+        usage_sink=usage_sink,
     )

python/mllmcelltype/providers/minimax.py

@@ -11,6 +11,7 @@
 from ..url_utils import get_working_minimax_endpoint
 from .common import (
     NonRetryableProviderError,
+    UsageSink,
     build_chat_completions_body,
     call_openai_compatible_api,
     ensure_api_key,
@@ -43,7 +44,11 @@ def _parse_minimax_response(content: dict[str, Any]) -> list[str]:
 
 
 def process_minimax(
-    prompt: str, model: str, api_key: str, base_url: str | None = None
+    prompt: str,
+    model: str,
+    api_key: str,
+    base_url: str | None = None,
+    usage_sink: UsageSink | None = None,
 ) -> list[str]:
     """Process request using MiniMax models.
 
@@ -52,6 +57,7 @@ def process_minimax(
         model: The model name (e.g., 'MiniMax-M2.7', 'MiniMax-M2.7-highspeed', 'MiniMax-M2.5')
         api_key: MiniMax API key
         base_url: Optional custom base URL
+        usage_sink: Optional dict populated in place with token usage.
 
     Returns:
         List[str]: Processed responses, one per cluster
@@ -87,4 +93,5 @@ def process_minimax(
         post_func=requests.post,
         response_parser=_parse_minimax_response,
         non_retry_exceptions=(NonRetryableProviderError,),
+        usage_sink=usage_sink,
     )