diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index c471582dc9e..2875b30232e 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -75,24 +75,23 @@ class GigaChatModelResponseIterator: ) finish_reason = "tool_calls" - if chunk_finish_reason == "stop": - usage_data: Final = chunk.get("usage") or {} # mutable-ok: empty dict default - if usage_data and isinstance(usage_data, dict): - validated_usage: Final = {k: int(v) for k, v in usage_data.items()} - usage = convert_usage(validated_usage) - _prompt_details: dict | None = ( - usage.prompt_tokens_details.model_dump() if usage.prompt_tokens_details else None - ) # rebind-ok: conditional - _completion_details: dict | None = ( - usage.completion_tokens_details.model_dump() if usage.completion_tokens_details else None - ) # rebind-ok: conditional - usage_block = ChatCompletionUsageBlock( # pyright: ignore[reportCallIssue] # TypedDict kwarg constructor - prompt_tokens=usage.prompt_tokens, - completion_tokens=usage.completion_tokens, - total_tokens=usage.total_tokens, - prompt_tokens_details=_prompt_details, - completion_tokens_details=_completion_details, - ) + usage_data: Final = chunk.get("usage") or {} # mutable-ok: empty dict default + if usage_data and isinstance(usage_data, dict): + validated_usage: Final = {k: int(v) for k, v in usage_data.items()} + usage = convert_usage(validated_usage) + _prompt_details: dict | None = ( + usage.prompt_tokens_details.model_dump() if usage.prompt_tokens_details else None + ) # rebind-ok: conditional + _completion_details: dict | None = ( + usage.completion_tokens_details.model_dump() if usage.completion_tokens_details else None + ) # rebind-ok: conditional + usage_block = ChatCompletionUsageBlock( # pyright: ignore[reportCallIssue] # TypedDict kwarg constructor + prompt_tokens=usage.prompt_tokens, + completion_tokens=usage.completion_tokens, + total_tokens=usage.total_tokens, + prompt_tokens_details=_prompt_details, + completion_tokens_details=_completion_details, + ) return GenericStreamingChunk( text=str(text), diff --git a/litellm/llms/gigachat/utils.py b/litellm/llms/gigachat/utils.py index b1db8f685e5..2072e6003e3 100644 --- a/litellm/llms/gigachat/utils.py +++ b/litellm/llms/gigachat/utils.py @@ -9,27 +9,16 @@ GIGACHAT_BASE_URL: Final = "https://gigachat.devices.sberbank.ru/api/v1" def convert_usage(usage_data: Mapping[str, int]) -> Usage: - prompt_tokens: Final = usage_data.get("prompt_tokens", 0) - completion_tokens: Final = usage_data.get("completion_tokens", 0) precached_prompt_tokens: Final = usage_data.get("precached_prompt_tokens", 0) - total_tokens: Final = usage_data.get("total_tokens", 0) - - prompt_tokens_total: Final = prompt_tokens + precached_prompt_tokens - total_tokens_total: Final = total_tokens + precached_prompt_tokens - - prompt_tokens_details: PromptTokensDetailsWrapper | None = ( - None # rebind-ok: conditionally assigned when cached tokens exist + prompt_tokens_details: Final = ( + PromptTokensDetailsWrapper(cached_tokens=precached_prompt_tokens) if precached_prompt_tokens > 0 else None ) - if precached_prompt_tokens > 0: - prompt_tokens_details = PromptTokensDetailsWrapper( - cached_tokens=precached_prompt_tokens - ) # rebind-ok: conditionally assigned when cached tokens exist return Usage( - prompt_tokens=prompt_tokens_total, - completion_tokens=completion_tokens, + prompt_tokens=usage_data.get("prompt_tokens", 0), + completion_tokens=usage_data.get("completion_tokens", 0), prompt_tokens_details=prompt_tokens_details, - total_tokens=total_tokens_total, + total_tokens=usage_data.get("total_tokens", 0), ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 6c7d0c92e58..d3df490d43c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2957,16 +2957,21 @@ async def handle_gigachat_passthrough_router_model( request=request ) # mutable-ok: mutated in place by proxy pipeline; pyright: ignore[reportExplicitAny] # Any needed for proxy pipeline if user_api_key_dict is not None: - if data.get("metadata") is None: - data["metadata"] = {} # mutable-ok: metadata dict mutated in place - if hasattr(user_api_key_dict, "user_id") and user_api_key_dict.user_id is not None: - data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id - if hasattr(user_api_key_dict, "team_id") and user_api_key_dict.team_id is not None: - data["metadata"]["user_api_key_team_id"] = user_api_key_dict.team_id - if hasattr(user_api_key_dict, "org_id") and user_api_key_dict.org_id is not None: - data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id - if hasattr(user_api_key_dict, "agent_id") and user_api_key_dict.agent_id is not None: - data["metadata"]["agent_id"] = user_api_key_dict.agent_id + auth_metadata: Final = { + metadata_key: value + for metadata_key, value in ( + ("user_api_key_user_id", getattr(user_api_key_dict, "user_id", None)), + ("user_api_key_team_id", getattr(user_api_key_dict, "team_id", None)), + ("user_api_key_org_id", getattr(user_api_key_dict, "org_id", None)), + ("agent_id", getattr(user_api_key_dict, "agent_id", None)), + ) + if value is not None + } + existing_metadata: Final = data.get("metadata") + data["metadata"] = { + **(existing_metadata if isinstance(existing_metadata, dict) else {}), + **auth_metadata, + } verbose_proxy_logger.debug( "Gigachat router passthrough: model='%s', endpoint='%s', streaming=%s", model, endpoint, is_streaming diff --git a/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_streaming.py b/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_streaming.py new file mode 100644 index 00000000000..500aff264f2 --- /dev/null +++ b/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_streaming.py @@ -0,0 +1,85 @@ +""" +Tests for litellm.llms.gigachat.chat.streaming +""" + +from litellm.llms.gigachat.chat.streaming import GigaChatModelResponseIterator + + +def _parse(chunk: dict) -> dict: + iterator = GigaChatModelResponseIterator(streaming_response=None, sync_stream=True) + return dict(iterator.chunk_parser(chunk=chunk)) + + +class TestChunkParserUsage: + def test_usage_on_stop_chunk(self): + parsed = _parse( + { + "choices": [{"delta": {"content": ""}, "index": 0, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 25, "completion_tokens": 7, "total_tokens": 32}, + } + ) + + assert parsed["finish_reason"] == "stop" + assert parsed["usage"] is not None + assert parsed["usage"]["prompt_tokens"] == 25 + assert parsed["usage"]["completion_tokens"] == 7 + assert parsed["usage"]["total_tokens"] == 32 + + def test_usage_on_function_call_chunk(self): + """Regression: a final chunk ending in function_call still carries usage; it must not be dropped.""" + parsed = _parse( + { + "choices": [ + { + "delta": {"function_call": {"name": "get_weather", "arguments": {"city": "Moscow"}}}, + "index": 0, + "finish_reason": "function_call", + } + ], + "usage": {"prompt_tokens": 40, "completion_tokens": 12, "total_tokens": 52}, + } + ) + + assert parsed["finish_reason"] == "tool_calls" + assert parsed["tool_use"] is not None + assert parsed["usage"] is not None + assert parsed["usage"]["prompt_tokens"] == 40 + assert parsed["usage"]["completion_tokens"] == 12 + assert parsed["usage"]["total_tokens"] == 52 + + def test_usage_on_length_chunk(self): + parsed = _parse( + { + "choices": [{"delta": {"content": "truncated"}, "index": 0, "finish_reason": "length"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 128, "total_tokens": 138}, + } + ) + + assert parsed["usage"] is not None + assert parsed["usage"]["total_tokens"] == 138 + + def test_no_usage_on_interim_chunk(self): + parsed = _parse({"choices": [{"delta": {"content": "hello"}, "index": 0, "finish_reason": None}]}) + + assert parsed["text"] == "hello" + assert parsed["is_finished"] is False + assert parsed["usage"] is None + + def test_cache_hit_usage_not_inflated(self): + """precached_prompt_tokens is a subset of prompt_tokens; totals must not be inflated on cache hits.""" + parsed = _parse( + { + "choices": [{"delta": {"content": ""}, "index": 0, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": 25, + "completion_tokens": 7, + "total_tokens": 32, + "precached_prompt_tokens": 20, + }, + } + ) + + assert parsed["usage"] is not None + assert parsed["usage"]["prompt_tokens"] == 25 + assert parsed["usage"]["total_tokens"] == 32 + assert parsed["usage"]["prompt_tokens_details"]["cached_tokens"] == 20 diff --git a/tests/test_litellm/llms/gigachat/test_utils.py b/tests/test_litellm/llms/gigachat/test_utils.py index 3d45b1a7465..3d6b7d842e7 100644 --- a/tests/test_litellm/llms/gigachat/test_utils.py +++ b/tests/test_litellm/llms/gigachat/test_utils.py @@ -26,7 +26,7 @@ class TestConvertUsage: ) def test_usage_with_precached_prompt_tokens(self): - """Test convert_usage adds precached_prompt_tokens to prompt_tokens and total_tokens.""" + """precached_prompt_tokens is a subset of prompt_tokens (OpenAI cached_tokens semantics), never additive.""" result = convert_usage( { "prompt_tokens": 10, @@ -37,9 +37,9 @@ class TestConvertUsage: ) assert result == Usage( - prompt_tokens=13, + prompt_tokens=10, completion_tokens=5, - total_tokens=18, + total_tokens=15, prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=3), ) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 303c71a1630..1d4b0264879 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -2123,6 +2123,77 @@ class TestGigachatProxyRoute: mock_llm_router.allm_passthrough_route.assert_awaited_once() assert isinstance(result, Response) + @pytest.mark.asyncio + async def test_gigachat_router_handler_keeps_cached_body_and_payload_metadata_pristine(self): + """Regression: auth-metadata injection must not leak into the cached parsed body or the upstream payload.""" + from litellm.proxy.common_utils.http_parsing_utils import get_request_body + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + handle_gigachat_passthrough_router_model, + ) + + body = json.dumps( + { + "model": "gigachat-router", + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"client_tag": "user-supplied"}, + } + ).encode() + scope = { + "type": "http", + "method": "POST", + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + "path": "/gigachat/chat/completions", + } + + async def receive(): + return {"type": "http.request", "body": body, "more_body": False} + + request = Request(scope, receive) + request_body = await get_request_body(request) + + captured: dict = {} + + class _CapturingProcessor: + def __init__(self, data: dict): + captured["data"] = data + + async def base_passthrough_process_llm_request(self, **kwargs): + return Response(content=b"{}", status_code=200) + + with patch( # test-quality-ok: patching litellm internal for unit test isolation + "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing", + _CapturingProcessor, + ): + await handle_gigachat_passthrough_router_model( + model="gigachat-router", + endpoint="/chat/completions", + request=request, + request_body=request_body, + fastapi_response=Response(), + llm_router=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-1", team_id="team-1"), + proxy_logging_obj=MagicMock(), + general_settings={}, + proxy_config=MagicMock(), + select_data_generator=MagicMock(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + version=None, + ) + + data = captured["data"] + assert data["json"] is request_body + assert request_body["metadata"] == {"client_tag": "user-supplied"} + assert data["metadata"]["client_tag"] == "user-supplied" + assert data["metadata"]["user_api_key_user_id"] == "user-1" + assert data["metadata"]["user_api_key_team_id"] == "team-1" + cached_reread = await get_request_body(request) + assert cached_reread["metadata"] == {"client_tag": "user-supplied"} + @pytest.mark.asyncio @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body",