diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 29f9f00b..51a6aa9e 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -381,6 +381,11 @@ behavior matches shared transport policy. Provider-specific gateway quirks, such as Cohere's supported `reasoning_effort` values, GitHub's API headers/catalog filtering, Hugging Face's disabled prior reasoning replay, and unsupported compatibility fields, stay in that provider package. +The OpenAI-chat transport owns standard streamed usage handling: it requests +`stream_options.include_usage`, consumes provider `prompt_tokens` and +`completion_tokens` when present, and falls back to local estimates when +providers omit or reject optional usage metadata. Provider modules only own true +usage quirks such as DeepSeek prompt-cache counters. ### Adding A Provider diff --git a/core/anthropic/streaming/ledger.py b/core/anthropic/streaming/ledger.py index 241e3733..cae5f80f 100644 --- a/core/anthropic/streaming/ledger.py +++ b/core/anthropic/streaming/ledger.py @@ -210,10 +210,13 @@ class AnthropicStreamLedger: stop_reason: str, output_tokens: int | None, *, + input_tokens: int | None = None, usage_fields: Mapping[str, int] | None = None, ) -> str: self.stop_reason = stop_reason - safe_in = _safe_usage_int(self.input_tokens) + safe_in = _safe_usage_int( + self.input_tokens if input_tokens is None else input_tokens + ) safe_out = output_tokens if isinstance(output_tokens, int) else 0 usage = {"input_tokens": safe_in, "output_tokens": safe_out} if usage_fields: diff --git a/providers/deepseek/client.py b/providers/deepseek/client.py index b6fffc15..2c991176 100644 --- a/providers/deepseek/client.py +++ b/providers/deepseek/client.py @@ -1,11 +1,11 @@ """DeepSeek provider implementation (OpenAI-compatible Chat Completions).""" -from collections.abc import Mapping from typing import Any from providers.base import ProviderConfig from providers.defaults import DEEPSEEK_DEFAULT_BASE from providers.transports.openai_chat import OpenAIChatTransport +from providers.transports.openai_chat.usage import usage_int from .compat import build_deepseek_request_body @@ -31,24 +31,10 @@ class DeepSeekProvider(OpenAIChatTransport): def _anthropic_usage_fields(self, usage_info: Any) -> dict[str, int]: usage_fields: dict[str, int] = {} - cache_hit_tokens = _usage_int(usage_info, "prompt_cache_hit_tokens") + cache_hit_tokens = usage_int(usage_info, "prompt_cache_hit_tokens") if cache_hit_tokens is not None: usage_fields["cache_read_input_tokens"] = cache_hit_tokens - cache_miss_tokens = _usage_int(usage_info, "prompt_cache_miss_tokens") + cache_miss_tokens = usage_int(usage_info, "prompt_cache_miss_tokens") if cache_miss_tokens is not None: usage_fields["cache_creation_input_tokens"] = cache_miss_tokens return usage_fields - - -def _usage_int(usage_info: Any, key: str) -> int | None: - if usage_info is None: - return None - if isinstance(usage_info, Mapping): - value = usage_info.get(key) - else: - value = getattr(usage_info, key, None) - if value is None: - extra = getattr(usage_info, "model_extra", None) - if isinstance(extra, Mapping): - value = extra.get(key) - return value if isinstance(value, int) else None diff --git a/providers/deepseek/compat.py b/providers/deepseek/compat.py index 5448744b..f755332f 100644 --- a/providers/deepseek/compat.py +++ b/providers/deepseek/compat.py @@ -428,10 +428,6 @@ def _request_from_dict(request_data: Any, data: dict[str, Any]) -> Any: def _apply_deepseek_chat_extras( body: dict[str, Any], _request_data: Any, thinking_enabled: bool ) -> None: - stream_options = body.setdefault("stream_options", {}) - if isinstance(stream_options, dict): - stream_options["include_usage"] = True - if not thinking_enabled or body.get("model") == "deepseek-reasoner": return extra_body = body.setdefault("extra_body", {}) diff --git a/providers/transports/openai_chat/stream.py b/providers/transports/openai_chat/stream.py index be719e63..e24de9fd 100644 --- a/providers/transports/openai_chat/stream.py +++ b/providers/transports/openai_chat/stream.py @@ -31,6 +31,7 @@ from .tool_calls import ( iter_heuristic_tool_use_sse, tool_call_extra_content, ) +from .usage import request_stream_usage, usage_int class OpenAIChatStreamAdapter: @@ -79,6 +80,7 @@ class OpenAIChatStreamAdapter: body = self._transport._build_request_body( self._request, thinking_enabled=self._thinking_enabled ) + request_stream_usage(body) thinking_enabled = self._transport._is_thinking_enabled( self._request, self._thinking_enabled ) @@ -113,8 +115,9 @@ class OpenAIChatStreamAdapter: stream_opened = True tool_argument_aliases = self._transport._tool_argument_aliases(body) async for chunk in stream: - if getattr(chunk, "usage", None): - usage_info = chunk.usage + chunk_usage = getattr(chunk, "usage", None) + if chunk_usage is not None: + usage_info = chunk_usage if not chunk.choices: continue @@ -350,24 +353,22 @@ class OpenAIChatStreamAdapter: for event in hold_events(ledger.close_all_blocks()): yield event - completion = ( - getattr(usage_info, "completion_tokens", None) - if usage_info is not None - else None - ) + completion = usage_int(usage_info, "completion_tokens") if isinstance(completion, int): output_tokens = completion else: output_tokens = ledger.estimate_output_tokens() - if usage_info and hasattr(usage_info, "prompt_tokens"): - provider_input = usage_info.prompt_tokens - if isinstance(provider_input, int): - logger.debug( - "TOKEN_ESTIMATE: our={} provider={} diff={:+d}", - self._input_tokens, - provider_input, - provider_input - self._input_tokens, - ) + provider_input = usage_int(usage_info, "prompt_tokens") + if provider_input is not None: + logger.debug( + "TOKEN_ESTIMATE: our={} provider={} diff={:+d}", + self._input_tokens, + provider_input, + provider_input - self._input_tokens, + ) + input_tokens = ( + provider_input if provider_input is not None else self._input_tokens + ) trace_event( stage="provider", event="provider.response.completed", @@ -375,12 +376,14 @@ class OpenAIChatStreamAdapter: provider=tag, finish_reason=(None if finish_reason is None else str(finish_reason)), output_tokens=output_tokens, + prompt_tokens=input_tokens, prompt_tokens_estimate=self._input_tokens, ) for event in hold_event( ledger.message_delta( ledger.final_stop_reason(map_stop_reason(finish_reason)), output_tokens, + input_tokens=input_tokens, usage_fields=self._transport._anthropic_usage_fields(usage_info), ) ): diff --git a/providers/transports/openai_chat/transport.py b/providers/transports/openai_chat/transport.py index e3726c17..ae95fcfe 100644 --- a/providers/transports/openai_chat/transport.py +++ b/providers/transports/openai_chat/transport.py @@ -20,6 +20,7 @@ from providers.rate_limit import GlobalRateLimiter from .output_cap import clamp_output_tokens, parse_output_token_cap from .stream import OpenAIChatStreamAdapter +from .usage import clone_without_stream_usage, is_stream_usage_rejection class OpenAIChatTransport(BaseProvider): @@ -117,26 +118,51 @@ class OpenAIChatTransport(BaseProvider): return {} async def _create_stream(self, body: dict) -> tuple[Any, dict]: - """Create a streaming chat completion, optionally retrying once.""" + """Create a streaming chat completion with bounded request fallbacks.""" body = self._apply_learned_output_cap(body) - try: - create_body = self._prepare_create_body(body) - stream = await self._global_rate_limiter.execute_with_retry( - self._client.chat.completions.create, **create_body, stream=True - ) - return stream, body - except Exception as error: - retry_body = self._retry_body_for_output_cap(error, body) - if retry_body is None: - retry_body = self._get_retry_request_body(error, body) - if retry_body is None: - raise + used_retry_kinds: set[str] = set() - create_retry_body = self._prepare_create_body(retry_body) - stream = await self._global_rate_limiter.execute_with_retry( - self._client.chat.completions.create, **create_retry_body, stream=True - ) - return stream, retry_body + while True: + try: + create_body = self._prepare_create_body(body) + stream = await self._global_rate_limiter.execute_with_retry( + self._client.chat.completions.create, **create_body, stream=True + ) + return stream, body + except Exception as error: + retry_body = self._next_create_retry_body(error, body, used_retry_kinds) + if retry_body is None: + raise + body = retry_body + + def _next_create_retry_body( + self, + error: Exception, + body: dict, + used_retry_kinds: set[str], + ) -> dict | None: + retry_body = self._retry_body_for_output_cap(error, body) + if retry_body is not None: + return retry_body + + if "stream_usage" not in used_retry_kinds and is_stream_usage_rejection(error): + retry_body = clone_without_stream_usage(body) + if retry_body is not None: + used_retry_kinds.add("stream_usage") + logger.warning( + "{}_STREAM: retrying without stream_options.include_usage " + "after upstream rejection", + self._provider_name, + ) + return retry_body + + if "provider_specific" not in used_retry_kinds: + retry_body = self._get_retry_request_body(error, body) + if retry_body is not None: + used_retry_kinds.add("provider_specific") + return retry_body + + return None def _apply_learned_output_cap(self, body: dict) -> dict: """Clamp output tokens to a previously learned cap for this model.""" diff --git a/providers/transports/openai_chat/usage.py b/providers/transports/openai_chat/usage.py new file mode 100644 index 00000000..4be8dff1 --- /dev/null +++ b/providers/transports/openai_chat/usage.py @@ -0,0 +1,98 @@ +"""OpenAI-chat streamed usage request and extraction helpers.""" + +import json +from collections.abc import Mapping +from typing import Any + +import openai + +_USAGE_OPTION_KEYS = ("stream_options", "include_usage") +_USAGE_REJECTION_WORDS = ( + "unsupported", + "not supported", + "unknown", + "unrecognized", + "unexpected", + "invalid", + "extra", + "forbidden", + "not permitted", +) + + +def request_stream_usage(body: dict[str, Any]) -> None: + """Ask an OpenAI-compatible streaming endpoint for its final usage chunk.""" + stream_options = body.get("stream_options") + if stream_options is None: + body["stream_options"] = {"include_usage": True} + return + if isinstance(stream_options, dict): + stream_options["include_usage"] = True + + +def clone_without_stream_usage(body: dict[str, Any]) -> dict[str, Any] | None: + """Return a clone with only ``include_usage`` removed from stream options.""" + stream_options = body.get("stream_options") + if not isinstance(stream_options, dict): + return None + if "include_usage" not in stream_options: + return None + + retry_body = dict(body) + retry_stream_options = dict(stream_options) + retry_stream_options.pop("include_usage", None) + if retry_stream_options: + retry_body["stream_options"] = retry_stream_options + else: + retry_body.pop("stream_options", None) + return retry_body + + +def is_stream_usage_rejection(error: Exception) -> bool: + """Return whether upstream rejected the optional streamed-usage request.""" + if not _is_bad_request_like(error): + return False + text = _error_text(error) + if not any(key in text for key in _USAGE_OPTION_KEYS): + return False + return any(word in text for word in _USAGE_REJECTION_WORDS) + + +def usage_int(usage_info: Any, key: str) -> int | None: + """Extract an integer usage field from OpenAI SDK objects or plain dicts.""" + if usage_info is None: + return None + if isinstance(usage_info, Mapping): + value = usage_info.get(key) + else: + value = getattr(usage_info, key, None) + if value is None: + extra = getattr(usage_info, "model_extra", None) + if isinstance(extra, Mapping): + value = extra.get(key) + return value if isinstance(value, int) and not isinstance(value, bool) else None + + +def _is_bad_request_like(error: Exception) -> bool: + if isinstance(error, openai.BadRequestError): + return True + status = getattr(error, "status_code", None) + if status is None: + response = getattr(error, "response", None) + status = ( + getattr(response, "status_code", None) if response is not None else None + ) + return status in (400, 422) + + +def _error_text(error: Exception) -> str: + parts = [str(error)] + body = getattr(error, "body", None) + if body is not None: + parts.append(json.dumps(body, default=str)) + response = getattr(error, "response", None) + if response is not None: + text = getattr(response, "text", None) + if isinstance(text, str) and text: + parts.append(text) + return " ".join(parts).lower() diff --git a/pyproject.toml b/pyproject.toml index c36c375b..829bceee 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "free-claude-code" -version = "3.4.8" +version = "3.4.9" description = "Middleware between Claude Code CLI (Anthropic API) and NVIDIA NIM" readme = "README.md" requires-python = ">=3.14.0" diff --git a/tests/providers/test_deepseek.py b/tests/providers/test_deepseek.py index 7a976861..fd560c87 100644 --- a/tests/providers/test_deepseek.py +++ b/tests/providers/test_deepseek.py @@ -82,7 +82,7 @@ def test_build_request_body_openai_chat_shape(deepseek_provider): assert body["messages"][1]["role"] == "user" assert body["messages"][1] == {"role": "user", "content": "Hello"} assert body["max_tokens"] == 100 - assert body["stream_options"] == {"include_usage": True} + assert "stream_options" not in body def test_build_request_body_default_max_tokens(deepseek_provider): @@ -188,7 +188,7 @@ def test_build_request_body_respects_global_thinking_disable(): ) body = provider._build_request_body(request) assert "extra_body" not in body - assert body["stream_options"] == {"include_usage": True} + assert "stream_options" not in body def test_preserve_unsigned_thinking_when_thinking_on(deepseek_provider): @@ -714,7 +714,7 @@ async def test_stream_uses_chat_completions_and_maps_cache_usage(deepseek_provid event.data["usage"] for event in parsed if event.event == "message_delta" ) assert usage == { - "input_tokens": 7, + "input_tokens": 30, "output_tokens": 3, "cache_read_input_tokens": 10, "cache_creation_input_tokens": 20, diff --git a/tests/providers/test_openai_chat_usage.py b/tests/providers/test_openai_chat_usage.py new file mode 100644 index 00000000..d07995e8 --- /dev/null +++ b/tests/providers/test_openai_chat_usage.py @@ -0,0 +1,218 @@ +"""OpenAI-chat streamed usage helper tests.""" + +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, patch + +import openai +import pytest +from httpx import Request, Response + +from core.anthropic.stream_contracts import parse_sse_text +from providers.base import ProviderConfig +from providers.rate_limit import GlobalRateLimiter +from providers.transports.openai_chat import OpenAIChatTransport +from providers.transports.openai_chat.usage import ( + clone_without_stream_usage, + is_stream_usage_rejection, + request_stream_usage, + usage_int, +) + + +class _UsageTestProvider(OpenAIChatTransport): + def __init__(self): + super().__init__( + ProviderConfig( + api_key="test_key", + base_url="https://provider.example/v1", + rate_limit=100, + rate_window=60, + ), + provider_name="USAGE_TEST", + base_url="https://provider.example/v1", + api_key="test_key", + ) + + def _build_request_body( + self, request: Any, thinking_enabled: bool | None = None + ) -> dict: + return {"model": request.model, "messages": [{"role": "user", "content": "x"}]} + + +def _bad_request(message: str, body: object | None = None) -> openai.BadRequestError: + response = Response( + 400, + request=Request("POST", "https://provider.example/v1/chat/completions"), + ) + return openai.BadRequestError(message, response=response, body=body) + + +async def _stream(chunks): + for chunk in chunks: + yield chunk + + +def _chunk( + *, + content: str | None = None, + finish_reason: str | None = None, + usage: Any = None, +): + if content is None and finish_reason is None: + return SimpleNamespace(choices=[], usage=usage) + return SimpleNamespace( + choices=[ + SimpleNamespace( + delta=SimpleNamespace( + content=content, + reasoning_content=None, + tool_calls=None, + ), + finish_reason=finish_reason, + ) + ], + usage=usage, + ) + + +def test_request_stream_usage_adds_stream_options_when_absent(): + body = {"model": "m"} + + request_stream_usage(body) + + assert body["stream_options"] == {"include_usage": True} + + +def test_request_stream_usage_preserves_existing_stream_options(): + stream_options = {"foo": "bar"} + body = {"model": "m", "stream_options": stream_options} + + request_stream_usage(body) + + assert body["stream_options"] == {"foo": "bar", "include_usage": True} + assert body["stream_options"] is stream_options + + +def test_clone_without_stream_usage_removes_only_include_usage(): + body = { + "model": "m", + "stream_options": {"foo": "bar", "include_usage": True}, + } + + retry_body = clone_without_stream_usage(body) + + assert retry_body == {"model": "m", "stream_options": {"foo": "bar"}} + assert body["stream_options"] == {"foo": "bar", "include_usage": True} + + +def test_clone_without_stream_usage_drops_empty_stream_options(): + body = {"model": "m", "stream_options": {"include_usage": True}} + + retry_body = clone_without_stream_usage(body) + + assert retry_body == {"model": "m"} + + +def test_usage_int_reads_dict_object_and_model_extra(): + assert usage_int({"prompt_tokens": 11}, "prompt_tokens") == 11 + assert usage_int(SimpleNamespace(completion_tokens=7), "completion_tokens") == 7 + assert ( + usage_int( + SimpleNamespace(model_extra={"prompt_cache_hit_tokens": 3}), + "prompt_cache_hit_tokens", + ) + == 3 + ) + assert usage_int(SimpleNamespace(prompt_tokens=None), "prompt_tokens") is None + assert usage_int({"prompt_tokens": True}, "prompt_tokens") is None + + +def test_stream_usage_rejection_matches_usage_option_400(): + error = _bad_request( + "Unrecognized request argument supplied: stream_options", + {"error": {"message": "stream_options is unsupported"}}, + ) + + assert is_stream_usage_rejection(error) + + +def test_stream_usage_rejection_does_not_match_unrelated_400(): + error = _bad_request( + "messages: invalid role", + {"error": {"message": "messages contains invalid role"}}, + ) + + assert not is_stream_usage_rejection(error) + + +@pytest.mark.asyncio +async def test_openai_chat_stream_requests_usage_and_uses_provider_prompt_tokens(): + GlobalRateLimiter.reset_instance() + try: + provider = _UsageTestProvider() + request = SimpleNamespace(model="m") + usage = SimpleNamespace(prompt_tokens=22, completion_tokens=4) + create = AsyncMock( + return_value=_stream( + [ + _chunk(content="hello"), + _chunk(finish_reason="stop"), + _chunk(usage=usage), + ] + ) + ) + + with patch.object(provider._client.chat.completions, "create", create): + events = [ + event + async for event in provider.stream_response(request, input_tokens=7) + ] + + create.assert_awaited_once() + await_args = create.await_args + assert await_args is not None + assert await_args.kwargs["stream_options"] == {"include_usage": True} + parsed = parse_sse_text("".join(events)) + start_usage = next( + event.data["message"]["usage"] + for event in parsed + if event.event == "message_start" + ) + final_usage = next( + event.data["usage"] for event in parsed if event.event == "message_delta" + ) + assert start_usage["input_tokens"] == 7 + assert final_usage == {"input_tokens": 22, "output_tokens": 4} + finally: + GlobalRateLimiter.reset_instance() + + +@pytest.mark.asyncio +async def test_openai_chat_stream_retries_without_usage_when_option_is_rejected(): + GlobalRateLimiter.reset_instance() + try: + provider = _UsageTestProvider() + body = {"model": "m", "messages": [{"role": "user", "content": "x"}]} + request_stream_usage(body) + create = AsyncMock( + side_effect=[ + _bad_request( + "stream_options is unsupported", + {"error": {"message": "stream_options is unsupported"}}, + ), + object(), + ] + ) + + with patch.object(provider._client.chat.completions, "create", create): + _stream_obj, used_body = await provider._create_stream(body) + + assert create.await_count == 2 + assert create.await_args_list[0].kwargs["stream_options"] == { + "include_usage": True + } + assert "stream_options" not in create.await_args_list[1].kwargs + assert "stream_options" not in used_body + finally: + GlobalRateLimiter.reset_instance() diff --git a/uv.lock b/uv.lock index 192d4b36..adc31cef 100644 --- a/uv.lock +++ b/uv.lock @@ -561,7 +561,7 @@ wheels = [ [[package]] name = "free-claude-code" -version = "3.4.8" +version = "3.4.9" source = { editable = "." } dependencies = [ { name = "aiohttp" },