diff --git a/providers/transports/openai_chat/output_cap.py b/providers/transports/openai_chat/output_cap.py new file mode 100644 index 00000000..c5a13ef9 --- /dev/null +++ b/providers/transports/openai_chat/output_cap.py @@ -0,0 +1,83 @@ +"""Recover from upstream ``max_(completion_)tokens`` too-large 400 rejections. + +Some OpenAI-compatible providers (Groq, NVIDIA NIM, ...) cap the per-request +output token count below what Claude Code asks for and reject the whole request +with an HTTP 400 that names the allowed maximum, e.g.:: + + max_completion_tokens must be less than or equal to 40960, ... + +This module parses that maximum and clamps the request body so the transport can +retry once and succeed. The transport also remembers the learned cap per model +so later requests clamp proactively instead of paying the 400 every time. +""" + +import json +import re +from typing import Any + +import openai + +# Body keys that carry the output-token budget across OpenAI-compatible policies. +_OUTPUT_TOKEN_FIELDS = ("max_completion_tokens", "max_tokens") + +# Only treat a 400 as an output-cap rejection when it names one of these fields. +_OUTPUT_TOKEN_KEYWORDS = ("max_completion_tokens", "max_tokens") + +# Comparator phrases that precede the allowed maximum in provider error text. +_CAP_PATTERNS: tuple[re.Pattern[str], ...] = ( + re.compile(r"less than or equal to\s+(\d+)"), + re.compile(r"smaller than or equal to\s+(\d+)"), + re.compile(r"<=\s*(\d+)"), + re.compile(r"at most\s+(\d+)"), + re.compile(r"must not exceed\s+(\d+)"), + re.compile(r"maximum(?:\s+value)?(?:\s+for\s+\S+)?\s+is\s+(\d+)"), + re.compile(r"maximum(?:\s+allowed)?(?:\s+value)?\s+of\s+(\d+)"), +) + + +def _is_bad_request(error: Exception) -> bool: + return isinstance(error, openai.BadRequestError) or ( + getattr(error, "status_code", None) == 400 + ) + + +def _error_text(error: Exception) -> str: + text = str(error) + body = getattr(error, "body", None) + if body is not None: + text = f"{text} {json.dumps(body, default=str)}" + return text.lower() + + +def parse_output_token_cap(error: Exception) -> int | None: + """Return the allowed output-token maximum named in a 400 rejection, if any.""" + if not _is_bad_request(error): + return None + + text = _error_text(error) + if not any(keyword in text for keyword in _OUTPUT_TOKEN_KEYWORDS): + return None + + for pattern in _CAP_PATTERNS: + match = pattern.search(text) + if match: + cap = int(match.group(1)) + if cap > 0: + return cap + return None + + +def clamp_output_tokens(body: dict[str, Any], cap: int) -> dict[str, Any] | None: + """Return a shallow clone with output-token fields clamped to ``cap``. + + Returns ``None`` when nothing needs clamping (no output field exceeds the + cap), so callers can avoid a pointless identical retry. + """ + clamped: dict[str, Any] | None = None + for field in _OUTPUT_TOKEN_FIELDS: + value = body.get(field) + if isinstance(value, int) and not isinstance(value, bool) and value > cap: + if clamped is None: + clamped = dict(body) + clamped[field] = cap + return clamped diff --git a/providers/transports/openai_chat/transport.py b/providers/transports/openai_chat/transport.py index f40665a1..e3726c17 100644 --- a/providers/transports/openai_chat/transport.py +++ b/providers/transports/openai_chat/transport.py @@ -5,6 +5,7 @@ from collections.abc import AsyncIterator, Iterator, Mapping from typing import Any import httpx +from loguru import logger from openai import AsyncOpenAI from core.anthropic.streaming import AnthropicStreamLedger @@ -17,6 +18,7 @@ from providers.error_mapping import ( from providers.model_listing import extract_openai_model_ids from providers.rate_limit import GlobalRateLimiter +from .output_cap import clamp_output_tokens, parse_output_token_cap from .stream import OpenAIChatStreamAdapter @@ -36,6 +38,9 @@ class OpenAIChatTransport(BaseProvider): self._provider_name = provider_name self._api_key = api_key self._base_url = base_url.rstrip("/") + # Learned per-model output-token caps from upstream 400 rejections, so + # later requests clamp proactively instead of paying the 400 each time. + self._model_output_caps: dict[str, int] = {} self._global_rate_limiter = GlobalRateLimiter.get_scoped_instance( provider_name.lower(), rate_limit=config.rate_limit, @@ -113,6 +118,7 @@ class OpenAIChatTransport(BaseProvider): async def _create_stream(self, body: dict) -> tuple[Any, dict]: """Create a streaming chat completion, optionally retrying once.""" + body = self._apply_learned_output_cap(body) try: create_body = self._prepare_create_body(body) stream = await self._global_rate_limiter.execute_with_retry( @@ -120,7 +126,9 @@ class OpenAIChatTransport(BaseProvider): ) return stream, body except Exception as error: - retry_body = self._get_retry_request_body(error, body) + 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 @@ -130,6 +138,37 @@ class OpenAIChatTransport(BaseProvider): ) return stream, retry_body + def _apply_learned_output_cap(self, body: dict) -> dict: + """Clamp output tokens to a previously learned cap for this model.""" + model = body.get("model") + if not isinstance(model, str): + return body + cap = self._model_output_caps.get(model) + if cap is None: + return body + clamped = clamp_output_tokens(body, cap) + return clamped if clamped is not None else body + + def _retry_body_for_output_cap(self, error: Exception, body: dict) -> dict | None: + """Learn an upstream output-token cap from a 400 and clamp for one retry.""" + cap = parse_output_token_cap(error) + if cap is None: + return None + model = body.get("model") + if isinstance(model, str): + previous = self._model_output_caps.get(model) + cap = cap if previous is None else min(previous, cap) + self._model_output_caps[model] = cap + clamped = clamp_output_tokens(body, cap) + if clamped is None: + return None + logger.warning( + "{}_STREAM: clamping output tokens to {} after upstream cap rejection", + self._provider_name, + cap, + ) + return clamped + def _openai_error_message(self, error: Exception, request_id: str | None) -> str: mapped_error = map_error(error, rate_limiter=self._global_rate_limiter) return user_visible_message_for_mapped_provider_error( diff --git a/pyproject.toml b/pyproject.toml index 2de42e1b..999ff2fa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "free-claude-code" -version = "3.2.0" +version = "3.2.1" 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_openai_chat_output_cap.py b/tests/providers/test_openai_chat_output_cap.py new file mode 100644 index 00000000..fdab5e54 --- /dev/null +++ b/tests/providers/test_openai_chat_output_cap.py @@ -0,0 +1,207 @@ +"""Tests for OpenAI-compatible output-token cap recovery (issue #955). + +Covers the pure parse/clamp helpers and the transport behavior that clamps +``max_completion_tokens``/``max_tokens`` to the upstream maximum, retries once, +and learns the cap so later requests clamp proactively. +""" + +from contextlib import asynccontextmanager +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from providers.base import ProviderConfig +from providers.groq import GROQ_DEFAULT_BASE, GroqProvider +from providers.transports.openai_chat.output_cap import ( + clamp_output_tokens, + parse_output_token_cap, +) + + +class _BadRequest(Exception): + """Stand-in for openai.BadRequestError (status_code + optional JSON body).""" + + def __init__(self, message: str, body: object | None = None): + super().__init__(message) + self.status_code = 400 + self.body = body + + +# --------------------------------------------------------------------------- # +# Pure helpers +# --------------------------------------------------------------------------- # + + +def test_parse_cap_from_groq_message(): + error = _BadRequest( + "max_completion_tokens must be less than or equal to 40960, the maximum " + "value for max_completion_tokens is less than the context_window for this model" + ) + assert parse_output_token_cap(error) == 40960 + + +@pytest.mark.parametrize( + "message,expected", + [ + ("max_tokens: maximum value is 8192", 8192), + ("max_tokens must not exceed 16000", 16000), + ("`max_completion_tokens` <= 4096 required", 4096), + ("max_tokens at most 2048 allowed", 2048), + ("maximum allowed value of 32768 for max_tokens", 32768), + ], +) +def test_parse_cap_various_phrasings(message, expected): + assert parse_output_token_cap(_BadRequest(message)) == expected + + +def test_parse_cap_reads_json_body(): + error = _BadRequest( + "invalid request", + body={"error": {"param": "max_completion_tokens", "message": "<= 12000"}}, + ) + assert parse_output_token_cap(error) == 12000 + + +def test_parse_cap_ignores_non_400(): + error = _BadRequest("max_tokens must be less than or equal to 40960") + error.status_code = 500 + assert parse_output_token_cap(error) is None + + +def test_parse_cap_ignores_unrelated_400(): + assert parse_output_token_cap(_BadRequest("temperature must be <= 2")) is None + + +def test_parse_cap_returns_none_without_number(): + assert ( + parse_output_token_cap(_BadRequest("max_tokens is larger than allowed")) is None + ) + + +def test_clamp_reduces_max_completion_tokens(): + assert clamp_output_tokens({"max_completion_tokens": 64000}, 40960) == { + "max_completion_tokens": 40960 + } + + +def test_clamp_reduces_max_tokens(): + assert clamp_output_tokens({"max_tokens": 100000}, 8192) == {"max_tokens": 8192} + + +def test_clamp_noop_when_within_cap_returns_none(): + assert clamp_output_tokens({"max_completion_tokens": 1000}, 40960) is None + + +def test_clamp_does_not_mutate_input(): + body = {"max_tokens": 99999, "model": "m"} + clamped = clamp_output_tokens(body, 8192) + assert body["max_tokens"] == 99999 + assert clamped is not None + assert clamped["max_tokens"] == 8192 + + +def test_clamp_ignores_bool_values(): + assert clamp_output_tokens({"max_tokens": True}, 8192) is None + + +# --------------------------------------------------------------------------- # +# Transport integration (via GroqProvider, which uses max_completion_tokens) +# --------------------------------------------------------------------------- # + + +class MockMessage: + def __init__(self, role, content): + self.role = role + self.content = content + + +class MockRequest: + def __init__(self, max_tokens=64000): + self.model = "llama-3.3-70b-versatile" + self.messages = [MockMessage("user", "Hello")] + self.max_tokens = max_tokens + self.temperature = 0.5 + self.top_p = 0.9 + self.system = "System prompt" + self.stop_sequences = None + self.tools = [] + self.thinking = MagicMock() + self.thinking.enabled = False + + +@pytest.fixture(autouse=True) +def mock_rate_limiter(): + @asynccontextmanager + async def _slot(): + yield + + with patch("providers.transports.openai_chat.transport.GlobalRateLimiter") as mock: + instance = mock.get_scoped_instance.return_value + + async def _passthrough(fn, *args, **kwargs): + return await fn(*args, **kwargs) + + instance.execute_with_retry = AsyncMock(side_effect=_passthrough) + instance.concurrency_slot.side_effect = _slot + yield instance + + +@pytest.fixture +def groq_provider(): + return GroqProvider( + ProviderConfig( + api_key="test_groq_key", + base_url=GROQ_DEFAULT_BASE, + rate_limit=10, + rate_window=60, + enable_thinking=False, + ) + ) + + +@pytest.mark.asyncio +async def test_create_stream_clamps_and_learns_on_cap_rejection(groq_provider): + body = groq_provider._build_request_body(MockRequest(max_tokens=64000)) + assert body["max_completion_tokens"] == 64000 + model = body["model"] + + error = _BadRequest("max_completion_tokens must be less than or equal to 40960") + create = AsyncMock(side_effect=[error, object()]) + + with patch.object(groq_provider._client.chat.completions, "create", create): + _stream, used_body = await groq_provider._create_stream(body) + + assert create.call_count == 2 + assert create.call_args_list[1].kwargs["max_completion_tokens"] == 40960 + assert used_body["max_completion_tokens"] == 40960 + assert groq_provider._model_output_caps[model] == 40960 + + +@pytest.mark.asyncio +async def test_learned_cap_clamps_next_request_without_a_400(groq_provider): + body = groq_provider._build_request_body(MockRequest(max_tokens=64000)) + model = body["model"] + groq_provider._model_output_caps[model] = 40960 + + create = AsyncMock(return_value=object()) + with patch.object(groq_provider._client.chat.completions, "create", create): + _stream, used_body = await groq_provider._create_stream(body) + + assert create.call_count == 1 + assert create.call_args.kwargs["max_completion_tokens"] == 40960 + assert used_body["max_completion_tokens"] == 40960 + + +@pytest.mark.asyncio +async def test_unrelated_400_is_not_clamped_and_propagates(groq_provider): + body = groq_provider._build_request_body(MockRequest(max_tokens=100)) + create = AsyncMock(side_effect=_BadRequest("messages: invalid role 'wizard'")) + + with ( + patch.object(groq_provider._client.chat.completions, "create", create), + pytest.raises(Exception, match="wizard"), + ): + await groq_provider._create_stream(body) + + assert create.call_count == 1 + assert groq_provider._model_output_caps == {} diff --git a/uv.lock b/uv.lock index fa006ca7..1358a9a9 100644 --- a/uv.lock +++ b/uv.lock @@ -561,7 +561,7 @@ wheels = [ [[package]] name = "free-claude-code" -version = "3.2.0" +version = "3.2.1" source = { editable = "." } dependencies = [ { name = "aiohttp" },