scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 83

fix(gigachat): correct cached-token accounting, stream usage on all finish reasons, stop mutating cached request body

Browse source 
precached_prompt_tokens is a subset of prompt_tokens (OpenAI cached_tokens
semantics), so map it to prompt_tokens_details.cached_tokens instead of
adding it on top of prompt/total. Emit stream usage from any final chunk
carrying it rather than only finish_reason stop, which dropped tokens for
function_call and length streams. Merge auth metadata into a new dict in
the gigachat router handler instead of mutating the shared parsed-body
cache in place.
This commit is contained in:
mateo-berri 2026-08-31 14:28:30 -07:00
parent 15aa51a88a
commit 98b1e2e7b4
6 changed files with 196 additions and 47 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
litellm/llms/gigachat/chat/streaming.py
View file

@ -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),

21
litellm/llms/gigachat/utils.py
View file

@ -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),
)

25
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -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

85
tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_streaming.py Normal file
View file

@ -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

6
tests/test_litellm/llms/gigachat/test_utils.py
View file

@ -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),
)

71
tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
View file

@ -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",