From 4106999e55243f16f4d61f853d5bdd4d055a45bc Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Thu, 16 Apr 2026 12:57:07 +0000 Subject: [PATCH 001/120] fix(gigachat): improve usage reporting and config handling --- litellm/llms/gigachat/authenticator.py | 36 ++++++++++++++++--- litellm/llms/gigachat/chat/streaming.py | 12 +++++-- litellm/llms/gigachat/chat/transformation.py | 18 ++++------ .../llms/gigachat/embedding/transformation.py | 12 +++---- litellm/llms/gigachat/file_handler.py | 18 ++++++---- litellm/llms/gigachat/utils.py | 35 ++++++++++++++++++ 6 files changed, 101 insertions(+), 30 deletions(-) create mode 100644 litellm/llms/gigachat/utils.py diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index 59942a9c038..83824c72fa7 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -65,6 +65,7 @@ def get_access_token( credentials: Optional[str] = None, scope: Optional[str] = None, auth_url: Optional[str] = None, + litellm_params: Optional[dict] = None, ) -> str: """ Get valid access token, using cache if available. @@ -80,6 +81,15 @@ def get_access_token( Raises: GigaChatAuthError: If authentication fails """ + if not litellm_params: + litellm_params = {} + + access_token = litellm_params.get("gigachat_access_token") or get_secret_str( + "GIGACHAT_ACCESS_TOKEN" + ) + if access_token: + return access_token + credentials = credentials or _get_credentials() if not credentials: raise GigaChatAuthError( @@ -87,8 +97,8 @@ def get_access_token( message="GigaChat credentials not provided. Set GIGACHAT_CREDENTIALS or GIGACHAT_API_KEY environment variable.", ) - scope = scope or _get_scope() - auth_url = auth_url or _get_auth_url() + scope = scope or litellm_params.get("gigachat_scope") or _get_scope() + auth_url = auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url() # Check cache cache_key = f"gigachat_token:{credentials[:16]}" @@ -117,6 +127,7 @@ async def get_access_token_async( credentials: Optional[str] = None, scope: Optional[str] = None, auth_url: Optional[str] = None, + litellm_params: Optional[dict] = None, ) -> str: """Async version of get_access_token.""" credentials = credentials or _get_credentials() @@ -125,9 +136,26 @@ async def get_access_token_async( status_code=401, message="GigaChat credentials not provided. Set GIGACHAT_CREDENTIALS or GIGACHAT_API_KEY environment variable.", ) + if not litellm_params: + litellm_params = {} - scope = scope or _get_scope() - auth_url = auth_url or _get_auth_url() + access_token = litellm_params.get("gigachat_access_token") or get_secret_str( + "GIGACHAT_ACCESS_TOKEN" + ) + if access_token: + return access_token + + credentials = credentials or _get_credentials() + if not credentials: + raise GigaChatAuthError( + status_code=401, + message="GigaChat credentials not provided. Set GIGACHAT_CREDENTIALS or GIGACHAT_API_KEY environment variable.", + ) + + scope = scope or litellm_params.get("gigachat_scope") or _get_scope() + auth_url = ( + auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url() + ) # Check cache cache_key = f"gigachat_token:{credentials[:16]}" diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index 4f10f8bb658..2fd5fc22b01 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -6,11 +6,12 @@ import json import uuid from typing import Any, Optional +from litellm.llms.gigachat.utils import convert_usage from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, ) -from litellm.types.utils import GenericStreamingChunk +from litellm.types.utils import ChatCompletionUsageBlock, GenericStreamingChunk class GigaChatModelResponseIterator: @@ -70,6 +71,13 @@ class GigaChatModelResponseIterator: ) finish_reason = "tool_calls" + usage_block = None + if finish_reason == "stop": + usage_data = chunk.get("usage", {}) + if usage_data: + usage = convert_usage(usage_data) + usage_block = ChatCompletionUsageBlock(**usage.dict()) + if finish_reason is not None: is_finished = True @@ -78,7 +86,7 @@ class GigaChatModelResponseIterator: tool_use=tool_use, is_finished=is_finished, finish_reason=finish_reason or "", - usage=None, + usage=usage_block, index=choice.get("index", 0), ) diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index cef80768762..b3bb7bc5770 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -13,9 +13,10 @@ import httpx from litellm._logging import verbose_logger from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.gigachat.utils import convert_usage, get_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import Choices, Message, ModelResponse, Usage +from litellm.types.utils import Choices, Message, ModelResponse from ..authenticator import get_access_token from ..file_handler import upload_file_sync @@ -27,9 +28,6 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any -# GigaChat API endpoint -GIGACHAT_BASE_URL = "https://gigachat.devices.sberbank.ru/api/v1" - def is_valid_json(value: str) -> bool: """Checks whether the value passed is a valid serialized JSON string""" @@ -94,7 +92,7 @@ class GigaChatConfig(BaseConfig): stream: Optional[bool] = None, ) -> str: """Get complete API URL for chat completions.""" - base = api_base or get_secret_str("GIGACHAT_API_BASE") or GIGACHAT_BASE_URL + base = get_api_base(api_base) return f"{base}/chat/completions" def validate_environment( @@ -116,7 +114,9 @@ class GigaChatConfig(BaseConfig): or get_secret_str("GIGACHAT_CREDENTIALS") or get_secret_str("GIGACHAT_API_KEY") ) - access_token = get_access_token(credentials=credentials) + access_token = get_access_token( + credentials=credentials, litellm_params=litellm_params + ) # Store credentials for image uploads self._current_credentials = credentials @@ -467,11 +467,7 @@ class GigaChatConfig(BaseConfig): # Build usage usage_data = response_json.get("usage", {}) - usage = Usage( - prompt_tokens=usage_data.get("prompt_tokens", 0), - completion_tokens=usage_data.get("completion_tokens", 0), - total_tokens=usage_data.get("total_tokens", 0), - ) + usage = convert_usage(usage_data) model_response.id = response_json.get("id", f"chatcmpl-{uuid.uuid4().hex[:12]}") model_response.created = response_json.get("created", int(time.time())) diff --git a/litellm/llms/gigachat/embedding/transformation.py b/litellm/llms/gigachat/embedding/transformation.py index 0da6565050e..8dcccd49fcb 100644 --- a/litellm/llms/gigachat/embedding/transformation.py +++ b/litellm/llms/gigachat/embedding/transformation.py @@ -14,14 +14,12 @@ from litellm import LlmProviders from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.gigachat.utils import get_api_base from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues from litellm.types.utils import EmbeddingResponse from ..authenticator import get_access_token -# GigaChat API endpoint -GIGACHAT_BASE_URL = "https://gigachat.devices.sberbank.ru/api/v1" - class GigaChatEmbeddingError(BaseLLMException): """GigaChat Embedding API error.""" @@ -82,7 +80,7 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): Returns: Tuple of (custom_llm_provider, api_base, dynamic_api_key) """ - api_base = api_base or GIGACHAT_BASE_URL + api_base = get_api_base(api_base) return LlmProviders.GIGACHAT.value, api_base, api_key def get_complete_url( @@ -95,7 +93,7 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): stream: Optional[bool] = None, ) -> str: """Get the complete URL for embeddings endpoint.""" - base = api_base or GIGACHAT_BASE_URL + base = get_api_base(api_base) return f"{base}/embeddings" def transform_embedding_request( @@ -194,7 +192,9 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): Set up headers with OAuth token for GigaChat. """ # Get access token via OAuth - access_token = get_access_token(api_key) + access_token = get_access_token( + credentials=api_key, litellm_params=litellm_params + ) default_headers = { "Content-Type": "application/json", diff --git a/litellm/llms/gigachat/file_handler.py b/litellm/llms/gigachat/file_handler.py index 200428a747a..2054fe1f566 100644 --- a/litellm/llms/gigachat/file_handler.py +++ b/litellm/llms/gigachat/file_handler.py @@ -16,13 +16,11 @@ from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, ) +from litellm.llms.gigachat.utils import get_api_base from litellm.types.utils import LlmProviders from .authenticator import get_access_token, get_access_token_async -# GigaChat API endpoint -GIGACHAT_BASE_URL = "https://gigachat.devices.sberbank.ru/api/v1" - # Simple in-memory cache for file IDs _file_cache: Dict[str, str] = {} @@ -82,6 +80,7 @@ def upload_file_sync( image_url: str, credentials: Optional[str] = None, api_base: Optional[str] = None, + litellm_params: Optional[dict] = None, ) -> Optional[str]: """ Upload file to GigaChat and return file_id (sync). @@ -114,10 +113,12 @@ def upload_file_sync( filename = f"{uuid.uuid4()}.{ext}" # Get access token - access_token = get_access_token(credentials) + access_token = get_access_token( + credentials=credentials, litellm_params=litellm_params + ) # Upload to GigaChat - base_url = api_base or GIGACHAT_BASE_URL + base_url = get_api_base(api_base) upload_url = f"{base_url}/files" client = _get_httpx_client(params={"ssl_verify": False}) @@ -147,6 +148,7 @@ async def upload_file_async( image_url: str, credentials: Optional[str] = None, api_base: Optional[str] = None, + litellm_params: Optional[dict] = None, ) -> Optional[str]: """ Upload file to GigaChat and return file_id (async). @@ -179,10 +181,12 @@ async def upload_file_async( filename = f"{uuid.uuid4()}.{ext}" # Get access token - access_token = await get_access_token_async(credentials) + access_token = await get_access_token_async( + credentials=credentials, litellm_params=litellm_params + ) # Upload to GigaChat - base_url = api_base or GIGACHAT_BASE_URL + base_url = get_api_base(api_base) upload_url = f"{base_url}/files" client = get_async_httpx_client( diff --git a/litellm/llms/gigachat/utils.py b/litellm/llms/gigachat/utils.py new file mode 100644 index 00000000000..e79083f16fe --- /dev/null +++ b/litellm/llms/gigachat/utils.py @@ -0,0 +1,35 @@ +from typing import Optional + +from litellm.secret_managers.main import get_secret_str +from litellm.types.utils import PromptTokensDetailsWrapper, Usage + +# GigaChat API endpoint +GIGACHAT_BASE_URL = "https://gigachat.devices.sberbank.ru/api/v1" + + +def convert_usage(usage_data: dict[str, int]) -> Usage: + prompt_tokens = usage_data.get("prompt_tokens", 0) + completion_tokens = usage_data.get("completion_tokens", 0) + precached_prompt_tokens = usage_data.get("precached_prompt_tokens", 0) + total_tokens = usage_data.get("total_tokens", 0) + + prompt_tokens += precached_prompt_tokens + total_tokens += precached_prompt_tokens + + prompt_tokens_details = None + if precached_prompt_tokens > 0: + prompt_tokens_details = PromptTokensDetailsWrapper( + cached_tokens=precached_prompt_tokens + ) + + return Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + prompt_tokens_details=prompt_tokens_details, + total_tokens=total_tokens, + ) + + +@staticmethod +def get_api_base(api_base: Optional[str] = None) -> Optional[str]: + return api_base or get_secret_str("GIGACHAT_API_BASE") or GIGACHAT_BASE_URL From 718985144ad2828db642b2438160b48f54fe57a7 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Thu, 16 Apr 2026 15:22:19 +0000 Subject: [PATCH 002/120] fix(gigachat): fix GigaChat-2 model name --- docs/my-website/docs/providers/gigachat.md | 4 ++-- model_prices_and_context_window.json | 11 ++++++++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/docs/my-website/docs/providers/gigachat.md b/docs/my-website/docs/providers/gigachat.md index 13eec298c25..9b4bee80125 100644 --- a/docs/my-website/docs/providers/gigachat.md +++ b/docs/my-website/docs/providers/gigachat.md @@ -201,7 +201,7 @@ model_list: ssl_verify: false - model_name: gigachat-lite litellm_params: - model: gigachat/GigaChat-2-Lite + model: gigachat/GigaChat-2 api_key: "os.environ/GIGACHAT_CREDENTIALS" ssl_verify: false - model_name: gigachat-embeddings @@ -260,7 +260,7 @@ print(response) | Model Name | Context Window | Vision | Description | |------------|----------------|--------|-------------| -| gigachat/GigaChat-2-Lite | 128K | No | Fast, lightweight model | +| gigachat/GigaChat-2 | 128K | No | Fast, lightweight model | | gigachat/GigaChat-2-Pro | 128K | Yes | Professional model with vision | | gigachat/GigaChat-2-Max | 128K | Yes | Maximum capability model | diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c624736d6bf..47e2a5eb0ab 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -16731,7 +16731,7 @@ "supports_response_schema": true, "supports_vision": true }, - "gigachat/GigaChat-2-Lite": { + "gigachat/GigaChat-2": { "input_cost_per_token": 0.0, "litellm_provider": "gigachat", "max_input_tokens": 128000, @@ -16793,6 +16793,15 @@ "output_cost_per_token": 0.0, "output_vector_size": 2560 }, + "gigachat/GigaEmbeddings-3B-2025-09": { + "input_cost_per_token": 0.0, + "litellm_provider": "gigachat", + "max_input_tokens": 4096, + "max_tokens": 4096, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 2048 + }, "gmi/anthropic/claude-opus-4.5": { "input_cost_per_token": 5e-06, "litellm_provider": "gmi", From bf36953dcbeb8fbad0f632dccb71b216219e41e0 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Thu, 16 Apr 2026 15:28:27 +0000 Subject: [PATCH 003/120] feature(gigachat): add gigachat passthrough endpoint --- docs/my-website/docs/pass_through/gigachat.md | 122 +++++++++ .../litellm_core_utils/get_litellm_params.py | 3 + .../get_llm_provider_logic.py | 12 + litellm/llms/gigachat/__init__.py | 2 + litellm/llms/gigachat/passthrough/__init__.py | 7 + .../gigachat/passthrough/transformation.py | 196 ++++++++++++++ litellm/main.py | 3 + litellm/proxy/_types.py | 1 + .../llm_passthrough_endpoints.py | 242 ++++++++++++++++++ .../provider_create_fields.json | 62 +++++ litellm/utils.py | 6 + .../test_llm_pass_through_endpoints.py | 173 +++++++++++++ .../public/assets/logos/gigachat.svg | 27 ++ .../src/components/provider_info_helpers.tsx | 3 + 14 files changed, 859 insertions(+) create mode 100644 docs/my-website/docs/pass_through/gigachat.md create mode 100644 litellm/llms/gigachat/passthrough/__init__.py create mode 100644 litellm/llms/gigachat/passthrough/transformation.py create mode 100644 ui/litellm-dashboard/public/assets/logos/gigachat.svg diff --git a/docs/my-website/docs/pass_through/gigachat.md b/docs/my-website/docs/pass_through/gigachat.md new file mode 100644 index 00000000000..aa302fd104d --- /dev/null +++ b/docs/my-website/docs/pass_through/gigachat.md @@ -0,0 +1,122 @@ +# GigaChat Passthrough + +Pass-through endpoints for direct GigaChat API access via LiteLLM Proxy. + +## Overview + +| Feature | Supported | Notes | +|-------|-------|-------| +| Cost Tracking | ✅ | Works with proxy cost metadata and router models | +| Logging | ✅ | Logs requests and responses across LiteLLM integrations | +| Streaming | ✅ | Supported for streaming GigaChat chat completions | + +## When to use this + +- Use the native LiteLLM GigaChat provider for standard chat and embedding calls when possible. +- Use `/gigachat` passthrough when you need provider-specific GigaChat endpoints or raw GigaChat request shapes. +- This is useful for newer or less common GigaChat API endpoints that LiteLLM does not yet expose natively. + +## How it works + +Any path under `/gigachat` is treated as a provider-specific route and routed through LiteLLM's GigaChat passthrough path. +The proxy accepts the same request body shape as GigaChat and forwards it to the GigaChat backend. + +### Proxy base URL mapping + +| Original GigaChat URL | Proxy URL | +|-----------------------|-----------| +| `https://gigachat.devices.sberbank.ru/api/v1` | `http://0.0.0.0:4000/gigachat/api/v1` | + +## Request format + +The proxy requires a `model` field in the request body. For GigaChat passthrough, use the LightLLM model prefix format such as `gigachat/GigaChat-2-Max`. + +### Example: Chat completion + +```bash +curl --request POST \ + --url http://0.0.0.0:4000/gigachat/api/v1/chat/completions \ + --header 'accept: application/json' \ + --header 'content-type: application/json' \ + --header 'x-api-key: $LITELLM_API_KEY' \ + --data '{ + "model": "gigachat/GigaChat-2-Max", + "messages": [ + {"role": "user", "content": "Hello, world"} + ] + }' +``` + +### Python example + +```python +import requests +import os + +response = requests.post( + "http://0.0.0.0:4000/gigachat/api/v1/chat/completions", + headers={ + "Content-Type": "application/json", + "x-api-key": os.environ["LITELLM_API_KEY"], + }, + json={ + "model": "gigachat/GigaChat-2-Max", + "messages": [ + {"role": "user", "content": "Hello, world"} + ], + }, +) +print(response.json()) +``` + +## Authentication + +- Authenticate to the proxy with `x-api-key: $LITELLM_API_KEY` or `Authorization: Bearer $LITELLM_API_KEY`. +- The proxy then uses the configured GigaChat credentials to authenticate with the upstream GigaChat API. + +## Notes + +- GigaChat uses OAuth-style credentials. Configure your GigaChat credentials in LiteLLM using `GIGACHAT_CREDENTIALS` or `GIGACHAT_API_KEY` as described in the main GigaChat provider docs. +- The proxy automatically handles GigaChat's self-signed SSL setup when forwarding requests, so you do not need to disable SSL verification from the client side. +- The `model` field is required for passthrough requests. + +## Advanced + +### Use with router-backed GigaChat models + +If you define router models in `config.yaml`, you can use the passthrough endpoint with a router-backed GigaChat model: + +```bash +curl --request POST \ + --url http://0.0.0.0:4000/gigachat/api/v1/chat/completions \ + --header 'Content-Type: application/json' \ + --header 'x-api-key: $LITELLM_API_KEY' \ + --data '{ + "model": "gigachat/GigaChat-2-Max", + "messages": [ + {"role": "user", "content": "Hello, world"} + ] + }' +``` + +### Sending metadata + +You can attach LiteLLM metadata for cost tracking and tags using `litellm_metadata` in the request body: + +```bash +curl --request POST \ + --url http://0.0.0.0:4000/gigachat/api/v1/chat/completions \ + --header 'accept: application/json' \ + --header 'content-type: application/json' \ + --header 'x-api-key: $LITELLM_API_KEY' \ + --data '{ + "model": "gigachat/GigaChat-2-Max", + "messages": [ + {"role": "user", "content": "Hello, world"} + ], + "litellm_metadata": { + "tags": ["test-tag"], + "user": "test-user" + } + }' +``` diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index ad9538ac171..cf1d2b4d2f3 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -30,6 +30,9 @@ _OPTIONAL_KWARGS_KEYS = frozenset( "aws_sts_endpoint", "aws_external_id", "aws_bedrock_runtime_endpoint", + "gigachat_scope", + "gigachat_auth_url", + "gigachat_access_token", "tpm", "rpm", } diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 95bcd4d7186..962e2054205 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -324,6 +324,9 @@ def get_llm_provider( # noqa: PLR0915 elif endpoint == "https://api.inference.wandb.ai/v1": custom_llm_provider = "wandb" dynamic_api_key = get_secret_str("WANDB_API_KEY") + elif endpoint == "https://gigachat.devices.sberbank.ru/api/v1": + custom_llm_provider = "gigachat" + dynamic_api_key = get_secret_str("GIGACHAT_API_KEY") if api_base is not None and not isinstance(api_base, str): raise Exception( @@ -459,6 +462,8 @@ def get_llm_provider( # noqa: PLR0915 custom_llm_provider = "amazon_nova" elif model.startswith("sap/"): custom_llm_provider = "sap" + elif model in litellm.gigachat_models or model.startswith("gigachat/"): + custom_llm_provider = "gigachat" if not custom_llm_provider: if litellm.suppress_debug_info is False: print() # noqa @@ -944,6 +949,13 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 api_base or get_secret_str("MANUS_API_BASE") or "https://api.manus.im" ) dynamic_api_key = api_key or get_secret_str("MANUS_API_KEY") + elif custom_llm_provider == "gigachat": + api_base = ( + api_base + or get_secret_str("GIGACHAT_API_BASE") + or "https://gigachat.devices.sberbank.ru/api/v1" + ) + dynamic_api_key = api_key or get_secret_str("GIGACHAT_API_KEY") if api_base is not None and not isinstance(api_base, str): raise Exception("api base needs to be a string. api_base={}".format(api_base)) diff --git a/litellm/llms/gigachat/__init__.py b/litellm/llms/gigachat/__init__.py index 3ddbd7864d9..af5d2717643 100644 --- a/litellm/llms/gigachat/__init__.py +++ b/litellm/llms/gigachat/__init__.py @@ -15,9 +15,11 @@ API Documentation: https://developers.sber.ru/docs/ru/gigachat/api/overview from .chat.transformation import GigaChatConfig, GigaChatError from .embedding.transformation import GigaChatEmbeddingConfig +from .passthrough.transformation import GigaChatPassthroughConfig __all__ = [ "GigaChatConfig", "GigaChatEmbeddingConfig", "GigaChatError", + "GigaChatPassthroughConfig", ] diff --git a/litellm/llms/gigachat/passthrough/__init__.py b/litellm/llms/gigachat/passthrough/__init__.py new file mode 100644 index 00000000000..9e5f9b8ed77 --- /dev/null +++ b/litellm/llms/gigachat/passthrough/__init__.py @@ -0,0 +1,7 @@ +""" +GigaChat passthrough Module +""" + +from .transformation import GigaChatPassthroughConfig + +__all__ = ["GigaChatPassthroughConfig"] diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py new file mode 100644 index 00000000000..42bad43c79f --- /dev/null +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -0,0 +1,196 @@ +import json +from typing import TYPE_CHECKING, List, Optional, Tuple, cast + +import httpx + +from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig +from litellm.llms.gigachat.authenticator import get_access_token +from litellm.llms.gigachat.chat.streaming import GigaChatModelResponseIterator +from litellm.llms.gigachat.utils import get_api_base +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues + +if TYPE_CHECKING: + from httpx import URL, Response + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import CostResponseTypes + + +class GigaChatPassthroughConfig(BasePassthroughConfig): + def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: + return request_data.get("stream", False) + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + endpoint: str, + request_query_params: Optional[dict], + litellm_params: dict, + ) -> Tuple["URL", str]: + """Get complete API URL for chat completions.""" + base_target_url = self.get_api_base(api_base) + + if base_target_url is None: + raise Exception("GigaChat api base not found") + + litellm_metadata = litellm_params.get("litellm_metadata") or {} + model_group = litellm_metadata.get("model_group") + if model_group and model_group in endpoint: + endpoint = endpoint.replace(model_group, model) + + complete_url = f"{base_target_url}/chat/completions" + return ( + httpx.URL(complete_url), + base_target_url, + ) + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + """ + Set up headers with OAuth token. + """ + # Get access token + access_token = get_access_token( + credentials=api_key, litellm_params=litellm_params + ) + + headers["Authorization"] = f"Bearer {access_token}" + headers["Content-Type"] = "application/json" + headers["Accept"] = "application/json" + + return headers + + def logging_non_streaming_response( + self, + model: str, + custom_llm_provider: str, + httpx_response: "Response", + request_data: dict, + logging_obj: "LiteLLMLoggingObj", + endpoint: str, + ) -> Optional["CostResponseTypes"]: + from litellm import encoding + from litellm.types.utils import LlmProviders, ModelResponse + from litellm.utils import ProviderConfigManager + + provider_chat_config = ProviderConfigManager.get_provider_chat_config( + provider=LlmProviders(custom_llm_provider), + model=model, + ) + + if provider_chat_config is None: + raise ValueError(f"No provider config found for model: {model}") + + litellm_model_response: ModelResponse = provider_chat_config.transform_response( + model=model, + messages=request_data.get("messages", []), + raw_response=httpx_response, + model_response=ModelResponse(), + logging_obj=logging_obj, + optional_params={}, + litellm_params={}, + api_key="", + request_data=request_data, + encoding=encoding, + ) + + return litellm_model_response + + def handle_logging_collected_chunks( + self, + all_chunks: List[str], + litellm_logging_obj: "LiteLLMLoggingObj", + model: str, + custom_llm_provider: str, + endpoint: str, + ) -> Optional["CostResponseTypes"]: + """ + 1. Convert all_chunks to a ModelResponseStream + 2. combine model_response_stream to model_response + 3. Return the model_response + """ + + from litellm.litellm_core_utils.streaming_handler import ( + convert_generic_chunk_to_model_response_stream, + generic_chunk_has_all_required_fields, + ) + from litellm.main import stream_chunk_builder + from litellm.types.utils import GenericStreamingChunk, ModelResponseStream + + all_translated_chunks = [] + + for chunk in all_chunks: + if isinstance(chunk, bytes): + chunk = chunk.decode("utf-8", errors="ignore") + + if isinstance(chunk, str): + chunk = chunk.strip() + if not chunk or chunk == "[DONE]": + continue + if chunk.startswith("data: "): + chunk = chunk[6:] + try: + message = json.loads(chunk) + except json.JSONDecodeError: + continue + elif isinstance(chunk, dict): + message = chunk + else: + continue + + gigachat_iterator = GigaChatModelResponseIterator( + streaming_response=None, + sync_stream=False, + ) + translated_chunk = gigachat_iterator.chunk_parser(chunk=message) + + if isinstance( + translated_chunk, dict + ) and generic_chunk_has_all_required_fields(cast(dict, translated_chunk)): + chunk_obj = convert_generic_chunk_to_model_response_stream( + cast(GenericStreamingChunk, translated_chunk) + ) + elif isinstance(translated_chunk, ModelResponseStream): + chunk_obj = translated_chunk + else: + continue + + all_translated_chunks.append(chunk_obj) + + if len(all_translated_chunks) > 0: + model_response = stream_chunk_builder( + chunks=all_translated_chunks, + logging_obj=litellm_logging_obj, + ) + return model_response + return None + + @staticmethod + def get_api_base(api_base: Optional[str] = None) -> Optional[str]: + return get_api_base(api_base) + + @staticmethod + def get_api_key( + api_key: Optional[str] = None, + ) -> Optional[str]: + return api_key or get_secret_str("GIGACHAT_API_KEY") + + @staticmethod + def get_base_model(model: str) -> Optional[str]: + return model + + def get_models( + self, api_key: Optional[str] = None, api_base: Optional[str] = None + ) -> List[str]: + return super().get_models(api_key, api_base) diff --git a/litellm/main.py b/litellm/main.py index ddd37b47536..94326b2bc8c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1590,6 +1590,9 @@ def completion( # type: ignore # noqa: PLR0915 litellm_request_debug=kwargs.get("litellm_request_debug", False), tpm=kwargs.get("tpm"), rpm=kwargs.get("rpm"), + gigachat_scope=kwargs.get("gigachat_scope"), + gigachat_auth_url=kwargs.get("gigachat_auth_url"), + gigachat_access_token=kwargs.get("gigachat_access_token"), ) cast(LiteLLMLoggingObj, logging).update_environment_variables( model=model, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0bbee56d5e0..2389e62baf0 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -405,6 +405,7 @@ class LiteLLMRoutes(enum.Enum): "/vllm", "/mistral", "/milvus", + "/gigachat", ] ######################################################### diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 1ef866486ec..0a8f651522f 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -22,6 +22,7 @@ from litellm.constants import ( ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS, BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES, ) +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._types import * @@ -2369,3 +2370,244 @@ def create_generic_websocket_passthrough_endpoint( _forward_headers=forward_headers, cost_per_request=cost_per_request, ) + + +@router.api_route( + "/gigachat/{endpoint:path}", + methods=["GET", "POST", "PUT", "DELETE", "PATCH"], + tags=["Gigachat Pass-through", "pass-through"], +) +async def gigachat_proxy_route( + endpoint: str, + request: Request, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + [Docs](https://docs.litellm.ai/docs/pass_through/gigachat) + """ + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.proxy_server import ( + general_settings, + llm_router, + proxy_config, + proxy_logging_obj, + select_data_generator, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, + version, + ) + + ## check for streaming + request_body = await get_request_body(request) + is_router_model = is_passthrough_request_using_router_model( + request_body, llm_router + ) + + model = request_body.get("model") + if not model: + msg = "Model is required" + raise ValueError(msg) + + # If router model, use dedicated router passthrough handler + # This uses the same common processing path as non-router models + if is_router_model and llm_router: + return await handle_gigachat_passthrough_router_model( + model=model, + endpoint=endpoint, + request=request, + request_body=request_body, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + + # Fall back to existing implementation for direct GigaChat models + verbose_proxy_logger.debug( + f"Gigachat passthrough: Using direct Gigachat model '{model}' for endpoint '{endpoint}'" + ) + + data: Dict[str, Any] = {} + + data["method"] = request.method + data["endpoint"] = endpoint + data["json"] = request_body + data["custom_llm_provider"] = "gigachat" + + client = get_async_httpx_client( # type: ignore + llm_provider=LlmProviders.GIGACHAT, + params={ + "timeout": httpx.Timeout(timeout=600.0, connect=5.0), + "ssl_verify": False, + }, + ) + data["http_client"] = client + + base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) + + try: + result = await base_llm_response_processor.base_passthrough_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=model, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + + return result + except Exception as e: + raise await base_llm_response_processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def handle_gigachat_passthrough_router_model( + model: str, + endpoint: str, + request: Request, + request_body: dict, + llm_router: litellm.Router, + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj, + general_settings: dict, + proxy_config, + select_data_generator, + user_model: Optional[str], + user_temperature: Optional[float], + user_request_timeout: Optional[float], + user_max_tokens: Optional[int], + user_api_base: Optional[str], + version: Optional[str], +) -> Union[Response, StreamingResponse]: + """ + Handle Gigachat passthrough for router models (models defined in config.yaml). + + Uses the same common processing path as non-router models to ensure + metadata and hooks are properly initialized. + + Args: + model: The router model name (e.g., "gigachat/gigachat-2") + endpoint: The Gigachat endpoint path (e.g., "/chat/completions") + request: The FastAPI request object + request_body: The parsed request body + llm_router: The LiteLLM router instance + user_api_key_dict: The user API key authentication dictionary + (additional args for common processing) + + Returns: + Response or StreamingResponse depending on endpoint type + """ + from fastapi import Response as FastAPIResponse + + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + # Detect streaming based on request body + is_streaming = request_body.get("stream", False) + + data: Dict[str, Any] = await _read_request_body(request=request) + if user_api_key_dict is not None: + if data.get("metadata") is None: + data["metadata"] = {} + 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 + + verbose_proxy_logger.debug( + f"Gigachat router passthrough: model='{model}', endpoint='{endpoint}', streaming={is_streaming}" + ) + + # Use the common processing path (same as non-router models) + # This ensures all metadata, hooks, and logging are properly initialized + + data["model"] = model + data["method"] = request.method + data["endpoint"] = endpoint + data["json"] = request_body + data["custom_llm_provider"] = "gigachat" + + client = get_async_httpx_client( # type: ignore + llm_provider=LlmProviders.GIGACHAT, + params={ + "timeout": httpx.Timeout(timeout=600.0, connect=5.0), + "ssl_verify": False, + }, + ) + + data["http_client"] = client + base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) + + # Use the common passthrough processing to handle metadata and hooks + # This also handles all response formatting (streaming/non-streaming) and exceptions + try: + result = await base_llm_response_processor.base_passthrough_process_llm_request( + request=request, + fastapi_response=FastAPIResponse(), + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=model, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + + if isinstance(result, StreamingResponse): + if result.headers.get("Content-Type") is None: + result.headers["Content-Type"] = "text/event-stream; charset=utf-8" + return result + + return result + except Exception as e: + # Use common exception handling + raise await base_llm_response_processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 860593a6eab..bc3820b714f 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -1172,6 +1172,68 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "GIGACHAT", + "provider_display_name": "GigaChat", + "litellm_provider": "gigachat", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": null, + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": false, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "gigachat_scope", + "label": "Scope", + "placeholder": null, + "tooltip": null, + "required": false, + "field_type": "select", + "options": [ + "GIGACHAT_API_PERS", + "GIGACHAT_API_B2B", + "GIGACHAT_API_CORP" + ], + "default_value": "GIGACHAT_API_PERS" + }, + { + "key": "gigachat_auth_url", + "label": "Auth URL", + "placeholder": null, + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "gigachat_access_token", + "label": "Access token", + "placeholder": null, + "tooltip": "Disable OAuth, provide value to authorization.", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "GigaChat-2" + }, { "provider": "GITHUB", "provider_display_name": "Github", diff --git a/litellm/utils.py b/litellm/utils.py index 09df88f0ceb..efcba9357f2 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8675,6 +8675,12 @@ class ProviderConfigManager: ) return AzurePassthroughConfig() + elif LlmProviders.GIGACHAT == provider: + from litellm.llms.gigachat.passthrough.transformation import ( + GigaChatPassthroughConfig, + ) + + return GigaChatPassthroughConfig() return None @staticmethod 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 8acfa2231cc..637dc8e2f06 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 @@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest from fastapi import Request, Response +from fastapi.responses import StreamingResponse from fastapi.testclient import TestClient sys.path.insert( @@ -15,12 +16,14 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( BaseOpenAIPassThroughHandler, RouteChecks, bedrock_llm_proxy_route, create_pass_through_route, cursor_proxy_route, + gigachat_proxy_route, llm_passthrough_factory_proxy_route, milvus_proxy_route, openai_proxy_route, @@ -1466,6 +1469,176 @@ class TestVLLMProxyRoute: mock_factory_route.assert_awaited_once() +class TestGigachatProxyRoute: + @pytest.mark.asyncio + @patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body", + return_value={"model": "router-model", "stream": False}, + ) + @patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_passthrough_request_using_router_model", + return_value=True, + ) + @patch("litellm.proxy.proxy_server.llm_router") + async def test_gigachat_proxy_route_with_router_model( + self, mock_llm_router, mock_is_router, mock_get_body + ): + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.headers = {"content-type": "application/json"} + mock_request.query_params = {} + mock_fastapi_response = MagicMock(spec=Response) + mock_user_api_key_dict = MagicMock() + mock_llm_router.allm_passthrough_route = AsyncMock( + return_value=httpx.Response(200, json={"response": "success"}) + ) + + result = await gigachat_proxy_route( + endpoint="/chat/completions", + request=mock_request, + fastapi_response=mock_fastapi_response, + user_api_key_dict=mock_user_api_key_dict, + ) + + mock_is_router.assert_called_once() + mock_llm_router.allm_passthrough_route.assert_awaited_once() + assert isinstance(result, Response) + + @pytest.mark.asyncio + @patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body", + return_value={"model": "other-model"}, + ) + @patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_passthrough_request_using_router_model", + return_value=False, + ) + @patch( + "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing.base_passthrough_process_llm_request", + new_callable=AsyncMock, + ) + async def test_gigachat_proxy_route_fallback_to_http_pass_through( + self, + mock_base_passthrough, + mock_is_router, + mock_get_body, + ): + mock_request = MagicMock(spec=Request) + mock_fastapi_response = MagicMock(spec=Response) + mock_user_api_key_dict = MagicMock() + + expected_response = Response( + content=b'{"response": "success"}', + status_code=200, + media_type="application/json", + ) + mock_base_passthrough.return_value = expected_response + + result = await gigachat_proxy_route( + endpoint="/chat/completions", + request=mock_request, + fastapi_response=mock_fastapi_response, + user_api_key_dict=mock_user_api_key_dict, + ) + + assert isinstance(result, Response) + assert result.status_code == 200 + assert result.body == b'{"response": "success"}' + mock_base_passthrough.assert_awaited_once() + + @pytest.mark.asyncio + async def test_allm_passthrough_streaming_preserves_upstream_headers(self): + async def _stream() -> bytes: + yield b'data: {"id":"1"}\n\n' + + class MockPassthroughStreamingResponse: + def __init__(self): + self.status_code = 201 + self.headers = { + "content-type": "text/event-stream; charset=utf-8", + "x-request-id": "req-123", + "x-ratelimit-remaining-requests": "77", + "transfer-encoding": "chunked", + "content-encoding": "gzip", + } + self._iterator = _stream() + + def __aiter__(self): + return self + + async def __anext__(self): + return await self._iterator.__anext__() + + processor = ProxyBaseLLMRequestProcessing( + data={ + "model": "some-provider/model", + "stream": True, + "litellm_call_id": "call-123", + "litellm_logging_obj": MagicMock(litellm_call_id="call-123"), + } + ) + + mock_request = MagicMock(spec=Request) + mock_request.headers = {"content-type": "application/json"} + mock_fastapi_response = MagicMock(spec=Response) + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.allowed_model_region = "" + mock_user_api_key_dict.spend = 0.0 + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + mock_proxy_logging_obj.update_request_status = AsyncMock(return_value=None) + mock_proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value={"x-test-callback-header": "callback-value"} + ) + + streaming_response = MockPassthroughStreamingResponse() + + async def _fake_route_request(*args, **kwargs): + async def _inner(): + return streaming_response + + return _inner() + + with patch.object( + processor, + "common_processing_pre_call_logic", + new=AsyncMock( + return_value=( + processor.data, + processor.data["litellm_logging_obj"], + ) + ), + ), patch( + "litellm.proxy.common_request_processing.route_request", + new=_fake_route_request, + ), patch( + "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing.get_custom_headers", + return_value={"x-litellm-call-id": "call-123"}, + ): + result = await processor.base_passthrough_process_llm_request( + request=mock_request, + fastapi_response=mock_fastapi_response, + user_api_key_dict=mock_user_api_key_dict, + proxy_logging_obj=mock_proxy_logging_obj, + general_settings={}, + proxy_config=MagicMock(), + select_data_generator=MagicMock(), + llm_router=None, + model="some-provider/model", + version="test-version", + ) + + assert isinstance(result, StreamingResponse) + assert result.status_code == 201 + assert result.headers["content-type"] == "text/event-stream; charset=utf-8" + assert result.headers["x-request-id"] == "req-123" + assert result.headers["x-ratelimit-remaining-requests"] == "77" + assert result.headers["x-litellm-call-id"] == "call-123" + assert result.headers["x-test-callback-header"] == "callback-value" + assert "transfer-encoding" not in result.headers + assert "content-encoding" not in result.headers + + class TestForwardHeaders: """ Test cases for _forward_headers parameter in passthrough endpoints diff --git a/ui/litellm-dashboard/public/assets/logos/gigachat.svg b/ui/litellm-dashboard/public/assets/logos/gigachat.svg new file mode 100644 index 00000000000..e7abe47b221 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/gigachat.svg @@ -0,0 +1,27 @@ + + + + + diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index e833d0eb4fb..2b807c304f8 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -40,6 +40,7 @@ export enum Providers { FireworksAI = "Fireworks AI", FRIENDLIAI = "Friendliai", GALADRIEL = "Galadriel", + GIGACHAT = "GigaChat", GITHUB_COPILOT = "Github Copilot", Google_AI_Studio = "Google AI Studio", GradientAI = "GradientAI", @@ -146,6 +147,7 @@ export const provider_map: Record = { FireworksAI: "fireworks_ai", FRIENDLIAI: "friendliai", GALADRIEL: "galadriel", + GIGACHAT: "gigachat", GITHUB_COPILOT: "github_copilot", Google_AI_Studio: "gemini", GradientAI: "gradient_ai", @@ -247,6 +249,7 @@ export const providerLogoMap: Record = { [Providers.FEATHERLESS_AI]: `${asset_logos_folder}featherless.svg`, [Providers.FireworksAI]: `${asset_logos_folder}fireworks.svg`, [Providers.FRIENDLIAI]: `${asset_logos_folder}friendli.svg`, + [Providers.GIGACHAT]: `${asset_logos_folder}gigachat.svg`, [Providers.GITHUB_COPILOT]: `${asset_logos_folder}github_copilot.svg`, [Providers.Google_AI_Studio]: `${asset_logos_folder}google.svg`, [Providers.GradientAI]: `${asset_logos_folder}gradientai.svg`, From 52796fb060e40a4fb8b60e0cdf884a5e99b9fdcb Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Thu, 16 Apr 2026 15:33:40 +0000 Subject: [PATCH 004/120] feature: add logging when passthrough return Result --- litellm/litellm_core_utils/litellm_logging.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e84c1e13a8b..560f5478d76 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1852,6 +1852,11 @@ class Logging(LiteLLMLoggingBaseClass): logging_result = self.normalize_logging_result(result=result) + if isinstance(result, Response) and isinstance( + logging_result, ModelResponse + ): + result = logging_result + if ( standard_logging_object is None and result is not None @@ -5569,7 +5574,13 @@ def get_standard_logging_object_payload( def emit_standard_logging_payload(payload: StandardLoggingPayload): if os.getenv("LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD"): - print(json.dumps(payload, indent=4)) # noqa + try: + print(json.dumps(payload, indent=4, default=str)) # noqa + except Exception as e: + print( + "Error serializing standard logging payload for debug output:", + str(e), + ) def get_standard_logging_metadata( From 211388ac51b58383e425b3f547344da81054ddee Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Thu, 16 Apr 2026 17:03:02 +0000 Subject: [PATCH 005/120] fix: return headers for streaming passthrough and use HTTP client in llm_passthrough_route --- litellm/passthrough/main.py | 118 +++++++++++------- litellm/proxy/common_request_processing.py | 35 +++++- .../test_hosted_vllm_passthrough.py | 2 +- .../test_async_streaming_error_propagation.py | 95 +++++++------- .../passthrough/test_passthrough_main.py | 50 ++++---- 5 files changed, 181 insertions(+), 119 deletions(-) diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index edee50bdfc4..e58c1cab744 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -4,11 +4,11 @@ This module is used to pass through requests to the LLM APIs. import asyncio import contextvars +from collections.abc import AsyncIterator from functools import partial from typing import ( TYPE_CHECKING, Any, - AsyncGenerator, Coroutine, Generator, List, @@ -35,6 +35,61 @@ if TYPE_CHECKING: from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig +class _AsyncPassthroughStreamingResponse(AsyncIterator[bytes]): + """ + Async iterator wrapper that preserves upstream response metadata for streaming. + """ + + def __init__( + self, + response: httpx.Response, + litellm_logging_obj: "LiteLLMLoggingObj", + provider_config: "BasePassthroughConfig", + ) -> None: + self.response = response + self.headers = response.headers + self.status_code = response.status_code + self._litellm_logging_obj = litellm_logging_obj + self._provider_config = provider_config + self._iterator = response.aiter_bytes() + self._raw_bytes: List[bytes] = [] + self._flush_started = False + + def __aiter__(self) -> "_AsyncPassthroughStreamingResponse": + return self + + async def __anext__(self) -> bytes: + try: + chunk = await self._iterator.__anext__() + self._raw_bytes.append(chunk) + return chunk + except StopAsyncIteration: + self._start_flush() + raise + except Exception: + try: + await self.response.aclose() + except Exception: + pass + raise + + def _start_flush(self) -> None: + if self._flush_started: + return + + self._flush_started = True + asyncio.create_task( + self._litellm_logging_obj.async_flush_passthrough_collected_chunks( + raw_bytes=self._raw_bytes, + provider_config=self._provider_config, + ) + ) + + async def aclose(self) -> None: + self._start_flush() + await self.response.aclose() + + @client async def allm_passthrough_route( *, @@ -52,9 +107,9 @@ async def allm_passthrough_route( json: Optional[Any] = None, params: Optional[QueryParamTypes] = None, cookies: Optional[CookieTypes] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + http_client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, **kwargs, -) -> Union[httpx.Response, AsyncGenerator[Any, Any]]: +) -> Union[httpx.Response, AsyncIterator[bytes]]: """ Async: Reranks a list of documents based on their relevance to the query """ @@ -98,7 +153,7 @@ async def allm_passthrough_route( json=json, params=params, cookies=cookies, - client=client, + http_client=http_client, **kwargs, ) @@ -178,14 +233,14 @@ def llm_passthrough_route( json: Optional[Any] = None, params: Optional[QueryParamTypes] = None, cookies: Optional[CookieTypes] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + http_client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, **kwargs, ) -> Union[ httpx.Response, Coroutine[Any, Any, httpx.Response], - Coroutine[Any, Any, Union[httpx.Response, AsyncGenerator[Any, Any]]], + Coroutine[Any, Any, Union[httpx.Response, AsyncIterator[bytes]]], Generator[Any, Any, Any], - AsyncGenerator[Any, Any], + AsyncIterator[bytes], ]: """ Pass through requests to the LLM APIs. @@ -200,11 +255,11 @@ def llm_passthrough_route( _is_async = allm_passthrough_route - if client is None: + if http_client is None: if _is_async: - client = litellm.module_level_aclient + http_client = litellm.module_level_aclient else: - client = litellm.module_level_client + http_client = litellm.module_level_client litellm_logging_obj = cast("LiteLLMLoggingObj", kwargs.get("litellm_logging_obj")) @@ -286,7 +341,7 @@ def llm_passthrough_route( if json and isinstance(json, dict) and "model" in json: json["model"] = model - request = client.client.build_request( + request = http_client.client.build_request( method=method, url=updated_url, content=signed_json_body if signed_json_body is not None else content, @@ -323,7 +378,7 @@ def llm_passthrough_route( if _is_async: # Return the coroutine to be awaited by the caller return _async_passthrough_request( - client=client, + client=http_client, request=request, is_streaming_request=is_streaming_request, litellm_logging_obj=litellm_logging_obj, @@ -331,7 +386,7 @@ def llm_passthrough_route( ) else: # Sync path - client.client.send returns Response directly - response: httpx.Response = client.client.send(request=request, stream=is_streaming_request) # type: ignore + response: httpx.Response = http_client.client.send(request=request, stream=is_streaming_request) # type: ignore response.raise_for_status() if ( @@ -356,7 +411,7 @@ async def _async_passthrough_request( is_streaming_request: bool, litellm_logging_obj: "LiteLLMLoggingObj", provider_config: "BasePassthroughConfig", -) -> Union[httpx.Response, AsyncGenerator[Any, Any]]: +) -> Union[httpx.Response, AsyncIterator[bytes]]: """ Handle async passthrough requests. Uses async client to send request and properly handles streaming. @@ -367,9 +422,10 @@ async def _async_passthrough_request( # Check if it's a coroutine and await it if asyncio.iscoroutine(response_result): if is_streaming_request: - # Pass the coroutine to _async_streaming which will await it - return _async_streaming( - response=response_result, + iter_response = await response_result + iter_response.raise_for_status() + return _AsyncPassthroughStreamingResponse( + response=iter_response, litellm_logging_obj=litellm_logging_obj, provider_config=provider_config, ) @@ -403,31 +459,3 @@ def _sync_streaming( ) except Exception as e: raise e - - -async def _async_streaming( - response: Coroutine[Any, Any, httpx.Response], - litellm_logging_obj: "LiteLLMLoggingObj", - provider_config: "BasePassthroughConfig", -): - iter_response = await response - try: - iter_response.raise_for_status() - raw_bytes: List[bytes] = [] - - async for chunk in iter_response.aiter_bytes(): # type: ignore - raw_bytes.append(chunk) - yield chunk - - asyncio.create_task( - litellm_logging_obj.async_flush_passthrough_collected_chunks( - raw_bytes=raw_bytes, - provider_config=provider_config, - ) - ) - except Exception: - try: - await iter_response.aclose() - except Exception: - pass - raise diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 037f913ad07..e582e7c79a2 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -490,6 +490,26 @@ class ProxyBaseLLMRequestProcessing: def __init__(self, data: dict): self.data = data + @staticmethod + def _merge_passthrough_streaming_headers( + response_headers: Optional[Any], + custom_headers: dict, + ) -> dict: + """ + Merge upstream passthrough headers with proxy/custom headers. + + Proxy/custom headers win on key collisions. + """ + excluded_headers = {"transfer-encoding", "content-encoding"} + + merged_headers = { + key: value + for key, value in dict(response_headers or {}).items() + if key.lower() not in excluded_headers + } + merged_headers.update(custom_headers) + return merged_headers + @staticmethod def get_custom_headers( *, @@ -1169,6 +1189,13 @@ class ProxyBaseLLMRequestProcessing: logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete # type: ignore[union-attr] if route_type == "allm_passthrough_route": + streaming_headers = custom_headers + if hasattr(response, "headers"): + streaming_headers = ProxyBaseLLMRequestProcessing._merge_passthrough_streaming_headers( + response_headers=getattr(response, "headers", None), + custom_headers=custom_headers, + ) + # Check if response is an async generator if self._is_streaming_response(response): if asyncio.iscoroutine(response): @@ -1180,15 +1207,17 @@ class ProxyBaseLLMRequestProcessing: # since we're dealing with raw binary data (e.g., AWS event streams) return StreamingResponse( content=generator, # type: ignore[arg-type] - status_code=status.HTTP_200_OK, - headers=custom_headers, + status_code=getattr( + response, "status_code", status.HTTP_200_OK + ), + headers=streaming_headers, ) else: # Traditional HTTP response with aiter_bytes return StreamingResponse( content=response.aiter_bytes(), # type: ignore[union-attr] status_code=response.status_code, # type: ignore[union-attr] - headers=custom_headers, + headers=streaming_headers, ) elif route_type == "anthropic_messages": # Check if response is actually a streaming response (async generator) diff --git a/tests/pass_through_tests/test_hosted_vllm_passthrough.py b/tests/pass_through_tests/test_hosted_vllm_passthrough.py index 746f103cc9e..003d2570400 100644 --- a/tests/pass_through_tests/test_hosted_vllm_passthrough.py +++ b/tests/pass_through_tests/test_hosted_vllm_passthrough.py @@ -63,7 +63,7 @@ async def test_allm_passthrough_route_with_hosted_vllm_model_does_not_raise(): "model": "anything", # will be replaced internally with normalized model "messages": [{"role": "user", "content": "Hello"}], }, - client=client, + http_client=client, ) # Then it should not raise and return a successful httpx.Response diff --git a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py b/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py index 8148edb633f..b5ab8279bd1 100644 --- a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py +++ b/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py @@ -1,12 +1,12 @@ """ -Tests for error propagation in _async_streaming passthrough routes. +Tests for error propagation in async passthrough streaming routes. -Verifies that HTTP 4xx/5xx errors from upstream (e.g. Azure 429 rate limits) -raise exceptions instead of being silently forwarded as raw bytes under HTTP 200. - -See: litellm/passthrough/main.py _async_streaming() +Verifies that streaming passthrough wrappers preserve the previous guarantees: +HTTP 4xx/5xx failures must raise instead of being silently forwarded as bytes, +and successful streaming responses should still yield chunks normally. """ +import asyncio import json from unittest.mock import AsyncMock, MagicMock @@ -50,73 +50,82 @@ def _make_mock_logging_obj(): @pytest.mark.asyncio -async def test_async_streaming_429_raises(): - """429 from upstream should raise HTTPStatusError, not yield error bytes.""" - from litellm.passthrough.main import _async_streaming - +async def test_async_passthrough_wrapper_429_raises_before_iteration(): + """429 from upstream should be raised before the wrapper is constructed.""" error_body = json.dumps( {"error": {"code": "429", "message": "Rate limit exceeded."}} ).encode() mock_response = _make_mock_response(429, error_body) - async def response_coro(): - return mock_response - - chunks = [] with pytest.raises(httpx.HTTPStatusError) as exc_info: - async for chunk in _async_streaming( - response=response_coro(), - litellm_logging_obj=_make_mock_logging_obj(), - provider_config=MagicMock(), - ): - chunks.append(chunk) + mock_response.raise_for_status() assert exc_info.value.response.status_code == 429 - assert len(chunks) == 0 @pytest.mark.asyncio -async def test_async_streaming_500_raises(): - """500 from upstream should also raise, not yield error bytes.""" - from litellm.passthrough.main import _async_streaming - +async def test_async_passthrough_wrapper_500_raises_before_iteration(): + """500 from upstream should be raised before the wrapper is constructed.""" error_body = json.dumps( {"error": {"code": "500", "message": "Internal server error"}} ).encode() mock_response = _make_mock_response(500, error_body) - async def response_coro(): - return mock_response - with pytest.raises(httpx.HTTPStatusError) as exc_info: - async for _ in _async_streaming( - response=response_coro(), - litellm_logging_obj=_make_mock_logging_obj(), - provider_config=MagicMock(), - ): - pass + mock_response.raise_for_status() assert exc_info.value.response.status_code == 500 @pytest.mark.asyncio -async def test_async_streaming_200_yields_chunks(): +async def test_async_passthrough_wrapper_200_yields_chunks(): """Successful 200 streaming responses should continue to work normally.""" - from litellm.passthrough.main import _async_streaming + from litellm.passthrough.main import _AsyncPassthroughStreamingResponse sse_data = b'data: {"type":"response.created"}\n\ndata: [DONE]\n\n' mock_response = _make_mock_response(200, sse_data) - - async def response_coro(): - return mock_response + mock_logging_obj = _make_mock_logging_obj() + async_stream = _AsyncPassthroughStreamingResponse( + response=mock_response, + litellm_logging_obj=mock_logging_obj, + provider_config=MagicMock(), + ) chunks = [] - async for chunk in _async_streaming( - response=response_coro(), - litellm_logging_obj=_make_mock_logging_obj(), - provider_config=MagicMock(), - ): + async for chunk in async_stream: chunks.append(chunk) + await asyncio.sleep(0) + assert len(chunks) == 1 assert b"response.created" in chunks[0] + mock_logging_obj.async_flush_passthrough_collected_chunks.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_async_passthrough_wrapper_closes_response_on_iteration_error(): + """Wrapper should close the upstream response if iteration raises.""" + from litellm.passthrough.main import _AsyncPassthroughStreamingResponse + + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = httpx.Headers({"content-type": "text/event-stream"}) + mock_response.aclose = AsyncMock() + + async def _failing_aiter_bytes(): + raise RuntimeError("stream failed") + yield b"" + + mock_response.aiter_bytes = _failing_aiter_bytes + + async_stream = _AsyncPassthroughStreamingResponse( + response=mock_response, + litellm_logging_obj=_make_mock_logging_obj(), + provider_config=MagicMock(), + ) + + with pytest.raises(RuntimeError, match="stream failed"): + async for chunk in async_stream: + _ = chunk + + mock_response.aclose.assert_awaited_once() diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/test_litellm/passthrough/test_passthrough_main.py index 489357149c5..bdfb3a348b3 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/test_litellm/passthrough/test_passthrough_main.py @@ -40,7 +40,7 @@ def test_llm_passthrough_route(): "model": "my-custom-model", "messages": [{"role": "user", "content": "Hello, world!"}], }, - client=client, + http_client=client, ) mock_post.call_args.kwargs[ @@ -90,7 +90,7 @@ def test_bedrock_application_inference_profile_url_encoding(): endpoint="model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse", method="POST", custom_llm_provider="bedrock", - client=client, + http_client=client, litellm_logging_obj=mock_logging_obj, ) @@ -144,7 +144,7 @@ def test_bedrock_non_application_inference_profile_no_encoding(): endpoint="model/anthropic.claude-3-sonnet-20240229-v1:0/converse", method="POST", custom_llm_provider="bedrock", - client=client, + http_client=client, litellm_logging_obj=mock_logging_obj, ) @@ -486,7 +486,7 @@ def test_azure_with_custom_api_base_and_key(): "model": "gpt-4.1", "messages": [{"role": "user", "content": "Hello!"}], }, - client=client, + http_client=client, litellm_logging_obj=mock_logging_obj, ) @@ -558,7 +558,7 @@ def test_content_param_forwarded_to_build_request(): content=raw_content, data=None, json=None, - client=client, + http_client=client, litellm_logging_obj=mock_logging_obj, ) @@ -611,13 +611,14 @@ async def test_allm_passthrough_route_429_streaming_raises(): Regression test: Azure 429 during streaming must raise HTTPStatusError, not be silently forwarded as raw bytes under HTTP 200. - Before the fix, _async_streaming() would yield the 429 error JSON as - chunks and allm_passthrough_route returned an async generator. The - caller (azure_proxy_route) wrapped it in StreamingResponse(status_code=200), + Before the fix, the async passthrough streaming path would yield the 429 + error JSON as chunks and allm_passthrough_route returned a streaming + iterator. The caller (azure_proxy_route) wrapped it in + StreamingResponse(status_code=200), so the client saw HTTP 200 + unparseable SSE body → silent task_complete(null). - After the fix, raise_for_status() fires inside _async_streaming() before - any chunks are yielded, so the exception propagates all the way up. + After the fix, raise_for_status() fires before the streaming wrapper is + returned, so the exception propagates all the way up. """ mock_provider_config = MagicMock() mock_provider_config.get_complete_url.return_value = ( @@ -640,6 +641,7 @@ async def test_allm_passthrough_route_429_streaming_raises(): mock_logging_obj = MagicMock() mock_logging_obj.update_environment_variables = MagicMock() mock_logging_obj.async_flush_passthrough_collected_chunks = AsyncMock() + mock_logging_obj.async_failure_handler = AsyncMock() with patch( "litellm.utils.ProviderConfigManager.get_provider_passthrough_config", @@ -660,23 +662,17 @@ async def test_allm_passthrough_route_429_streaming_raises(): ), patch.object( async_client.client, "build_request", mock_build_request ): - result = await allm_passthrough_route( - model="azure/gpt-4", - endpoint="openai/deployments/gpt-4/responses", - method="POST", - custom_llm_provider="azure", - api_base="https://my-azure.openai.azure.com", - api_key="fake-azure-key", - json={"model": "gpt-4", "input": "hello", "stream": True}, - client=async_client, - litellm_logging_obj=mock_logging_obj, - ) - - # result is an async generator — consuming it must raise, not silently yield error bytes - chunks = [] with pytest.raises(httpx.HTTPStatusError) as exc_info: - async for chunk in result: # type: ignore[union-attr] - chunks.append(chunk) + await allm_passthrough_route( + model="azure/gpt-4", + endpoint="openai/deployments/gpt-4/responses", + method="POST", + custom_llm_provider="azure", + api_base="https://my-azure.openai.azure.com", + api_key="fake-azure-key", + json={"model": "gpt-4", "input": "hello", "stream": True}, + http_client=async_client, + litellm_logging_obj=mock_logging_obj, + ) assert exc_info.value.response.status_code == 429 - assert len(chunks) == 0, "No chunks should be yielded before the 429 raises" From c4e5dbc38bfe5c3c4172b8f95f738196d4b6c12c Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Fri, 17 Apr 2026 08:31:03 +0300 Subject: [PATCH 006/120] fix gigachat passthrough url, authenticator, get_api_base --- litellm/llms/gigachat/authenticator.py | 6 ------ litellm/llms/gigachat/passthrough/transformation.py | 6 +----- litellm/llms/gigachat/utils.py | 1 - 3 files changed, 1 insertion(+), 12 deletions(-) diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index 83824c72fa7..1821a209fed 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -130,12 +130,6 @@ async def get_access_token_async( litellm_params: Optional[dict] = None, ) -> str: """Async version of get_access_token.""" - credentials = credentials or _get_credentials() - if not credentials: - raise GigaChatAuthError( - status_code=401, - message="GigaChat credentials not provided. Set GIGACHAT_CREDENTIALS or GIGACHAT_API_KEY environment variable.", - ) if not litellm_params: litellm_params = {} diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index 42bad43c79f..baf287b161f 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -36,12 +36,8 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): if base_target_url is None: raise Exception("GigaChat api base not found") - litellm_metadata = litellm_params.get("litellm_metadata") or {} - model_group = litellm_metadata.get("model_group") - if model_group and model_group in endpoint: - endpoint = endpoint.replace(model_group, model) + complete_url = f"{base_target_url}/{endpoint.lstrip('/')}" - complete_url = f"{base_target_url}/chat/completions" return ( httpx.URL(complete_url), base_target_url, diff --git a/litellm/llms/gigachat/utils.py b/litellm/llms/gigachat/utils.py index e79083f16fe..b66d25cbc0b 100644 --- a/litellm/llms/gigachat/utils.py +++ b/litellm/llms/gigachat/utils.py @@ -30,6 +30,5 @@ def convert_usage(usage_data: dict[str, int]) -> Usage: ) -@staticmethod def get_api_base(api_base: Optional[str] = None) -> Optional[str]: return api_base or get_secret_str("GIGACHAT_API_BASE") or GIGACHAT_BASE_URL From 4c27c6772e19380c7a4738e04ed7f7c67ab32747 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Fri, 17 Apr 2026 09:19:27 +0300 Subject: [PATCH 007/120] fix SSRF via user-controlled OAuth URL --- .../pass_through_endpoints/llm_passthrough_endpoints.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 0a8f651522f..aa4951bc258 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2566,6 +2566,11 @@ async def handle_gigachat_passthrough_router_model( data["json"] = request_body data["custom_llm_provider"] = "gigachat" + # Remove sensitive keys from data + keys = ["gigachat_auth_url", "gigachat_access_token", "gigachat_scope"] + for key in keys: + data.pop(key, None) + client = get_async_httpx_client( # type: ignore llm_provider=LlmProviders.GIGACHAT, params={ From eff086d8e00b503c0dfc9ea0503f7df1a1cf9137 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Fri, 17 Apr 2026 10:05:09 +0300 Subject: [PATCH 008/120] improve gigachat proxy route error handling --- litellm/llms/gigachat/authenticator.py | 13 ++++++----- .../llm_passthrough_endpoints.py | 22 ++++++++++++------- 2 files changed, 21 insertions(+), 14 deletions(-) diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index 1821a209fed..ce6d4b3eca9 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -163,12 +163,13 @@ async def get_access_token_async( # Request new token token, expires_at = await _request_token_async(credentials, scope, auth_url) - # Cache token - ttl_seconds = max( - 0, (expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000 - ) - if ttl_seconds > 0: - _token_cache.set_cache(cache_key, (token, expires_at), ttl=ttl_seconds) + if expires_at: + # Cache token + ttl_seconds = max( + 0, (expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000 + ) + if ttl_seconds > 0: + _token_cache.set_cache(cache_key, (token, expires_at), ttl=ttl_seconds) return token diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index aa4951bc258..ed9a4ac3028 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2403,23 +2403,28 @@ async def gigachat_proxy_route( ## check for streaming request_body = await get_request_body(request) - is_router_model = is_passthrough_request_using_router_model( - request_body, llm_router - ) + is_router_model = False model = request_body.get("model") - if not model: - msg = "Model is required" - raise ValueError(msg) + if model: + is_router_model = is_passthrough_request_using_router_model( + request_body, llm_router + ) + elif any(word in endpoint for word in ("completions", "embeddings")): + raise HTTPException( + status_code=400, detail={"error": "Model is required in request body"} + ) + # If router model, use dedicated router passthrough handler # This uses the same common processing path as non-router models - if is_router_model and llm_router: + if model and is_router_model and llm_router: return await handle_gigachat_passthrough_router_model( model=model, endpoint=endpoint, request=request, request_body=request_body, + fastapi_response=fastapi_response, llm_router=llm_router, user_api_key_dict=user_api_key_dict, proxy_logging_obj=proxy_logging_obj, @@ -2490,6 +2495,7 @@ async def handle_gigachat_passthrough_router_model( endpoint: str, request: Request, request_body: dict, + fastapi_response: Response, llm_router: litellm.Router, user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj, @@ -2587,7 +2593,7 @@ async def handle_gigachat_passthrough_router_model( try: result = await base_llm_response_processor.base_passthrough_process_llm_request( request=request, - fastapi_response=FastAPIResponse(), + fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, proxy_logging_obj=proxy_logging_obj, llm_router=llm_router, From 120f51a39d5e30e55c6e2df41137cf86a6e85dc6 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Fri, 17 Apr 2026 12:13:54 +0300 Subject: [PATCH 009/120] add check for logging gigachat response, add sensitive keys for removing from data --- litellm/llms/gigachat/passthrough/transformation.py | 4 ++++ .../proxy/pass_through_endpoints/llm_passthrough_endpoints.py | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index baf287b161f..a9bc59fed0a 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -80,6 +80,10 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): from litellm.types.utils import LlmProviders, ModelResponse from litellm.utils import ProviderConfigManager + # cost tracking only for completions + if "completions" not in endpoint: + return None + provider_chat_config = ProviderConfigManager.get_provider_chat_config( provider=LlmProviders(custom_llm_provider), model=model, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index ed9a4ac3028..e5eb58941ad 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2573,7 +2573,7 @@ async def handle_gigachat_passthrough_router_model( data["custom_llm_provider"] = "gigachat" # Remove sensitive keys from data - keys = ["gigachat_auth_url", "gigachat_access_token", "gigachat_scope"] + keys = ["gigachat_auth_url", "gigachat_access_token", "gigachat_scope", "api_base", "api_key"] for key in keys: data.pop(key, None) From 9ecfac90debe64b9c5571b3f54737f31865be5ea Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Fri, 8 May 2026 14:35:16 +0000 Subject: [PATCH 010/120] add embedding cost tracking to gigachat proxy route --- .../gigachat/passthrough/transformation.py | 70 +++++++++++++------ 1 file changed, 48 insertions(+), 22 deletions(-) diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index a9bc59fed0a..0d3e440f253 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -9,6 +9,8 @@ from litellm.llms.gigachat.chat.streaming import GigaChatModelResponseIterator from litellm.llms.gigachat.utils import get_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import EmbeddingResponse + if TYPE_CHECKING: from httpx import URL, Response @@ -80,32 +82,56 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): from litellm.types.utils import LlmProviders, ModelResponse from litellm.utils import ProviderConfigManager - # cost tracking only for completions - if "completions" not in endpoint: - return None + # cost tracking only for completions and embeddings + if "completions" in endpoint: - provider_chat_config = ProviderConfigManager.get_provider_chat_config( - provider=LlmProviders(custom_llm_provider), - model=model, - ) + provider_chat_config = ProviderConfigManager.get_provider_chat_config( + provider=LlmProviders(custom_llm_provider), + model=model, + ) - if provider_chat_config is None: - raise ValueError(f"No provider config found for model: {model}") + if provider_chat_config is None: + raise ValueError(f"No provider config found for model: {model}") - litellm_model_response: ModelResponse = provider_chat_config.transform_response( - model=model, - messages=request_data.get("messages", []), - raw_response=httpx_response, - model_response=ModelResponse(), - logging_obj=logging_obj, - optional_params={}, - litellm_params={}, - api_key="", - request_data=request_data, - encoding=encoding, - ) + litellm_model_response: ModelResponse = provider_chat_config.transform_response( + model=model, + messages=request_data.get("messages", []), + raw_response=httpx_response, + model_response=ModelResponse(), + logging_obj=logging_obj, + optional_params={}, + litellm_params={}, + api_key="", + request_data=request_data, + encoding=encoding, + ) - return litellm_model_response + return litellm_model_response + + if "embeddings" in endpoint: + + provider_embedding_config = ProviderConfigManager.get_provider_embedding_config( + provider=LlmProviders(custom_llm_provider), + model=model, + ) + + if provider_embedding_config is None: + raise ValueError(f"No provider config found for model: {model}") + + litellm_embedding_response: EmbeddingResponse = provider_embedding_config.transform_embedding_response( + model=model, + raw_response=httpx_response, + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + optional_params={}, + api_key="", + request_data=request_data, + litellm_params={}, + ) + + return litellm_embedding_response + + return None def handle_logging_collected_chunks( self, From b334ebf35ea34cae9d618f67801b435e8ce7e985 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Tue, 19 May 2026 05:17:33 +0000 Subject: [PATCH 011/120] remove potential recursive function --- litellm/llms/gigachat/passthrough/transformation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index 0d3e440f253..43669355f2b 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -6,7 +6,7 @@ import httpx from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig from litellm.llms.gigachat.authenticator import get_access_token from litellm.llms.gigachat.chat.streaming import GigaChatModelResponseIterator -from litellm.llms.gigachat.utils import get_api_base +from litellm.llms.gigachat.utils import GIGACHAT_BASE_URL from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import EmbeddingResponse @@ -204,7 +204,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): @staticmethod def get_api_base(api_base: Optional[str] = None) -> Optional[str]: - return get_api_base(api_base) + return api_base or get_secret_str("GIGACHAT_API_BASE") or GIGACHAT_BASE_URL @staticmethod def get_api_key( From 25edb2f228864c434b07c59cfa68a3d6d0f9b819 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Tue, 19 May 2026 06:25:17 +0000 Subject: [PATCH 012/120] add gigachat to allowlist --- gateway/routes/allowlist.py | 1 + 1 file changed, 1 insertion(+) diff --git a/gateway/routes/allowlist.py b/gateway/routes/allowlist.py index cbbf55c9873..9f397c70ff7 100644 --- a/gateway/routes/allowlist.py +++ b/gateway/routes/allowlist.py @@ -84,6 +84,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = ( "/bedrock/", "/cohere/", "/gemini/", + "/gigachat/", "/google/", "/vertex_ai/", "/vertex-ai/", From ecd72752d57c1c9046b9f2bd13894c6d4f23006a Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Tue, 19 May 2026 18:22:26 +0000 Subject: [PATCH 013/120] fix format and lint --- litellm/litellm_core_utils/litellm_logging.py | 7 ++- litellm/llms/gigachat/authenticator.py | 4 +- litellm/llms/gigachat/chat/streaming.py | 6 +- .../gigachat/passthrough/transformation.py | 58 ++++++++++--------- .../llm_passthrough_endpoints.py | 13 +++-- 5 files changed, 50 insertions(+), 38 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 084b53cd904..70d5b443974 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5632,9 +5632,10 @@ def emit_standard_logging_payload(payload: StandardLoggingPayload): try: print(json.dumps(payload, indent=4, default=str)) # noqa except Exception as e: - print( - "Error serializing standard logging payload for debug output:", - str(e), + verbose_logger.exception( + "Error serializing standard logging payload for debug output: {}".format( + str(e) + ) ) diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index ce6d4b3eca9..a20da943ee1 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -147,9 +147,7 @@ async def get_access_token_async( ) scope = scope or litellm_params.get("gigachat_scope") or _get_scope() - auth_url = ( - auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url() - ) + auth_url = auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url() # Check cache cache_key = f"gigachat_token:{credentials[:16]}" diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index 2fd5fc22b01..9ef29b2c6cb 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -76,7 +76,11 @@ class GigaChatModelResponseIterator: usage_data = chunk.get("usage", {}) if usage_data: usage = convert_usage(usage_data) - usage_block = ChatCompletionUsageBlock(**usage.dict()) + usage_block = ChatCompletionUsageBlock( + prompt_tokens=usage.prompt_tokens, + completion_tokens=usage.completion_tokens, + total_tokens=usage.total_tokens, + ) if finish_reason is not None: is_finished = True diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index 43669355f2b..ff2c4ad4986 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -93,44 +93,50 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): if provider_chat_config is None: raise ValueError(f"No provider config found for model: {model}") - litellm_model_response: ModelResponse = provider_chat_config.transform_response( - model=model, - messages=request_data.get("messages", []), - raw_response=httpx_response, - model_response=ModelResponse(), - logging_obj=logging_obj, - optional_params={}, - litellm_params={}, - api_key="", - request_data=request_data, - encoding=encoding, + litellm_model_response: ModelResponse = ( + provider_chat_config.transform_response( + model=model, + messages=request_data.get("messages", []), + raw_response=httpx_response, + model_response=ModelResponse(), + logging_obj=logging_obj, + optional_params={}, + litellm_params={}, + api_key="", + request_data=request_data, + encoding=encoding, + ) ) return litellm_model_response - + if "embeddings" in endpoint: - - provider_embedding_config = ProviderConfigManager.get_provider_embedding_config( - provider=LlmProviders(custom_llm_provider), - model=model, + + provider_embedding_config = ( + ProviderConfigManager.get_provider_embedding_config( + provider=LlmProviders(custom_llm_provider), + model=model, + ) ) if provider_embedding_config is None: raise ValueError(f"No provider config found for model: {model}") - litellm_embedding_response: EmbeddingResponse = provider_embedding_config.transform_embedding_response( - model=model, - raw_response=httpx_response, - model_response=EmbeddingResponse(), - logging_obj=logging_obj, - optional_params={}, - api_key="", - request_data=request_data, - litellm_params={}, + litellm_embedding_response: EmbeddingResponse = ( + provider_embedding_config.transform_embedding_response( + model=model, + raw_response=httpx_response, + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + optional_params={}, + api_key="", + request_data=request_data, + litellm_params={}, + ) ) return litellm_embedding_response - + return None def handle_logging_collected_chunks( diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 15852890eec..46d8b32369f 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2461,7 +2461,7 @@ async def gigachat_proxy_route( is_router_model = False model = request_body.get("model") - if model: + if model: is_router_model = is_passthrough_request_using_router_model( request_body, llm_router ) @@ -2470,7 +2470,6 @@ async def gigachat_proxy_route( status_code=400, detail={"error": "Model is required in request body"} ) - # If router model, use dedicated router passthrough handler # This uses the same common processing path as non-router models if model and is_router_model and llm_router: @@ -2582,8 +2581,6 @@ async def handle_gigachat_passthrough_router_model( Returns: Response or StreamingResponse depending on endpoint type """ - from fastapi import Response as FastAPIResponse - from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing # Detect streaming based on request body @@ -2628,7 +2625,13 @@ async def handle_gigachat_passthrough_router_model( data["custom_llm_provider"] = "gigachat" # Remove sensitive keys from data - keys = ["gigachat_auth_url", "gigachat_access_token", "gigachat_scope", "api_base", "api_key"] + keys = [ + "gigachat_auth_url", + "gigachat_access_token", + "gigachat_scope", + "api_base", + "api_key", + ] for key in keys: data.pop(key, None) From e1e56961a31a1dde3f61b2eeffaa0cf776c3eae4 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Sun, 7 Jun 2026 12:19:11 +0000 Subject: [PATCH 014/120] revert rename http client in llm_passthrough_route --- litellm/passthrough/main.py | 18 +++++++++--------- .../test_hosted_vllm_passthrough.py | 2 +- .../passthrough/test_passthrough_main.py | 12 ++++++------ 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index c06297ec15b..8b20406397b 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -108,7 +108,7 @@ async def allm_passthrough_route( json: Optional[Any] = None, params: Optional[QueryParamTypes] = None, cookies: Optional[CookieTypes] = None, - http_client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, **kwargs, ) -> Union[httpx.Response, AsyncIterator[bytes]]: """ @@ -154,7 +154,7 @@ async def allm_passthrough_route( json=json, params=params, cookies=cookies, - http_client=http_client, + client=client, **kwargs, ) @@ -234,7 +234,7 @@ def llm_passthrough_route( json: Optional[Any] = None, params: Optional[QueryParamTypes] = None, cookies: Optional[CookieTypes] = None, - http_client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, **kwargs, ) -> Union[ httpx.Response, @@ -256,11 +256,11 @@ def llm_passthrough_route( _is_async = allm_passthrough_route - if http_client is None: + if client is None: if _is_async: - http_client = litellm.module_level_aclient + client = litellm.module_level_aclient else: - http_client = litellm.module_level_client + client = litellm.module_level_client litellm_logging_obj = cast("LiteLLMLoggingObj", kwargs.get("litellm_logging_obj")) @@ -342,7 +342,7 @@ def llm_passthrough_route( if json and isinstance(json, dict) and "model" in json: json["model"] = model - request = http_client.client.build_request( + request = client.client.build_request( method=method, url=updated_url, content=signed_json_body if signed_json_body is not None else content, @@ -379,7 +379,7 @@ def llm_passthrough_route( if _is_async: # Return the coroutine to be awaited by the caller return _async_passthrough_request( - client=http_client, + client=client, request=request, is_streaming_request=is_streaming_request, litellm_logging_obj=litellm_logging_obj, @@ -387,7 +387,7 @@ def llm_passthrough_route( ) else: # Sync path - client.client.send returns Response directly - response: httpx.Response = http_client.client.send(request=request, stream=is_streaming_request) # type: ignore + response: httpx.Response = client.client.send(request=request, stream=is_streaming_request) # type: ignore response.raise_for_status() if ( diff --git a/tests/pass_through_tests/test_hosted_vllm_passthrough.py b/tests/pass_through_tests/test_hosted_vllm_passthrough.py index e8d4c0d2ef0..272b4e1bb00 100644 --- a/tests/pass_through_tests/test_hosted_vllm_passthrough.py +++ b/tests/pass_through_tests/test_hosted_vllm_passthrough.py @@ -63,7 +63,7 @@ async def test_allm_passthrough_route_with_hosted_vllm_model_does_not_raise(): "model": "anything", # will be replaced internally with normalized model "messages": [{"role": "user", "content": "Hello"}], }, - http_client=client, + client=client, ) # Then it should not raise and return a successful httpx.Response diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/test_litellm/passthrough/test_passthrough_main.py index adcca96271e..3600aeed59e 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/test_litellm/passthrough/test_passthrough_main.py @@ -40,7 +40,7 @@ def test_llm_passthrough_route(): "model": "my-custom-model", "messages": [{"role": "user", "content": "Hello, world!"}], }, - http_client=client, + client=client, ) mock_post.call_args.kwargs[ @@ -94,7 +94,7 @@ def test_bedrock_application_inference_profile_url_encoding(): endpoint="model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse", method="POST", custom_llm_provider="bedrock", - http_client=client, + client=client, litellm_logging_obj=mock_logging_obj, ) @@ -152,7 +152,7 @@ def test_bedrock_non_application_inference_profile_no_encoding(): endpoint="model/anthropic.claude-3-sonnet-20240229-v1:0/converse", method="POST", custom_llm_provider="bedrock", - http_client=client, + client=client, litellm_logging_obj=mock_logging_obj, ) @@ -511,7 +511,7 @@ def test_azure_with_custom_api_base_and_key(): "model": "gpt-4.1", "messages": [{"role": "user", "content": "Hello!"}], }, - http_client=client, + client=client, litellm_logging_obj=mock_logging_obj, ) @@ -592,7 +592,7 @@ def test_content_param_forwarded_to_build_request(): content=raw_content, data=None, json=None, - http_client=client, + client=client, litellm_logging_obj=mock_logging_obj, ) @@ -717,7 +717,7 @@ async def test_allm_passthrough_route_429_streaming_raises(): api_base="https://my-azure.openai.azure.com", api_key="fake-azure-key", json={"model": "gpt-4", "input": "hello", "stream": True}, - http_client=async_client, + client=async_client, litellm_logging_obj=mock_logging_obj, ) From dfa07b7c277fb5aab03aa4c8571cef4d546dd8a8 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Fri, 12 Jun 2026 13:58:02 +0000 Subject: [PATCH 015/120] add http client propagation for allm_passthrough_route --- litellm/llms/gigachat/passthrough/transformation.py | 1 - .../proxy/pass_through_endpoints/llm_passthrough_endpoints.py | 4 ++-- litellm/router.py | 2 ++ 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index ff2c4ad4986..474622a01a8 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -11,7 +11,6 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import EmbeddingResponse - if TYPE_CHECKING: from httpx import URL, Response diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 7a0403194ab..04de332b295 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2521,7 +2521,7 @@ async def gigachat_proxy_route( "ssl_verify": False, }, ) - data["http_client"] = client + data["client"] = client base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) @@ -2652,7 +2652,7 @@ async def handle_gigachat_passthrough_router_model( }, ) - data["http_client"] = client + data["client"] = client base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) # Use the common passthrough processing to handle metadata and hooks diff --git a/litellm/router.py b/litellm/router.py index a92590d3dba..c9637d5e9f2 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6014,6 +6014,8 @@ class Router: **kwargs, ) elif call_type == "allm_passthrough_route": + if client: + kwargs["client"] = client return await self._ageneric_api_call_with_fallbacks( original_function=original_function, passthrough_on_no_deployment=True, From 81dbb3803c42c4f6c2f2eb1e22bb440758f5aaae Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Fri, 12 Jun 2026 14:43:40 +0000 Subject: [PATCH 016/120] fix gigachat stream usage --- litellm/llms/gigachat/chat/streaming.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index 9ef29b2c6cb..b42c7baa403 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -80,6 +80,16 @@ class GigaChatModelResponseIterator: prompt_tokens=usage.prompt_tokens, completion_tokens=usage.completion_tokens, total_tokens=usage.total_tokens, + prompt_tokens_details=( + usage.prompt_tokens_details.model_dump() + if usage.prompt_tokens_details + else None + ), + completion_tokens_details=( + usage.completion_tokens_details.model_dump() + if usage.completion_tokens_details + else None + ), ) if finish_reason is not None: From e3058e34c4937acf93cd25e0c264281dfb9b5af0 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Fri, 12 Jun 2026 15:22:01 +0000 Subject: [PATCH 017/120] fix return headers for stream response --- litellm/passthrough/main.py | 302 +++++++++++------- .../test_async_streaming_error_propagation.py | 81 +++-- ...test_streaming_interrupt_spend_tracking.py | 79 +++-- 3 files changed, 274 insertions(+), 188 deletions(-) diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 8b20406397b..2c654a3f23b 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -4,11 +4,11 @@ This module is used to pass through requests to the LLM APIs. import asyncio import contextvars -from collections.abc import AsyncIterator from functools import partial from typing import ( TYPE_CHECKING, Any, + AsyncGenerator, Coroutine, Generator, List, @@ -36,59 +36,203 @@ if TYPE_CHECKING: from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig -class _AsyncPassthroughStreamingResponse(AsyncIterator[bytes]): - """ - Async iterator wrapper that preserves upstream response metadata for streaming. - """ +class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): + def __init__( + self, + response: Coroutine[Any, Any, httpx.Response], + litellm_logging_obj: "LiteLLMLoggingObj", + provider_config: "BasePassthroughConfig", + ) -> None: + self._initialized = False + self._status_code: int = 0 + self._headers = httpx.Headers() + self._response_coro = response + self._response: httpx.Response + self._iterator: AsyncGenerator[bytes, Any] + self._litellm_logging_obj = litellm_logging_obj + self._provider_config = provider_config + self._raw_bytes: List[bytes] = [] + self._flush_scheduled = False + self._background_tasks: set[asyncio.Task] = set() + @property + def status_code(self) -> int: + if not self._initialized: + raise RuntimeError( + "AsyncPassthroughStreamingResponse must be awaited " + "before accessing status_code" + ) + return self._status_code + + @status_code.setter + def status_code(self, value: int) -> None: + self._status_code = value + + @property + def headers(self) -> httpx.Headers: + if not self._initialized: + raise RuntimeError( + "AsyncPassthroughStreamingResponse must be awaited " + "before accessing headers" + ) + return self._headers + + @headers.setter + def headers(self, value: httpx.Headers) -> None: + self._headers = value + + def __await__(self): + async def _init(): + if not self._initialized: + self._response = await self._response_coro + self.headers = self._response.headers + self.status_code = self._response.status_code + self._initialized = True + try: + self._response.raise_for_status() + self._iterator = cast( + AsyncGenerator[bytes, Any], self._response.aiter_bytes() + ) + except Exception: + try: + await self._response.aclose() + except Exception: + pass + raise + return self + + return _init().__await__() + + def _start_flush(self) -> None: + if self._flush_scheduled or not self._raw_bytes: + return + self._flush_scheduled = True + + try: + task = asyncio.create_task( + self._litellm_logging_obj.async_flush_passthrough_collected_chunks( + raw_bytes=self._raw_bytes, + provider_config=self._provider_config, + ) + ) + + # Compliant: Save a strong reference to prevent GC + self._background_tasks.add(task) + + # Remove the task from the set when it finishes to avoid memory leaks + task.add_done_callback(self._background_tasks.discard) + except Exception as e: + verbose_logger.exception( + "Failed to schedule passthrough spend-tracking flush; " + "%d buffered chunks dropped: %s", + len(self._raw_bytes), + e, + ) + + def __aiter__(self) -> "AsyncPassthroughStreamingResponse": + return self + + async def __anext__(self) -> bytes: + if not self._initialized: + await self + try: + chunk = await self._iterator.__anext__() + self._raw_bytes.append(chunk) + return chunk + except Exception: + self._start_flush() + try: + await self._response.aclose() + except Exception: + pass + raise + + async def asend(self, value: Any) -> bytes: + if not self._initialized: + await self + return await self._iterator.asend(value) + + async def athrow(self, typ: Any, val: Any = None, tb: Any = None) -> bytes: + if not self._initialized: + await self + return await self._iterator.athrow(typ, val, tb) + + async def aclose(self) -> None: + self._start_flush() + try: + if self._initialized: + await self._response.aclose() + except Exception: + pass + + +class PassthroughStreamingResponse(Generator[Any, Any, Any]): def __init__( self, response: httpx.Response, litellm_logging_obj: "LiteLLMLoggingObj", provider_config: "BasePassthroughConfig", ) -> None: - self.response = response + self._response = response self.headers = response.headers self.status_code = response.status_code self._litellm_logging_obj = litellm_logging_obj self._provider_config = provider_config - self._iterator = response.aiter_bytes() + self._iterator: Generator[bytes, Any, Any] = cast( + Generator[bytes, Any, Any], response.iter_bytes() + ) self._raw_bytes: List[bytes] = [] - self._flush_started = False + self._flush_scheduled = False - def __aiter__(self) -> "_AsyncPassthroughStreamingResponse": + def _start_flush(self) -> None: + if self._flush_scheduled or not self._raw_bytes: + return + self._flush_scheduled = True + + from litellm.utils import executor + + try: + executor.submit( + self._litellm_logging_obj.flush_passthrough_collected_chunks, + raw_bytes=self._raw_bytes, + provider_config=self._provider_config, + ) + except Exception as e: + verbose_logger.exception( + "Failed to schedule passthrough spend-tracking flush; " + "%d buffered chunks dropped: %s", + len(self._raw_bytes), + e, + ) + + def __iter__(self) -> "PassthroughStreamingResponse": return self - async def __anext__(self) -> bytes: + def __next__(self) -> bytes: try: - chunk = await self._iterator.__anext__() + chunk = next(self._iterator) self._raw_bytes.append(chunk) return chunk - except StopAsyncIteration: - self._start_flush() - raise except Exception: + self._start_flush() try: - await self.response.aclose() + self._response.close() except Exception: pass raise - def _start_flush(self) -> None: - if self._flush_started: - return + def send(self, value: Any) -> bytes: + return self._iterator.send(value) - self._flush_started = True - asyncio.create_task( - self._litellm_logging_obj.async_flush_passthrough_collected_chunks( - raw_bytes=self._raw_bytes, - provider_config=self._provider_config, - ) - ) + def throw(self, typ: Any, val: Any = None, tb: Any = None) -> bytes: + return self._iterator.throw(typ, val, tb) - async def aclose(self) -> None: + def close(self) -> None: self._start_flush() - await self.response.aclose() + try: + self._response.close() + except Exception: + pass @client @@ -110,7 +254,7 @@ async def allm_passthrough_route( cookies: Optional[CookieTypes] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, **kwargs, -) -> Union[httpx.Response, AsyncIterator[bytes]]: +) -> Union[httpx.Response, AsyncGenerator[Any, Any]]: """ Async: Reranks a list of documents based on their relevance to the query """ @@ -239,9 +383,9 @@ def llm_passthrough_route( ) -> Union[ httpx.Response, Coroutine[Any, Any, httpx.Response], - Coroutine[Any, Any, Union[httpx.Response, AsyncIterator[bytes]]], + Coroutine[Any, Any, Union[httpx.Response, AsyncGenerator[Any, Any]]], Generator[Any, Any, Any], - AsyncIterator[bytes], + AsyncGenerator[Any, Any], ]: """ Pass through requests to the LLM APIs. @@ -390,12 +534,11 @@ def llm_passthrough_route( response: httpx.Response = client.client.send(request=request, stream=is_streaming_request) # type: ignore response.raise_for_status() - if ( - hasattr(response, "iter_bytes") and is_streaming_request - ): # yield the chunk, so we can store it in the logging object - return _sync_streaming(response, litellm_logging_obj, provider_config) + if hasattr(response, "iter_bytes") and is_streaming_request: + return PassthroughStreamingResponse( + response, litellm_logging_obj, provider_config + ) else: - # For non-streaming responses, yield the entire response return response except Exception as e: if provider_config is None: @@ -412,7 +555,7 @@ async def _async_passthrough_request( is_streaming_request: bool, litellm_logging_obj: "LiteLLMLoggingObj", provider_config: "BasePassthroughConfig", -) -> Union[httpx.Response, AsyncIterator[bytes]]: +) -> Union[httpx.Response, AsyncGenerator[Any, Any]]: """ Handle async passthrough requests. Uses async client to send request and properly handles streaming. @@ -423,10 +566,8 @@ async def _async_passthrough_request( # Check if it's a coroutine and await it if asyncio.iscoroutine(response_result): if is_streaming_request: - iter_response = await response_result - iter_response.raise_for_status() - return _AsyncPassthroughStreamingResponse( - response=iter_response, + return await AsyncPassthroughStreamingResponse( + response=response_result, litellm_logging_obj=litellm_logging_obj, provider_config=provider_config, ) @@ -438,84 +579,3 @@ async def _async_passthrough_request( else: # Fallback for sync-like behavior (shouldn't happen in async path) raise Exception("Expected coroutine from async client") - - -def _sync_streaming( - response: httpx.Response, - litellm_logging_obj: "LiteLLMLoggingObj", - provider_config: "BasePassthroughConfig", -): - from litellm.utils import executor - - raw_bytes: List[bytes] = [] - flush_scheduled = False - try: - for chunk in response.iter_bytes(): # type: ignore - raw_bytes.append(chunk) - yield chunk - finally: - if not flush_scheduled and raw_bytes: - flush_scheduled = True - try: - executor.submit( - litellm_logging_obj.flush_passthrough_collected_chunks, - raw_bytes=raw_bytes, - provider_config=provider_config, - ) - except Exception as e: - verbose_logger.exception( - "Failed to schedule passthrough spend-tracking flush " - "in _sync_streaming; %d buffered chunks dropped: %s", - len(raw_bytes), - e, - ) - - -async def _async_streaming( - response: Coroutine[Any, Any, httpx.Response], - litellm_logging_obj: "LiteLLMLoggingObj", - provider_config: "BasePassthroughConfig", -): - iter_response = await response - - try: - iter_response.raise_for_status() - except Exception: - try: - await iter_response.aclose() - except Exception: - pass - raise - - raw_bytes: List[bytes] = [] - flush_scheduled = False - try: - async for chunk in iter_response.aiter_bytes(): # type: ignore - raw_bytes.append(chunk) - yield chunk - except Exception: - try: - await iter_response.aclose() - except Exception: - pass - raise - finally: - # GeneratorExit (raised on client disconnect) is not caught by - # `except Exception`; the finally block ensures partial usage - # still gets flushed for spend tracking. See LIT-2642. - if not flush_scheduled and raw_bytes: - flush_scheduled = True - try: - asyncio.create_task( - litellm_logging_obj.async_flush_passthrough_collected_chunks( - raw_bytes=raw_bytes, - provider_config=provider_config, - ) - ) - except Exception as e: - verbose_logger.exception( - "Failed to schedule passthrough spend-tracking flush " - "in _async_streaming; %d buffered chunks dropped: %s", - len(raw_bytes), - e, - ) diff --git a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py b/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py index 631037b3357..f7052fee3b9 100644 --- a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py +++ b/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py @@ -52,43 +52,69 @@ def _make_mock_logging_obj(): @pytest.mark.asyncio -async def test_async_passthrough_wrapper_429_raises_before_iteration(): - """429 from upstream should be raised before the wrapper is constructed.""" +async def test_async_streaming_429_raises(): + """429 from upstream should raise HTTPStatusError, not yield error bytes.""" + from litellm.passthrough.main import AsyncPassthroughStreamingResponse + error_body = json.dumps( {"error": {"code": "429", "message": "Rate limit exceeded."}} ).encode() mock_response = _make_mock_response(429, error_body) - + + async def response_coro(): + return mock_response + + chunks = [] with pytest.raises(httpx.HTTPStatusError) as exc_info: - mock_response.raise_for_status() - + async for chunk in AsyncPassthroughStreamingResponse( + response=response_coro(), + litellm_logging_obj=_make_mock_logging_obj(), + provider_config=MagicMock(), + ): + chunks.append(chunk) + assert exc_info.value.response.status_code == 429 + assert len(chunks) == 0 @pytest.mark.asyncio -async def test_async_passthrough_wrapper_500_raises_before_iteration(): - """500 from upstream should be raised before the wrapper is constructed.""" +async def test_async_streaming_500_raises(): + """500 from upstream should also raise, not yield error bytes.""" + from litellm.passthrough.main import AsyncPassthroughStreamingResponse + error_body = json.dumps( {"error": {"code": "500", "message": "Internal server error"}} ).encode() mock_response = _make_mock_response(500, error_body) - + + async def response_coro(): + return mock_response + with pytest.raises(httpx.HTTPStatusError) as exc_info: - mock_response.raise_for_status() - + async for _ in AsyncPassthroughStreamingResponse( + response=response_coro(), + litellm_logging_obj=_make_mock_logging_obj(), + provider_config=MagicMock(), + ): + pass + assert exc_info.value.response.status_code == 500 @pytest.mark.asyncio async def test_async_passthrough_wrapper_200_yields_chunks(): """Successful 200 streaming responses should continue to work normally.""" - from litellm.passthrough.main import _AsyncPassthroughStreamingResponse + from litellm.passthrough.main import AsyncPassthroughStreamingResponse sse_data = b'data: {"type":"response.created"}\n\ndata: [DONE]\n\n' mock_response = _make_mock_response(200, sse_data) mock_logging_obj = _make_mock_logging_obj() - async_stream = _AsyncPassthroughStreamingResponse( - response=mock_response, + + async def response_coro(): + return mock_response + + async_stream = AsyncPassthroughStreamingResponse( + response=response_coro(), litellm_logging_obj=mock_logging_obj, provider_config=MagicMock(), ) @@ -102,32 +128,3 @@ async def test_async_passthrough_wrapper_200_yields_chunks(): assert len(chunks) == 1 assert b"response.created" in chunks[0] mock_logging_obj.async_flush_passthrough_collected_chunks.assert_awaited_once() - - -@pytest.mark.asyncio -async def test_async_passthrough_wrapper_closes_response_on_iteration_error(): - """Wrapper should close the upstream response if iteration raises.""" - from litellm.passthrough.main import _AsyncPassthroughStreamingResponse - - mock_response = MagicMock(spec=httpx.Response) - mock_response.status_code = 200 - mock_response.headers = httpx.Headers({"content-type": "text/event-stream"}) - mock_response.aclose = AsyncMock() - - async def _failing_aiter_bytes(): - raise RuntimeError("stream failed") - yield b"" - - mock_response.aiter_bytes = _failing_aiter_bytes - - async_stream = _AsyncPassthroughStreamingResponse( - response=mock_response, - litellm_logging_obj=_make_mock_logging_obj(), - provider_config=MagicMock(), - ) - - with pytest.raises(RuntimeError, match="stream failed"): - async for chunk in async_stream: - _ = chunk - - mock_response.aclose.assert_awaited_once() diff --git a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py b/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py index f3fe3ae5c38..27d265e6f5d 100644 --- a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py +++ b/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py @@ -35,11 +35,14 @@ class _ImmediateExecutor: @pytest.mark.asyncio -async def test_async_streaming_flushes_on_normal_completion(): - from litellm.passthrough.main import _async_streaming +async def test_asyncpassthroughstreamingresponse_flushes_on_normal_completion(): + from litellm.passthrough.main import AsyncPassthroughStreamingResponse chunks = [b"chunk-1", b"chunk-2", b"chunk-3"] mock_response = _make_streaming_response(chunks) + mock_response.headers = httpx.Headers( + {"content-type": "application/octet-stream", "x-request-id": "req-123"} + ) async def response_coro(): return mock_response @@ -48,14 +51,19 @@ async def test_async_streaming_flushes_on_normal_completion(): provider_config = MagicMock() received = [] - async for chunk in _async_streaming( + received_response = AsyncPassthroughStreamingResponse( response=response_coro(), litellm_logging_obj=mock_logging_obj, provider_config=provider_config, - ): + ) + + async for chunk in received_response: received.append(chunk) assert received == chunks + + assert received_response.headers["content-type"] == "application/octet-stream" + assert received_response.headers["x-request-id"] == "req-123" await asyncio.sleep(0) @@ -68,8 +76,8 @@ async def test_async_streaming_flushes_on_normal_completion(): @pytest.mark.asyncio -async def test_async_streaming_flushes_on_client_disconnect(): - from litellm.passthrough.main import _async_streaming +async def test_asyncpassthroughstreamingresponse_flushes_on_client_disconnect(): + from litellm.passthrough.main import AsyncPassthroughStreamingResponse chunks = [ b'{"chunk": 1, "outputTokens": 10}', @@ -77,6 +85,9 @@ async def test_async_streaming_flushes_on_client_disconnect(): b'{"chunk": 3, "outputTokens": 8}', ] mock_response = _make_streaming_response(chunks) + mock_response.headers = httpx.Headers( + {"content-type": "application/octet-stream", "x-request-id": "req-123"} + ) async def response_coro(): return mock_response @@ -84,7 +95,7 @@ async def test_async_streaming_flushes_on_client_disconnect(): mock_logging_obj = _make_logging_obj() provider_config = MagicMock() - gen = _async_streaming( + gen = AsyncPassthroughStreamingResponse( response=response_coro(), litellm_logging_obj=mock_logging_obj, provider_config=provider_config, @@ -105,11 +116,14 @@ async def test_async_streaming_flushes_on_client_disconnect(): @pytest.mark.asyncio -async def test_async_streaming_does_not_flush_on_4xx(): - from litellm.passthrough.main import _async_streaming +async def test_asyncpassthroughstreamingresponse_does_not_flush_on_4xx(): + from litellm.passthrough.main import AsyncPassthroughStreamingResponse err_response = MagicMock(spec=httpx.Response) err_response.status_code = 429 + err_response.headers = httpx.Headers( + {"content-type": "application/octet-stream"} + ) def _raise(): raise httpx.HTTPStatusError( @@ -129,7 +143,7 @@ async def test_async_streaming_does_not_flush_on_4xx(): mock_logging_obj = _make_logging_obj() with pytest.raises(httpx.HTTPStatusError): - async for _ in _async_streaming( + async for _ in AsyncPassthroughStreamingResponse( response=response_coro(), litellm_logging_obj=mock_logging_obj, provider_config=MagicMock(), @@ -140,8 +154,8 @@ async def test_async_streaming_does_not_flush_on_4xx(): @pytest.mark.asyncio -async def test_async_streaming_flushes_on_upstream_exception_with_partial_data(): - from litellm.passthrough.main import _async_streaming +async def test_asyncpassthroughstreamingresponse_flushes_on_upstream_exception_with_partial_data(): + from litellm.passthrough.main import AsyncPassthroughStreamingResponse partial_chunks = [b"partial-chunk-1", b"partial-chunk-2"] @@ -149,6 +163,9 @@ async def test_async_streaming_flushes_on_upstream_exception_with_partial_data() mock_response.status_code = 200 mock_response.raise_for_status = MagicMock(return_value=None) mock_response.aclose = AsyncMock() + mock_response.headers = httpx.Headers( + {"content-type": "application/octet-stream", "x-request-id": "req-123"} + ) async def _aiter_bytes_then_raise(): for c in partial_chunks: @@ -165,7 +182,7 @@ async def test_async_streaming_flushes_on_upstream_exception_with_partial_data() received = [] with pytest.raises(httpx.ReadError): - async for chunk in _async_streaming( + async for chunk in AsyncPassthroughStreamingResponse( response=response_coro(), litellm_logging_obj=mock_logging_obj, provider_config=provider_config, @@ -183,12 +200,16 @@ async def test_async_streaming_flushes_on_upstream_exception_with_partial_data() assert call_kwargs["raw_bytes"] == partial_chunks -def test_sync_streaming_flushes_on_normal_completion(): - from litellm.passthrough.main import _sync_streaming +def test_passthroughstreamingresponse_flushes_on_normal_completion(): + from litellm.passthrough.main import PassthroughStreamingResponse chunks = [b"a", b"b", b"c"] mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = httpx.Headers( + {"content-type": "application/octet-stream", "x-request-id": "req-123"} + ) def _iter_bytes(): yield from chunks @@ -199,25 +220,33 @@ def test_sync_streaming_flushes_on_normal_completion(): mock_logging_obj.flush_passthrough_collected_chunks = MagicMock() provider_config = MagicMock() + received_responce = PassthroughStreamingResponse( + response=mock_response, + litellm_logging_obj=mock_logging_obj, + provider_config=provider_config, + ) + with patch("litellm.utils.executor", _ImmediateExecutor()): - received = list( - _sync_streaming( - response=mock_response, - litellm_logging_obj=mock_logging_obj, - provider_config=provider_config, - ) - ) + received = list(received_responce) assert received == chunks + + assert received_responce.headers["content-type"] == "application/octet-stream" + assert received_responce.headers["x-request-id"] == "req-123" + mock_logging_obj.flush_passthrough_collected_chunks.assert_called_once() -def test_sync_streaming_flushes_on_early_close(): - from litellm.passthrough.main import _sync_streaming +def test_passthroughstreamingresponse_flushes_on_early_close(): + from litellm.passthrough.main import PassthroughStreamingResponse chunks = [b"first", b"second", b"third"] mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = httpx.Headers( + {"content-type": "application/octet-stream", "x-request-id": "req-123"} + ) def _iter_bytes(): yield from chunks @@ -229,7 +258,7 @@ def test_sync_streaming_flushes_on_early_close(): provider_config = MagicMock() with patch("litellm.utils.executor", _ImmediateExecutor()): - gen = _sync_streaming( + gen = PassthroughStreamingResponse( response=mock_response, litellm_logging_obj=mock_logging_obj, provider_config=provider_config, From 9307d1df74c2d76e9ee4b8bb238e9b97a091437f Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Tue, 16 Jun 2026 19:13:06 +0000 Subject: [PATCH 018/120] update OpenAPI schema --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 191 ++++++++++++++++++ 1 file changed, 191 insertions(+) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 4797e41a62b..e28dba160b3 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -4148,6 +4148,42 @@ export interface paths { patch?: never; trace?: never; }; + "/gigachat/{endpoint}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Gigachat Proxy Route + * @description [Docs](https://docs.litellm.ai/docs/pass_through/gigachat) + */ + get: operations["gigachat_proxy_route_gigachat__endpoint__get"]; + /** + * Gigachat Proxy Route + * @description [Docs](https://docs.litellm.ai/docs/pass_through/gigachat) + */ + put: operations["gigachat_proxy_route_gigachat__endpoint__put"]; + /** + * Gigachat Proxy Route + * @description [Docs](https://docs.litellm.ai/docs/pass_through/gigachat) + */ + post: operations["gigachat_proxy_route_gigachat__endpoint__post"]; + /** + * Gigachat Proxy Route + * @description [Docs](https://docs.litellm.ai/docs/pass_through/gigachat) + */ + delete: operations["gigachat_proxy_route_gigachat__endpoint__delete"]; + options?: never; + head?: never; + /** + * Gigachat Proxy Route + * @description [Docs](https://docs.litellm.ai/docs/pass_through/gigachat) + */ + patch: operations["gigachat_proxy_route_gigachat__endpoint__patch"]; + trace?: never; + }; "/global/activity": { parameters: { query?: never; @@ -39050,6 +39086,161 @@ export interface operations { }; }; }; + gigachat_proxy_route_gigachat__endpoint__get: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + gigachat_proxy_route_gigachat__endpoint__put: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + gigachat_proxy_route_gigachat__endpoint__post: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + gigachat_proxy_route_gigachat__endpoint__delete: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + gigachat_proxy_route_gigachat__endpoint__patch: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_global_activity_global_activity_get: { parameters: { query?: { From 93d272f40c73cb3925787c1d05baed987a4a9dd1 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Thu, 18 Jun 2026 14:35:46 +0000 Subject: [PATCH 019/120] fix some lint issues --- litellm/litellm_core_utils/litellm_logging.py | 2 +- litellm/llms/gigachat/authenticator.py | 27 ++-- litellm/llms/gigachat/chat/streaming.py | 8 +- litellm/llms/gigachat/chat/transformation.py | 58 ++++---- .../llms/gigachat/embedding/transformation.py | 26 ++-- litellm/llms/gigachat/file_handler.py | 25 ++-- .../gigachat/passthrough/transformation.py | 34 ++--- litellm/llms/gigachat/utils.py | 4 +- litellm/passthrough/main.py | 68 ++++----- litellm/proxy/common_request_processing.py | 2 +- .../llm_passthrough_endpoints.py | 139 ++++++++++-------- 11 files changed, 199 insertions(+), 194 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 347efbcbc97..e0102a60994 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5936,7 +5936,7 @@ def emit_standard_logging_payload(payload: StandardLoggingPayload): if os.getenv("LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD"): try: print(json.dumps(payload, indent=4, default=str)) # noqa: T201 - except Exception as e: + except Exception as e: # noqa: BLE001 verbose_logger.exception( "Error serializing standard logging payload for debug output: {}".format( str(e) diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index a20da943ee1..5418b3bcb8f 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -7,7 +7,6 @@ Based on official GigaChat SDK authentication flow. import time import uuid -from typing import Optional, Tuple import httpx @@ -41,7 +40,7 @@ class GigaChatAuthError(BaseLLMException): pass -def _get_credentials() -> Optional[str]: +def _get_credentials() -> str | None: """Get GigaChat credentials from environment.""" return get_secret_str("GIGACHAT_CREDENTIALS") or get_secret_str("GIGACHAT_API_KEY") @@ -62,10 +61,10 @@ def _get_http_client() -> HTTPHandler: def get_access_token( - credentials: Optional[str] = None, - scope: Optional[str] = None, - auth_url: Optional[str] = None, - litellm_params: Optional[dict] = None, + credentials: str | None = None, + scope: str | None = None, + auth_url: str | None = None, + litellm_params: dict | None = None, ) -> str: """ Get valid access token, using cache if available. @@ -124,10 +123,10 @@ def get_access_token( async def get_access_token_async( - credentials: Optional[str] = None, - scope: Optional[str] = None, - auth_url: Optional[str] = None, - litellm_params: Optional[dict] = None, + credentials: str | None = None, + scope: str | None = None, + auth_url: str | None = None, + litellm_params: dict | None = None, ) -> str: """Async version of get_access_token.""" if not litellm_params: @@ -176,12 +175,12 @@ def _request_token_sync( credentials: str, scope: str, auth_url: str, -) -> Tuple[str, int]: +) -> tuple[str, int]: """ Request new access token from GigaChat OAuth endpoint (sync). Returns: - Tuple of (access_token, expires_at_ms) + tuple of (access_token, expires_at_ms) """ headers = { "Authorization": f"Basic {credentials}", @@ -213,7 +212,7 @@ async def _request_token_async( credentials: str, scope: str, auth_url: str, -) -> Tuple[str, int]: +) -> tuple[str, int]: """Async version of _request_token_sync.""" headers = { "Authorization": f"Basic {credentials}", @@ -244,7 +243,7 @@ async def _request_token_async( ) -def _parse_token_response(response: httpx.Response) -> Tuple[str, int]: +def _parse_token_response(response: httpx.Response) -> tuple[str, int]: """Parse OAuth token response.""" data = response.json() diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index b42c7baa403..5c56a0c5d87 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -4,7 +4,7 @@ GigaChat Streaming Response Handler import json import uuid -from typing import Any, Optional +from typing import Any from litellm.llms.gigachat.utils import convert_usage from litellm.types.llms.openai import ( @@ -21,7 +21,7 @@ class GigaChatModelResponseIterator: self, streaming_response: Any, sync_stream: bool, - json_mode: Optional[bool] = False, + json_mode: bool | None = False, ): self.streaming_response = streaming_response self.response_iterator = self.streaming_response @@ -30,9 +30,9 @@ class GigaChatModelResponseIterator: def chunk_parser(self, chunk: dict) -> GenericStreamingChunk: """Parse a single streaming chunk from GigaChat.""" text = "" - tool_use: Optional[ChatCompletionToolCallChunk] = None + tool_use: ChatCompletionToolCallChunk | None = None is_finished = False - finish_reason: Optional[str] = None + finish_reason: str | None = None choices = chunk.get("choices", []) if not choices: diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index b3bb7bc5770..9ffb7cd8a16 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -7,7 +7,7 @@ Transforms OpenAI-format requests to GigaChat format and back. import json import time import uuid -from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, List, Optional, Union +from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, Union import httpx @@ -60,36 +60,36 @@ class GigaChatConfig(BaseConfig): stream: Enable streaming """ - temperature: Optional[float] = None - top_p: Optional[float] = None - max_tokens: Optional[int] = None - repetition_penalty: Optional[float] = None - profanity_check: Optional[bool] = None + temperature: float | None = None + top_p: float | None = None + max_tokens: int | None = None + repetition_penalty: float | None = None + profanity_check: bool | None = None def __init__( self, - temperature: Optional[float] = None, - top_p: Optional[float] = None, - max_tokens: Optional[int] = None, - repetition_penalty: Optional[float] = None, - profanity_check: Optional[bool] = None, + temperature: float | None = None, + top_p: float | None = None, + max_tokens: int | None = None, + repetition_penalty: float | None = None, + profanity_check: bool | None = None, ) -> None: locals_ = locals().copy() for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) # Instance variables for current request context - self._current_credentials: Optional[str] = None - self._current_api_base: Optional[str] = None + self._current_credentials: str | None = None + self._current_api_base: str | None = None def get_complete_url( self, - api_base: Optional[str], - api_key: Optional[str], + api_base: str | None, + api_key: str | None, model: str, optional_params: dict, litellm_params: dict, - stream: Optional[bool] = None, + stream: bool | None = None, ) -> str: """Get complete API URL for chat completions.""" base = get_api_base(api_base) @@ -99,11 +99,11 @@ class GigaChatConfig(BaseConfig): self, headers: dict, model: str, - messages: List[AllMessageValues], + messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - api_key: Optional[str] = None, - api_base: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, ) -> dict: """ Set up headers with OAuth token. @@ -128,7 +128,7 @@ class GigaChatConfig(BaseConfig): return headers - def get_supported_openai_params(self, model: str) -> List[str]: + def get_supported_openai_params(self, model: str) -> list[str]: """Return list of supported OpenAI parameters.""" return [ "stream", @@ -201,7 +201,7 @@ class GigaChatConfig(BaseConfig): return optional_params - def _convert_tools_to_functions(self, tools: List[dict]) -> List[dict]: + def _convert_tools_to_functions(self, tools: list[dict]) -> list[dict]: """Convert OpenAI tools format to GigaChat functions format.""" functions = [] for tool in tools: @@ -218,7 +218,7 @@ class GigaChatConfig(BaseConfig): def _map_tool_choice( self, tool_choice: Union[str, dict] - ) -> Optional[Union[str, dict]]: + ) -> Union[str, dict] | None: """ Map OpenAI tool_choice to GigaChat function_call format. @@ -258,7 +258,7 @@ class GigaChatConfig(BaseConfig): # Default to None (don't set function_call) return None - def _upload_image(self, image_url: str) -> Optional[str]: + def _upload_image(self, image_url: str) -> str | None: """ Upload image to GigaChat and return file_id. @@ -281,7 +281,7 @@ class GigaChatConfig(BaseConfig): def transform_request( self, model: str, - messages: List[AllMessageValues], + messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, headers: dict, @@ -316,7 +316,7 @@ class GigaChatConfig(BaseConfig): return request_data - def _transform_messages(self, messages: List[AllMessageValues]) -> List[dict]: + def _transform_messages(self, messages: list[AllMessageValues]) -> list[dict]: """Transform OpenAI messages to GigaChat format.""" transformed = [] @@ -395,12 +395,12 @@ class GigaChatConfig(BaseConfig): model_response: ModelResponse, logging_obj: LiteLLMLoggingObj, request_data: dict, - messages: List[AllMessageValues], + messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, encoding: Any, - api_key: Optional[str] = None, - json_mode: Optional[bool] = None, + api_key: str | None = None, + json_mode: bool | None = None, ) -> ModelResponse: """Transform GigaChat response to OpenAI format.""" try: @@ -494,7 +494,7 @@ class GigaChatConfig(BaseConfig): self, streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse], sync_stream: bool, - json_mode: Optional[bool] = False, + json_mode: bool | None = False, ): """Return streaming response iterator.""" from .streaming import GigaChatModelResponseIterator diff --git a/litellm/llms/gigachat/embedding/transformation.py b/litellm/llms/gigachat/embedding/transformation.py index 8dcccd49fcb..6e4698120ee 100644 --- a/litellm/llms/gigachat/embedding/transformation.py +++ b/litellm/llms/gigachat/embedding/transformation.py @@ -6,7 +6,7 @@ API Documentation: https://developers.sber.ru/docs/ru/gigachat/api/reference/res """ import types -from typing import List, Optional, Tuple, Union +from typing import Union import httpx @@ -55,7 +55,7 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): and v is not None } - def get_supported_openai_params(self, model: str) -> List[str]: + def get_supported_openai_params(self, model: str) -> list[str]: """GigaChat embeddings don't support additional parameters.""" return [] @@ -71,26 +71,26 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): def _get_openai_compatible_provider_info( self, - api_base: Optional[str], - api_key: Optional[str], - ) -> Tuple[str, Optional[str], Optional[str]]: + api_base: str | None, + api_key: str | None, + ) -> tuple[str, str | None, str | None]: """ Returns provider info for GigaChat. Returns: - Tuple of (custom_llm_provider, api_base, dynamic_api_key) + tuple of (custom_llm_provider, api_base, dynamic_api_key) """ api_base = get_api_base(api_base) return LlmProviders.GIGACHAT.value, api_base, api_key def get_complete_url( self, - api_base: Optional[str], - api_key: Optional[str], + api_base: str | None, + api_key: str | None, model: str, optional_params: dict, litellm_params: dict, - stream: Optional[bool] = None, + stream: bool | None = None, ) -> str: """Get the complete URL for embeddings endpoint.""" base = get_api_base(api_base) @@ -135,7 +135,7 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): raw_response: httpx.Response, model_response: EmbeddingResponse, logging_obj: LiteLLMLoggingObj, - api_key: Optional[str], + api_key: str | None, request_data: dict, optional_params: dict, litellm_params: dict, @@ -182,11 +182,11 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): self, headers: dict, model: str, - messages: List[AllMessageValues], + messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - api_key: Optional[str] = None, - api_base: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, ) -> dict: """ Set up headers with OAuth token for GigaChat. diff --git a/litellm/llms/gigachat/file_handler.py b/litellm/llms/gigachat/file_handler.py index 2054fe1f566..8c7075f5599 100644 --- a/litellm/llms/gigachat/file_handler.py +++ b/litellm/llms/gigachat/file_handler.py @@ -9,7 +9,6 @@ import base64 import hashlib import re import uuid -from typing import Dict, Optional, Tuple from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( @@ -22,7 +21,7 @@ from litellm.types.utils import LlmProviders from .authenticator import get_access_token, get_access_token_async # Simple in-memory cache for file IDs -_file_cache: Dict[str, str] = {} +_file_cache: dict[str, str] = {} def _get_url_hash(url: str) -> str: @@ -30,7 +29,7 @@ def _get_url_hash(url: str) -> str: return hashlib.sha256(url.encode()).hexdigest() -def _parse_data_url(data_url: str) -> Optional[Tuple[bytes, str, str]]: +def _parse_data_url(data_url: str) -> tuple[bytes, str, str] | None: """ Parse data URL (base64 image). @@ -49,7 +48,7 @@ def _parse_data_url(data_url: str) -> Optional[Tuple[bytes, str, str]]: return content_bytes, content_type, ext -def _download_image_sync(url: str) -> Tuple[bytes, str, str]: +def _download_image_sync(url: str) -> tuple[bytes, str, str]: """Download image from URL synchronously.""" client = _get_httpx_client(params={"ssl_verify": False}) response = client.get(url) @@ -61,7 +60,7 @@ def _download_image_sync(url: str) -> Tuple[bytes, str, str]: return response.content, content_type, ext -async def _download_image_async(url: str) -> Tuple[bytes, str, str]: +async def _download_image_async(url: str) -> tuple[bytes, str, str]: """Download image from URL asynchronously.""" client = get_async_httpx_client( llm_provider=LlmProviders.GIGACHAT, @@ -78,10 +77,10 @@ async def _download_image_async(url: str) -> Tuple[bytes, str, str]: def upload_file_sync( image_url: str, - credentials: Optional[str] = None, - api_base: Optional[str] = None, - litellm_params: Optional[dict] = None, -) -> Optional[str]: + credentials: str | None = None, + api_base: str | None = None, + litellm_params: dict | None = None, +) -> str | None: """ Upload file to GigaChat and return file_id (sync). @@ -146,10 +145,10 @@ def upload_file_sync( async def upload_file_async( image_url: str, - credentials: Optional[str] = None, - api_base: Optional[str] = None, - litellm_params: Optional[dict] = None, -) -> Optional[str]: + credentials: str | None = None, + api_base: str | None = None, + litellm_params: dict | None = None, +) -> str | None: """ Upload file to GigaChat and return file_id (async). diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index 474622a01a8..d2ba43b0f47 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -1,5 +1,5 @@ import json -from typing import TYPE_CHECKING, List, Optional, Tuple, cast +from typing import TYPE_CHECKING, cast import httpx @@ -24,13 +24,13 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): def get_complete_url( self, - api_base: Optional[str], - api_key: Optional[str], + api_base: str | None, + api_key: str | None, model: str, endpoint: str, - request_query_params: Optional[dict], + request_query_params: dict | None, litellm_params: dict, - ) -> Tuple["URL", str]: + ) -> tuple["URL", str]: """Get complete API URL for chat completions.""" base_target_url = self.get_api_base(api_base) @@ -48,11 +48,11 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): self, headers: dict, model: str, - messages: List[AllMessageValues], + messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - api_key: Optional[str] = None, - api_base: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, ) -> dict: """ Set up headers with OAuth token. @@ -76,7 +76,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): request_data: dict, logging_obj: "LiteLLMLoggingObj", endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> "CostResponseTypes" | None: from litellm import encoding from litellm.types.utils import LlmProviders, ModelResponse from litellm.utils import ProviderConfigManager @@ -140,12 +140,12 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): def handle_logging_collected_chunks( self, - all_chunks: List[str], + all_chunks: list[str], litellm_logging_obj: "LiteLLMLoggingObj", model: str, custom_llm_provider: str, endpoint: str, - ) -> Optional["CostResponseTypes"]: + ) -> "CostResponseTypes" | None: """ 1. Convert all_chunks to a ModelResponseStream 2. combine model_response_stream to model_response @@ -208,20 +208,20 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): return None @staticmethod - def get_api_base(api_base: Optional[str] = None) -> Optional[str]: + def get_api_base(api_base: str | None = None) -> str | None: return api_base or get_secret_str("GIGACHAT_API_BASE") or GIGACHAT_BASE_URL @staticmethod def get_api_key( - api_key: Optional[str] = None, - ) -> Optional[str]: + api_key: str | None = None, + ) -> str | None: return api_key or get_secret_str("GIGACHAT_API_KEY") @staticmethod - def get_base_model(model: str) -> Optional[str]: + def get_base_model(model: str) -> str | None: return model def get_models( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> List[str]: + self, api_key: str | None = None, api_base: str | None = None + ) -> list[str]: return super().get_models(api_key, api_base) diff --git a/litellm/llms/gigachat/utils.py b/litellm/llms/gigachat/utils.py index b66d25cbc0b..0895d1a4992 100644 --- a/litellm/llms/gigachat/utils.py +++ b/litellm/llms/gigachat/utils.py @@ -1,5 +1,3 @@ -from typing import Optional - from litellm.secret_managers.main import get_secret_str from litellm.types.utils import PromptTokensDetailsWrapper, Usage @@ -30,5 +28,5 @@ def convert_usage(usage_data: dict[str, int]) -> Usage: ) -def get_api_base(api_base: Optional[str] = None) -> Optional[str]: +def get_api_base(api_base: str | None = None) -> str | None: return api_base or get_secret_str("GIGACHAT_API_BASE") or GIGACHAT_BASE_URL diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 5c3e7aa1884..c3ebc81541d 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -50,7 +50,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): self._iterator: AsyncGenerator[bytes, Any] self._litellm_logging_obj = litellm_logging_obj self._provider_config = provider_config - self._raw_bytes: List[bytes] = [] + self._raw_bytes: list[bytes] = [] self._flush_scheduled = False self._background_tasks: set[asyncio.Task] = set() @@ -92,10 +92,10 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): self._iterator = cast( AsyncGenerator[bytes, Any], self._response.aiter_bytes() ) - except Exception: + except Exception: # noqa: BLE001 try: await self._response.aclose() - except Exception: + except Exception: # noqa: BLE001 pass raise return self @@ -120,7 +120,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): # Remove the task from the set when it finishes to avoid memory leaks task.add_done_callback(self._background_tasks.discard) - except Exception as e: + except Exception as e: # noqa: BLE001 verbose_logger.exception( "Failed to schedule passthrough spend-tracking flush; " "%d buffered chunks dropped: %s", @@ -138,11 +138,11 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): chunk = await self._iterator.__anext__() self._raw_bytes.append(chunk) return chunk - except Exception: + except Exception: # noqa: BLE001 self._start_flush() try: await self._response.aclose() - except Exception: + except Exception: # noqa: BLE001 pass raise @@ -161,7 +161,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): try: if self._initialized: await self._response.aclose() - except Exception: + except Exception: # noqa: BLE001 pass @@ -196,7 +196,7 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): raw_bytes=self._raw_bytes, provider_config=self._provider_config, ) - except Exception as e: + except Exception as e: # noqa: BLE001 verbose_logger.exception( "Failed to schedule passthrough spend-tracking flush; " "%d buffered chunks dropped: %s", @@ -212,11 +212,11 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): chunk = next(self._iterator) self._raw_bytes.append(chunk) return chunk - except Exception: + except Exception: # noqa: BLE001 self._start_flush() try: self._response.close() - except Exception: + except Exception: # noqa: BLE001 pass raise @@ -230,7 +230,7 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): self._start_flush() try: self._response.close() - except Exception: + except Exception: # noqa: BLE001 pass @@ -240,17 +240,17 @@ async def allm_passthrough_route( method: str, endpoint: str, model: str, - custom_llm_provider: Optional[str] = None, - api_base: Optional[str] = None, - api_key: Optional[str] = None, - request_query_params: Optional[dict] = None, - request_headers: Optional[dict] = None, - content: Optional[Any] = None, - data: Optional[dict] = None, - files: Optional[RequestFiles] = None, - json: Optional[Any] = None, - params: Optional[QueryParamTypes] = None, - cookies: Optional[CookieTypes] = None, + custom_llm_provider: str | None = None, + api_base: str | None = None, + api_key: str | None = None, + request_query_params: dict | None = None, + request_headers: dict | None = None, + content: Any | None = None, + data: dict | None = None, + files: RequestFiles | None = None, + json: Any | None = None, + params: QueryParamTypes | None = None, + cookies: CookieTypes | None = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, **kwargs, ) -> Union[httpx.Response, AsyncGenerator[Any, Any]]: @@ -365,19 +365,19 @@ def llm_passthrough_route( method: str, endpoint: str, model: str, - custom_llm_provider: Optional[str] = None, - api_base: Optional[str] = None, - api_key: Optional[str] = None, - request_query_params: Optional[dict] = None, - request_headers: Optional[dict] = None, + custom_llm_provider: str | None = None, + api_base: str | None = None, + api_key: str | None = None, + request_query_params: dict | None = None, + request_headers: dict | None = None, allm_passthrough_route: bool = False, - content: Optional[Any] = None, - data: Optional[dict] = None, - files: Optional[RequestFiles] = None, - json: Optional[Any] = None, - params: Optional[QueryParamTypes] = None, - cookies: Optional[CookieTypes] = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + content: Any | None = None, + data: dict | None = None, + files: RequestFiles | None = None, + json: Any | None = None, + params: QueryParamTypes | None = None, + cookies: CookieTypes | None = None, + client: Union[HTTPHandler, AsyncHTTPHandler] | None = None, **kwargs, ) -> Union[ httpx.Response, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index bfeaf84930d..6cb65692099 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -732,7 +732,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def _merge_passthrough_streaming_headers( - response_headers: Optional[Any], + response_headers: Any | None, custom_headers: dict, ) -> dict: """ diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 0f2bc79e9fb..3d3bab14020 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -9,7 +9,7 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc. import json import os import re -from typing import Any, Optional, Tuple, Union, cast +from typing import TYPE_CHECKING, Any, Callable, Union, cast import httpx from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket @@ -59,6 +59,11 @@ from litellm.utils import ProviderConfigManager from .passthrough_endpoint_router import PassthroughEndpointRouter +if TYPE_CHECKING: + from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.utils import ProxyLogging + + vertex_llm_base = VertexBase() router = APIRouter() default_vertex_config = None @@ -77,7 +82,7 @@ def create_request_copy(request: Request): def is_passthrough_request_using_router_model( - request_body: dict, llm_router: Optional[litellm.Router] + request_body: dict, llm_router: litellm.Router | None ) -> bool: """ Returns True if the model is in the llm_router model names @@ -225,7 +230,7 @@ async def gemini_proxy_route( ) # Add or update query parameters - gemini_api_key: Optional[str] = passthrough_endpoint_router.get_credentials( + gemini_api_key: str | None = passthrough_endpoint_router.get_credentials( custom_llm_provider="gemini", region_name=None, ) @@ -473,9 +478,9 @@ async def milvus_proxy_route( request_body = await get_request_body(request) # check collectionName - collection_name = cast(Optional[str], request_body.get("collectionName")) + collection_name = cast(str | None, request_body.get("collectionName")) extra_headers = {} - base_target_url: Optional[str] = None + base_target_url: str | None = None if not collection_name: raise HTTPException( status_code=400, @@ -760,12 +765,12 @@ async def handle_bedrock_passthrough_router_model( general_settings: dict, proxy_config, select_data_generator, - user_model: Optional[str], - user_temperature: Optional[float], - user_request_timeout: Optional[float], - user_max_tokens: Optional[int], - user_api_base: Optional[str], - version: Optional[str], + user_model: str | None, + user_temperature: float | None, + user_request_timeout: float | None, + user_max_tokens: int | None, + user_api_base: str | None, + version: str | None, ) -> Union[Response, StreamingResponse]: """ Handle Bedrock passthrough for router models (models defined in config.yaml). @@ -1134,12 +1139,12 @@ async def bedrock_proxy_route( def _resolve_vertex_model_from_router( model_id: str, - llm_router: Optional[litellm.Router], + llm_router: litellm.Router | None, encoded_endpoint: str, endpoint: str, - vertex_project: Optional[str], - vertex_location: Optional[str], -) -> Tuple[str, str, Optional[str], Optional[str]]: + vertex_project: str | None, + vertex_location: str | None, +) -> tuple[str, str, str | None, str | None]: """ Resolve Vertex AI model configuration from router. @@ -1152,7 +1157,7 @@ def _resolve_vertex_model_from_router( vertex_location: Current vertex location (may be from URL) Returns: - Tuple of (encoded_endpoint, endpoint, vertex_project, vertex_location) + tuple of (encoded_endpoint, endpoint, vertex_project, vertex_location) with resolved values from router config """ if not llm_router: @@ -1501,42 +1506,42 @@ from abc import ABC, abstractmethod class BaseVertexAIPassThroughHandler(ABC): @staticmethod @abstractmethod - def get_default_base_target_url(vertex_location: Optional[str]) -> str: + def get_default_base_target_url(vertex_location: str | None) -> str: pass @staticmethod @abstractmethod def update_base_target_url_with_credential_location( - base_target_url: str, vertex_location: Optional[str] + base_target_url: str, vertex_location: str | None ) -> str: pass class VertexAIDiscoveryPassThroughHandler(BaseVertexAIPassThroughHandler): @staticmethod - def get_default_base_target_url(vertex_location: Optional[str]) -> str: + def get_default_base_target_url(vertex_location: str | None) -> str: return "https://discoveryengine.googleapis.com/" @staticmethod def update_base_target_url_with_credential_location( - base_target_url: str, vertex_location: Optional[str] + base_target_url: str, vertex_location: str | None ) -> str: return base_target_url class VertexAIPassThroughHandler(BaseVertexAIPassThroughHandler): @staticmethod - def get_default_base_target_url(vertex_location: Optional[str]) -> str: + def get_default_base_target_url(vertex_location: str | None) -> str: return get_vertex_base_url(vertex_location) @staticmethod def update_base_target_url_with_credential_location( - base_target_url: str, vertex_location: Optional[str] + base_target_url: str, vertex_location: str | None ) -> str: return get_vertex_base_url(vertex_location) -def get_vertex_base_url(vertex_location: Optional[str]) -> str: +def get_vertex_base_url(vertex_location: str | None) -> str: """ Base URL for Vertex AI pass-through (trailing slash for URL joining). @@ -1586,10 +1591,10 @@ def get_vertex_pass_through_handler( def _override_vertex_params_from_router_credentials( - router_credentials: Optional[Any], - vertex_project: Optional[str], - vertex_location: Optional[str], -) -> Tuple[Optional[str], Optional[str]]: + router_credentials: Any | None, + vertex_project: str | None, + vertex_location: str | None, +) -> tuple[str | None, str | None]: """ Override vertex_project and vertex_location with values from router_credentials if available. @@ -1599,7 +1604,7 @@ def _override_vertex_params_from_router_credentials( vertex_location: Current vertex location (from URL) Returns: - Tuple of (vertex_project, vertex_location) with overridden values if applicable + tuple of (vertex_project, vertex_location) with overridden values if applicable """ if router_credentials is None: return vertex_project, vertex_location @@ -1648,13 +1653,13 @@ def _override_vertex_params_from_router_credentials( async def _prepare_vertex_auth_headers( request: Request, - vertex_credentials: Optional[Any], - router_credentials: Optional[Any], - vertex_project: Optional[str], - vertex_location: Optional[str], - base_target_url: Optional[str], + vertex_credentials: Any | None, + router_credentials: Any | None, + vertex_project: str | None, + vertex_location: str | None, + base_target_url: str | None, get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler, -) -> Tuple[dict, Optional[str], bool, Optional[str], Optional[str]]: +) -> tuple[dict, str | None, bool, str | None, str | None]: """ Prepare authentication headers for Vertex AI pass-through requests. @@ -1668,12 +1673,12 @@ async def _prepare_vertex_auth_headers( get_vertex_pass_through_handler: Handler for the specific Vertex AI service Returns: - Tuple containing: + tuple containing: - headers: dict - Authentication headers to use - - base_target_url: Optional[str] - Updated base target URL + - base_target_url: str | None - Updated base target URL - headers_passed_through: bool - Whether headers were passed through from request - - vertex_project: Optional[str] - Updated vertex project ID - - vertex_location: Optional[str] - Updated vertex location + - vertex_project: str | None - Updated vertex project ID + - vertex_location: str | None - Updated vertex location """ vertex_llm_base = VertexBase() headers_passed_through = False @@ -1746,8 +1751,8 @@ async def _base_vertex_proxy_route( request: Request, fastapi_response: Response, get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler, - user_api_key_dict: Optional[UserAPIKeyAuth] = None, - router_credentials: Optional[Any] = None, + user_api_key_dict: UserAPIKeyAuth | None = None, + router_credentials: Any | None = None, ): """ Base function for Vertex AI passthrough routes. @@ -1793,8 +1798,8 @@ async def _base_vertex_proxy_route( user_api_key_dict=user_api_key_dict, ) - vertex_project: Optional[str] = get_vertex_project_id_from_url(endpoint) - vertex_location: Optional[str] = get_vertex_location_from_url(endpoint) + vertex_project: str | None = get_vertex_project_id_from_url(endpoint) + vertex_location: str | None = get_vertex_location_from_url(endpoint) # Override with vector store credentials if available vertex_project, vertex_location = _override_vertex_params_from_router_credentials( @@ -1919,7 +1924,7 @@ async def vertex_discovery_proxy_route( from litellm.types.vector_stores import LiteLLM_ManagedVectorStore # Extract vector store ID from endpoint if present (e.g., dataStores/test-litellm-app_1761094730750) - vector_store_credentials: Optional[LiteLLM_ManagedVectorStore] = None + vector_store_credentials: LiteLLM_ManagedVectorStore | None = None vector_store_id_match = re.search(r"dataStores/([^/]+)", endpoint) if vector_store_id_match: @@ -2057,9 +2062,9 @@ class BaseOpenAIPassThroughHandler: fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth, base_target_url: str, - api_key: Optional[str], + api_key: str | None, custom_llm_provider: litellm.LlmProviders, - extra_headers: Optional[dict] = None, + extra_headers: dict | None = None, ): encoded_endpoint = httpx.URL(endpoint).path # Ensure endpoint starts with '/' for proper URL construction @@ -2115,7 +2120,7 @@ class BaseOpenAIPassThroughHandler: @staticmethod def _assemble_headers( - api_key: Optional[str], request: Request, extra_headers: Optional[dict] = None + api_key: str | None, request: Request, extra_headers: dict | None = None ) -> dict: base_headers = {} if api_key is not None: @@ -2251,10 +2256,10 @@ async def cursor_proxy_route( async def vertex_ai_live_websocket_passthrough( websocket: WebSocket, - model: Optional[str] = None, - vertex_project: Optional[str] = None, - vertex_location: Optional[str] = None, - user_api_key_dict: Optional[UserAPIKeyAuth] = None, + model: str | None = None, + vertex_project: str | None = None, + vertex_location: str | None = None, + user_api_key_dict: UserAPIKeyAuth | None = None, ): """ Vertex AI Live API WebSocket Pass-through Function @@ -2286,8 +2291,8 @@ async def vertex_ai_live_websocket_passthrough( ) resolved_project = vertex_project - resolved_location: Optional[str] = vertex_location - credentials_value: Optional[str] = None + resolved_location: str | None = vertex_location + credentials_value: str | None = None if vertex_credentials_config is not None: resolved_project = resolved_project or vertex_credentials_config.vertex_project @@ -2391,9 +2396,9 @@ def create_vertex_ai_live_websocket_endpoint(): def create_generic_websocket_passthrough_endpoint( provider: str, target_url: str, - custom_headers: Optional[dict] = None, + custom_headers: dict | None = None, forward_headers: bool = False, - cost_per_request: Optional[float] = None, + cost_per_request: float | None = None, ): """ Create a generic WebSocket passthrough endpoint for any provider. @@ -2540,7 +2545,7 @@ async def gigachat_proxy_route( ) return result - except Exception as e: + except Exception as e: # noqa: BLE001 raise await base_llm_response_processor._handle_llm_api_exception( e=e, user_api_key_dict=user_api_key_dict, @@ -2556,16 +2561,16 @@ async def handle_gigachat_passthrough_router_model( fastapi_response: Response, llm_router: litellm.Router, user_api_key_dict: UserAPIKeyAuth, - proxy_logging_obj, + proxy_logging_obj: ProxyLogging, general_settings: dict, - proxy_config, - select_data_generator, - user_model: Optional[str], - user_temperature: Optional[float], - user_request_timeout: Optional[float], - user_max_tokens: Optional[int], - user_api_base: Optional[str], - version: Optional[str], + proxy_config: ProxyConfig, + select_data_generator: Callable, + user_model: str | None, + user_temperature: float | None, + user_request_timeout: float | None, + user_max_tokens: int | None, + user_api_base: str | None, + version: str | None, ) -> Union[Response, StreamingResponse]: """ Handle Gigachat passthrough for router models (models defined in config.yaml). @@ -2580,6 +2585,10 @@ async def handle_gigachat_passthrough_router_model( request_body: The parsed request body llm_router: The LiteLLM router instance user_api_key_dict: The user API key authentication dictionary + proxy_logging_obj: Proxy logging + general_settings: Proxy general settings + proxy_config: Proxy config + select_data_generator: Select data generator function (additional args for common processing) Returns: @@ -2677,7 +2686,7 @@ async def handle_gigachat_passthrough_router_model( return result return result - except Exception as e: + except Exception as e: # noqa: BLE001 # Use common exception handling raise await base_llm_response_processor._handle_llm_api_exception( e=e, From 4791f003990aea2c52e0e666307e0cf2a9b8312f Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Sat, 20 Jun 2026 16:37:55 +0000 Subject: [PATCH 020/120] fix import ProxyConfig --- .../pass_through_endpoints/llm_passthrough_endpoints.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 3d3bab14020..01449d8ab26 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -47,7 +47,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, ) -from litellm.proxy.utils import is_known_model +from litellm.proxy.utils import is_known_model, ProxyLogging from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store, get_litellm_managed_vector_store, @@ -60,8 +60,11 @@ from litellm.utils import ProviderConfigManager from .passthrough_endpoint_router import PassthroughEndpointRouter if TYPE_CHECKING: - from litellm.proxy.proxy_server import ProxyConfig - from litellm.proxy.utils import ProxyLogging + from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig + + ProxyConfig = _ProxyConfig +else: + ProxyConfig = Any vertex_llm_base = VertexBase() From 0cc9147f4992e9389bdf1f027aa52be05a25b504 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Sat, 20 Jun 2026 22:04:59 +0000 Subject: [PATCH 021/120] Add gigachat tests --- .../gigachat/passthrough/transformation.py | 14 +- .../llms/gigachat/test_authenticator.py | 428 ++++++++++++++++ tests/litellm/llms/gigachat/test_utils.py | 84 +++ .../llms/gigachat/passthrough/__init__.py | 0 ...est_gigachat_passthrough_transformation.py | 481 ++++++++++++++++++ 5 files changed, 1001 insertions(+), 6 deletions(-) create mode 100644 tests/litellm/llms/gigachat/test_authenticator.py create mode 100644 tests/litellm/llms/gigachat/test_utils.py create mode 100644 tests/test_litellm/llms/gigachat/passthrough/__init__.py create mode 100644 tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index d2ba43b0f47..54277081398 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import json from typing import TYPE_CHECKING, cast @@ -30,7 +32,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): endpoint: str, request_query_params: dict | None, litellm_params: dict, - ) -> tuple["URL", str]: + ) -> tuple[URL, str]: """Get complete API URL for chat completions.""" base_target_url = self.get_api_base(api_base) @@ -72,11 +74,11 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): self, model: str, custom_llm_provider: str, - httpx_response: "Response", + httpx_response: Response, request_data: dict, - logging_obj: "LiteLLMLoggingObj", + logging_obj: LiteLLMLoggingObj, endpoint: str, - ) -> "CostResponseTypes" | None: + ) -> CostResponseTypes | None: from litellm import encoding from litellm.types.utils import LlmProviders, ModelResponse from litellm.utils import ProviderConfigManager @@ -141,11 +143,11 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): def handle_logging_collected_chunks( self, all_chunks: list[str], - litellm_logging_obj: "LiteLLMLoggingObj", + litellm_logging_obj: LiteLLMLoggingObj, model: str, custom_llm_provider: str, endpoint: str, - ) -> "CostResponseTypes" | None: + ) -> CostResponseTypes | None: """ 1. Convert all_chunks to a ModelResponseStream 2. combine model_response_stream to model_response diff --git a/tests/litellm/llms/gigachat/test_authenticator.py b/tests/litellm/llms/gigachat/test_authenticator.py new file mode 100644 index 00000000000..ed52ae986f5 --- /dev/null +++ b/tests/litellm/llms/gigachat/test_authenticator.py @@ -0,0 +1,428 @@ +""" +Tests for litellm.llms.gigachat.authenticator +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../../../")) + +from litellm.llms.gigachat.authenticator import ( + GIGACHAT_AUTH_URL, + GIGACHAT_SCOPE, + GigaChatAuthError, + _get_auth_url, + _get_credentials, + _get_scope, + _parse_token_response, + _request_token_async, + _request_token_sync, + get_access_token, + get_access_token_async, +) + + +class TestParseTokenResponse: + def test_parse_with_tok_and_exp(self): + response = MagicMock() + response.json.return_value = {"tok": "token123", "exp": 1234567890000} + token, expires_at = _parse_token_response(response) + assert token == "token123" + assert expires_at == 1234567890000 + + def test_parse_with_access_token_and_expires_at(self): + response = MagicMock() + response.json.return_value = { + "access_token": "token456", + "expires_at": 9876543210000, + } + token, expires_at = _parse_token_response(response) + assert token == "token456" + assert expires_at == 9876543210000 + + def test_parse_with_string_expires_at(self): + response = MagicMock() + response.json.return_value = { + "access_token": "token789", + "expires_at": "1234567890000", + } + token, expires_at = _parse_token_response(response) + assert token == "token789" + assert expires_at == 1234567890000 + + def test_parse_prefers_tok_over_access_token(self): + response = MagicMock() + response.json.return_value = { + "tok": "preferred", + "access_token": "fallback", + "exp": 111111, + } + token, expires_at = _parse_token_response(response) + assert token == "preferred" + assert expires_at == 111111 + + def test_parse_missing_token_raises(self): + response = MagicMock() + response.json.return_value = {"expires_at": 1234567890000} + with pytest.raises(GigaChatAuthError) as exc_info: + _parse_token_response(response) + assert "Invalid token response" in str(exc_info.value) + assert exc_info.value.status_code == 500 + + +class TestGetCredentials: + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + def test_get_credentials_from_gigachat_credentials(self, mock_get_secret): + mock_get_secret.side_effect = lambda key: "cred123" if key == "GIGACHAT_CREDENTIALS" else None + result = _get_credentials() + assert result == "cred123" + + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + def test_get_credentials_fallback_to_api_key(self, mock_get_secret): + mock_get_secret.side_effect = lambda key: ( + "apikey456" if key == "GIGACHAT_API_KEY" else None + ) + result = _get_credentials() + assert result == "apikey456" + + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + def test_get_credentials_returns_none(self, mock_get_secret): + mock_get_secret.return_value = None + result = _get_credentials() + assert result is None + + +class TestGetAuthUrl: + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + def test_get_auth_url_from_env(self, mock_get_secret): + mock_get_secret.return_value = "https://custom.auth.url" + result = _get_auth_url() + assert result == "https://custom.auth.url" + + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + def test_get_auth_url_default(self, mock_get_secret): + mock_get_secret.return_value = None + result = _get_auth_url() + assert result == GIGACHAT_AUTH_URL + + +class TestGetScope: + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + def test_get_scope_from_env(self, mock_get_secret): + mock_get_secret.return_value = "CUSTOM_SCOPE" + result = _get_scope() + assert result == "CUSTOM_SCOPE" + + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + def test_get_scope_default(self, mock_get_secret): + mock_get_secret.return_value = None + result = _get_scope() + assert result == GIGACHAT_SCOPE + + +class TestRequestTokenSync: + @patch("litellm.llms.gigachat.authenticator.uuid.uuid4") + @patch("litellm.llms.gigachat.authenticator._get_http_client") + def test_request_token_success(self, mock_get_client, mock_uuid): + mock_uuid.return_value = "test-uuid-123" + mock_response = MagicMock() + mock_response.json.return_value = {"tok": "newtoken", "exp": 9999999999999} + mock_response.raise_for_status = MagicMock() + mock_client = MagicMock() + mock_client.post.return_value = mock_response + mock_get_client.return_value = mock_client + + token, expires_at = _request_token_sync("creds", "SCOPE", "https://auth.url") + + assert token == "newtoken" + assert expires_at == 9999999999999 + mock_client.post.assert_called_once_with( + "https://auth.url", + headers={ + "Authorization": "Basic creds", + "RqUID": "test-uuid-123", + "Content-Type": "application/x-www-form-urlencoded", + }, + data={"scope": "SCOPE"}, + timeout=30, + ) + + @patch("litellm.llms.gigachat.authenticator._get_http_client") + def test_request_token_http_status_error(self, mock_get_client): + mock_response = MagicMock() + mock_response.text = "Unauthorized" + mock_response.status_code = 401 + mock_client = MagicMock() + mock_client.post.return_value = mock_response + mock_get_client.return_value = mock_client + + mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( + "401 Unauthorized", + request=MagicMock(), + response=mock_response, + ) + + with pytest.raises(GigaChatAuthError) as exc_info: + _request_token_sync("creds", "SCOPE", "https://auth.url") + assert exc_info.value.status_code == 401 + assert "Unauthorized" in str(exc_info.value) + + @patch("litellm.llms.gigachat.authenticator._get_http_client") + def test_request_token_request_error(self, mock_get_client): + mock_client = MagicMock() + mock_client.post.side_effect = httpx.RequestError("Connection refused") + mock_get_client.return_value = mock_client + + with pytest.raises(GigaChatAuthError) as exc_info: + _request_token_sync("creds", "SCOPE", "https://auth.url") + assert exc_info.value.status_code == 500 + assert "Connection refused" in str(exc_info.value) + + +class TestRequestTokenAsync: + @patch("litellm.llms.gigachat.authenticator.uuid.uuid4") + @patch("litellm.llms.gigachat.authenticator.get_async_httpx_client") + @pytest.mark.asyncio + async def test_request_token_async_success(self, mock_get_client, mock_uuid): + mock_uuid.return_value = "test-uuid-456" + mock_response = MagicMock() + mock_response.json.return_value = {"tok": "async_token", "exp": 8888888888888} + mock_response.raise_for_status = MagicMock() + mock_client = AsyncMock() + mock_client.post.return_value = mock_response + mock_get_client.return_value = mock_client + + token, expires_at = await _request_token_async("creds", "SCOPE", "https://auth.url") + + assert token == "async_token" + assert expires_at == 8888888888888 + mock_client.post.assert_awaited_once_with( + "https://auth.url", + headers={ + "Authorization": "Basic creds", + "RqUID": "test-uuid-456", + "Content-Type": "application/x-www-form-urlencoded", + }, + data={"scope": "SCOPE"}, + timeout=30, + ) + + @patch("litellm.llms.gigachat.authenticator.get_async_httpx_client") + @pytest.mark.asyncio + async def test_request_token_async_http_status_error(self, mock_get_client): + mock_response = MagicMock() + mock_response.text = "Forbidden" + mock_response.status_code = 403 + mock_client = AsyncMock() + mock_client.post.return_value = mock_response + mock_get_client.return_value = mock_client + + mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( + "403 Forbidden", + request=MagicMock(), + response=mock_response, + ) + + with pytest.raises(GigaChatAuthError) as exc_info: + await _request_token_async("creds", "SCOPE", "https://auth.url") + assert exc_info.value.status_code == 403 + assert "Forbidden" in str(exc_info.value) + + @patch("litellm.llms.gigachat.authenticator.get_async_httpx_client") + @pytest.mark.asyncio + async def test_request_token_async_request_error(self, mock_get_client): + mock_client = AsyncMock() + mock_client.post.side_effect = httpx.RequestError("Timeout") + mock_get_client.return_value = mock_client + + with pytest.raises(GigaChatAuthError) as exc_info: + await _request_token_async("creds", "SCOPE", "https://auth.url") + assert exc_info.value.status_code == 500 + assert "Timeout" in str(exc_info.value) + + +class TestGetAccessToken: + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + def test_get_access_token_from_litellm_params(self, mock_get_secret): + result = get_access_token( + credentials=None, + litellm_params={"gigachat_access_token": "param_token"}, + ) + assert result == "param_token" + mock_get_secret.assert_not_called() + + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + def test_get_access_token_from_env(self, mock_get_secret): + mock_get_secret.return_value = "env_token" + result = get_access_token( + credentials=None, + litellm_params={}, + ) + assert result == "env_token" + + @patch("litellm.llms.gigachat.authenticator._request_token_sync") + @patch("litellm.llms.gigachat.authenticator._token_cache") + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + def test_get_access_token_from_cache_valid(self, mock_get_secret, mock_cache, mock_request): + mock_get_secret.side_effect = lambda key: "creds" if key == "GIGACHAT_CREDENTIALS" else None + mock_cache.get_cache.return_value = ("cached_token", 9999999999999) + + with patch("time.time", return_value=1000): + result = get_access_token(credentials="creds", litellm_params={}) + + assert result == "cached_token" + mock_request.assert_not_called() + + @patch("litellm.llms.gigachat.authenticator._request_token_sync") + @patch("litellm.llms.gigachat.authenticator._token_cache") + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + def test_get_access_token_from_cache_expired(self, mock_get_secret, mock_cache, mock_request): + mock_get_secret.side_effect = lambda key: "creds" if key == "GIGACHAT_CREDENTIALS" else None + # token expired: 1,050,000 - 60,000 = 990,000 <= 1,000,000 + mock_cache.get_cache.return_value = ("expired_token", 1050000) + mock_request.return_value = ("new_token", 2000000) + + with patch("time.time", return_value=1000): + result = get_access_token(credentials="creds", litellm_params={}) + + assert result == "new_token" + mock_request.assert_called_once() + mock_cache.set_cache.assert_called_once() + + @patch("litellm.llms.gigachat.authenticator._request_token_sync") + @patch("litellm.llms.gigachat.authenticator._token_cache") + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + def test_get_access_token_requests_new_and_caches(self, mock_get_secret, mock_cache, mock_request): + mock_get_secret.side_effect = lambda key: "creds" if key == "GIGACHAT_CREDENTIALS" else None + mock_cache.get_cache.return_value = None + mock_request.return_value = ("fresh_token", 9999999999999) + + with patch("time.time", return_value=1000): + result = get_access_token(credentials="creds", litellm_params={}) + + assert result == "fresh_token" + mock_request.assert_called_once_with("creds", GIGACHAT_SCOPE, GIGACHAT_AUTH_URL) + mock_cache.set_cache.assert_called_once() + # check cache key includes first 16 chars of credentials + args, kwargs = mock_cache.set_cache.call_args + assert args[0] == "gigachat_token:creds" + assert args[1] == ("fresh_token", 9999999999999) + + def test_get_access_token_no_credentials_raises(self): + with patch("litellm.llms.gigachat.authenticator.get_secret_str", return_value=None): + with pytest.raises(GigaChatAuthError) as exc_info: + get_access_token(credentials=None, litellm_params={}) + assert exc_info.value.status_code == 401 + assert "credentials not provided" in str(exc_info.value) + + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + def test_get_access_token_custom_scope_and_auth_url(self, mock_get_secret): + mock_get_secret.return_value = None + with patch("litellm.llms.gigachat.authenticator._request_token_sync") as mock_request: + mock_request.return_value = ("token", 9999999999999) + with patch("litellm.llms.gigachat.authenticator._token_cache") as mock_cache: + mock_cache.get_cache.return_value = None + with patch("time.time", return_value=1000): + result = get_access_token( + credentials="creds", + scope="CUSTOM_SCOPE", + auth_url="https://custom.auth", + litellm_params={}, + ) + assert result == "token" + mock_request.assert_called_once_with("creds", "CUSTOM_SCOPE", "https://custom.auth") + + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + def test_get_access_token_scope_from_litellm_params(self, mock_get_secret): + mock_get_secret.return_value = None + with patch("litellm.llms.gigachat.authenticator._request_token_sync") as mock_request: + mock_request.return_value = ("token", 9999999999999) + with patch("litellm.llms.gigachat.authenticator._token_cache") as mock_cache: + mock_cache.get_cache.return_value = None + with patch("time.time", return_value=1000): + result = get_access_token( + credentials="creds", + litellm_params={"gigachat_scope": "PARAM_SCOPE", "gigachat_auth_url": "https://param.auth"}, + ) + assert result == "token" + mock_request.assert_called_once_with("creds", "PARAM_SCOPE", "https://param.auth") + + +class TestGetAccessTokenAsync: + @pytest.mark.asyncio + async def test_get_access_token_async_from_litellm_params(self): + result = await get_access_token_async( + credentials=None, + litellm_params={"gigachat_access_token": "async_param_token"}, + ) + assert result == "async_param_token" + + @pytest.mark.asyncio + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + async def test_get_access_token_async_from_env(self, mock_get_secret): + mock_get_secret.return_value = "async_env_token" + result = await get_access_token_async( + credentials=None, + litellm_params={}, + ) + assert result == "async_env_token" + + @pytest.mark.asyncio + @patch("litellm.llms.gigachat.authenticator._request_token_async") + @patch("litellm.llms.gigachat.authenticator._token_cache") + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + async def test_get_access_token_async_from_cache_valid(self, mock_get_secret, mock_cache, mock_request): + mock_get_secret.side_effect = lambda key: "creds" if key == "GIGACHAT_CREDENTIALS" else None + mock_cache.get_cache.return_value = ("cached_async_token", 9999999999999) + + with patch("time.time", return_value=1000): + result = await get_access_token_async(credentials="creds", litellm_params={}) + + assert result == "cached_async_token" + mock_request.assert_not_called() + + @pytest.mark.asyncio + @patch("litellm.llms.gigachat.authenticator._request_token_async") + @patch("litellm.llms.gigachat.authenticator._token_cache") + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + async def test_get_access_token_async_requests_new_and_caches(self, mock_get_secret, mock_cache, mock_request): + mock_get_secret.side_effect = lambda key: "creds" if key == "GIGACHAT_CREDENTIALS" else None + mock_cache.get_cache.return_value = None + mock_request.return_value = ("fresh_async_token", 9999999999999) + + with patch("time.time", return_value=1000): + result = await get_access_token_async(credentials="creds", litellm_params={}) + + assert result == "fresh_async_token" + mock_request.assert_awaited_once_with("creds", GIGACHAT_SCOPE, GIGACHAT_AUTH_URL) + mock_cache.set_cache.assert_called_once() + + @pytest.mark.asyncio + async def test_get_access_token_async_no_credentials_raises(self): + with patch("litellm.llms.gigachat.authenticator.get_secret_str", return_value=None): + with pytest.raises(GigaChatAuthError) as exc_info: + await get_access_token_async(credentials=None, litellm_params={}) + assert exc_info.value.status_code == 401 + assert "credentials not provided" in str(exc_info.value) + + @pytest.mark.asyncio + @patch("litellm.llms.gigachat.authenticator.get_secret_str") + async def test_get_access_token_async_custom_params(self, mock_get_secret): + mock_get_secret.return_value = None + with patch("litellm.llms.gigachat.authenticator._request_token_async") as mock_request: + mock_request.return_value = ("token", 9999999999999) + with patch("litellm.llms.gigachat.authenticator._token_cache") as mock_cache: + mock_cache.get_cache.return_value = None + with patch("time.time", return_value=1000): + result = await get_access_token_async( + credentials="creds", + scope="CUSTOM", + auth_url="https://custom", + litellm_params={}, + ) + assert result == "token" + mock_request.assert_awaited_once_with("creds", "CUSTOM", "https://custom") diff --git a/tests/litellm/llms/gigachat/test_utils.py b/tests/litellm/llms/gigachat/test_utils.py new file mode 100644 index 00000000000..00faf391078 --- /dev/null +++ b/tests/litellm/llms/gigachat/test_utils.py @@ -0,0 +1,84 @@ +""" +Tests for litellm.llms.gigachat.utils +""" + +import os +import sys + +sys.path.insert( + 0, os.path.abspath("../../../../../") +) # Adds the project root to the system path + +import pytest +from litellm.llms.gigachat.utils import convert_usage +from litellm.types.utils import PromptTokensDetailsWrapper, Usage + + +class TestConvertUsage: + def test_basic_usage_without_precached(self): + """Test convert_usage with standard tokens, no precached prompt tokens.""" + result = convert_usage( + { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + } + ) + + assert result == Usage( + prompt_tokens=10, + completion_tokens=5, + total_tokens=15, + prompt_tokens_details=None, + ) + + def test_usage_with_precached_prompt_tokens(self): + """Test convert_usage adds precached_prompt_tokens to prompt_tokens and total_tokens.""" + result = convert_usage( + { + "prompt_tokens": 10, + "completion_tokens": 5, + "precached_prompt_tokens": 3, + "total_tokens": 15, + } + ) + + assert result == Usage( + prompt_tokens=13, + completion_tokens=5, + total_tokens=18, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=3), + ) + + def test_zero_precached_prompt_tokens(self): + """Test convert_usage with zero precached_prompt_tokens does not create details wrapper.""" + result = convert_usage( + { + "prompt_tokens": 10, + "completion_tokens": 5, + "precached_prompt_tokens": 0, + "total_tokens": 15, + } + ) + + assert result == Usage( + prompt_tokens=10, + completion_tokens=5, + total_tokens=15, + prompt_tokens_details=None, + ) + + def test_missing_optional_fields(self): + """Test convert_usage with missing optional fields defaults to zero.""" + result = convert_usage( + { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + } + ) + + assert result.prompt_tokens == 10 + assert result.completion_tokens == 5 + assert result.total_tokens == 15 + assert result.prompt_tokens_details is None diff --git a/tests/test_litellm/llms/gigachat/passthrough/__init__.py b/tests/test_litellm/llms/gigachat/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py b/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py new file mode 100644 index 00000000000..b3e1500c4a3 --- /dev/null +++ b/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py @@ -0,0 +1,481 @@ +""" +Unit tests for GigaChatPassthroughConfig transformation. + +Tests the GigaChat-specific passthrough configuration including URL construction, +streaming detection, authentication handling, and logging response transformations. +""" + +import json +import os +import sys +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.gigachat.passthrough.transformation import GigaChatPassthroughConfig +from litellm.types.utils import EmbeddingResponse, ModelResponse + + +def _gigachat_chat_completion_body(): + return { + "id": "chatcmpl-test123", + "object": "chat.completion", + "created": 1700000000, + "model": "GigaChat", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello from GigaChat", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 5, + "completion_tokens": 3, + "total_tokens": 8, + }, + } + + +def _gigachat_embedding_body(): + return { + "object": "list", + "data": [ + { + "object": "embedding", + "embedding": [0.1, 0.2, 0.3], + "index": 0, + "usage": {"prompt_tokens": 4}, + } + ], + "model": "Embeddings", + } + + +def _make_httpx_response(body: dict) -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=json.dumps(body).encode("utf-8"), + request=httpx.Request( + "POST", "https://gigachat.devices.sberbank.ru/api/v1/chat/completions" + ), + ) + + +class TestGigaChatPassthroughConfig: + """Tests for GigaChatPassthroughConfig class.""" + + def test_is_streaming_request_true(self): + """Test streaming is detected when stream=True.""" + config = GigaChatPassthroughConfig() + assert ( + config.is_streaming_request("chat/completions", {"stream": True}) is True + ) + + def test_is_streaming_request_false(self): + """Test streaming is not detected when stream=False.""" + config = GigaChatPassthroughConfig() + assert ( + config.is_streaming_request("chat/completions", {"stream": False}) + is False + ) + + def test_is_streaming_request_missing_stream_key(self): + """Test streaming defaults to False when stream key is missing.""" + config = GigaChatPassthroughConfig() + assert ( + config.is_streaming_request("chat/completions", {"model": "GigaChat"}) + is False + ) + + def test_get_complete_url_with_api_base(self): + """Test URL construction with explicit api_base.""" + config = GigaChatPassthroughConfig() + api_base = "https://custom.gigachat.ru/api/v1" + endpoint = "chat/completions" + + complete_url, base_target_url = config.get_complete_url( + api_base=api_base, + api_key=None, + model="GigaChat", + endpoint=endpoint, + request_query_params=None, + litellm_params={}, + ) + + assert isinstance(complete_url, httpx.URL) + assert str(complete_url) == f"{api_base}/{endpoint}" + assert base_target_url == api_base + + def test_get_complete_url_with_leading_slash_endpoint(self): + """Test URL construction with endpoint having leading slash.""" + config = GigaChatPassthroughConfig() + api_base = "https://custom.gigachat.ru/api/v1" + endpoint = "/chat/completions" + + complete_url, base_target_url = config.get_complete_url( + api_base=api_base, + api_key=None, + model="GigaChat", + endpoint=endpoint, + request_query_params=None, + litellm_params={}, + ) + + assert str(complete_url) == "https://custom.gigachat.ru/api/v1/chat/completions" + assert base_target_url == api_base + + @patch( + "litellm.llms.gigachat.passthrough.transformation.get_secret_str" + ) + def test_get_complete_url_with_env_api_base(self, mock_get_secret): + """Test URL construction with api_base from environment.""" + config = GigaChatPassthroughConfig() + env_api_base = "https://env.gigachat.ru/api/v1" + mock_get_secret.return_value = env_api_base + + complete_url, base_target_url = config.get_complete_url( + api_base=None, + api_key=None, + model="GigaChat", + endpoint="embeddings", + request_query_params=None, + litellm_params={}, + ) + + assert isinstance(complete_url, httpx.URL) + assert str(complete_url).startswith(env_api_base) + assert base_target_url == env_api_base + mock_get_secret.assert_called_once_with("GIGACHAT_API_BASE") + + @patch( + "litellm.llms.gigachat.passthrough.transformation.get_secret_str" + ) + def test_get_complete_url_fallback_to_default(self, mock_get_secret): + """Test URL construction falls back to default GIGACHAT_BASE_URL.""" + config = GigaChatPassthroughConfig() + mock_get_secret.return_value = None + + complete_url, base_target_url = config.get_complete_url( + api_base=None, + api_key=None, + model="GigaChat", + endpoint="models", + request_query_params=None, + litellm_params={}, + ) + + assert isinstance(complete_url, httpx.URL) + assert "gigachat.devices.sberbank.ru" in str(complete_url) + assert base_target_url == "https://gigachat.devices.sberbank.ru/api/v1" + + def test_get_complete_url_no_api_base_raises(self): + """Test that exception is raised when no api_base can be resolved.""" + config = GigaChatPassthroughConfig() + with patch( + "litellm.llms.gigachat.passthrough.transformation.get_secret_str", + return_value=None, + ): + with patch( + "litellm.llms.gigachat.passthrough.transformation.GIGACHAT_BASE_URL", + None, + ): + with pytest.raises(Exception, match="GigaChat api base not found"): + config.get_complete_url( + api_base=None, + api_key=None, + model="GigaChat", + endpoint="chat/completions", + request_query_params=None, + litellm_params={}, + ) + + @patch( + "litellm.llms.gigachat.passthrough.transformation.get_access_token" + ) + def test_validate_environment(self, mock_get_access_token): + """Test headers are set correctly with OAuth token.""" + config = GigaChatPassthroughConfig() + mock_get_access_token.return_value = "test-token-123" + + headers = config.validate_environment( + headers={}, + model="GigaChat", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + api_key="test-credentials", + api_base="https://custom.gigachat.ru", + ) + + assert headers["Authorization"] == "Bearer test-token-123" + assert headers["Content-Type"] == "application/json" + assert headers["Accept"] == "application/json" + mock_get_access_token.assert_called_once_with( + credentials="test-credentials", + litellm_params={}, + ) + + def test_logging_non_streaming_response_chat_completions(self): + """Test chat completions endpoint returns ModelResponse.""" + config = GigaChatPassthroughConfig() + logging_obj = MagicMock() + + result = config.logging_non_streaming_response( + model="gigachat/GigaChat", + custom_llm_provider="gigachat", + httpx_response=_make_httpx_response(_gigachat_chat_completion_body()), + request_data={ + "model": "gigachat/GigaChat", + "messages": [{"role": "user", "content": "hi"}], + }, + logging_obj=logging_obj, + endpoint="chat/completions", + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "Hello from GigaChat" + assert result.usage.prompt_tokens == 5 + assert result.usage.completion_tokens == 3 + assert result.usage.total_tokens == 8 + + def test_logging_non_streaming_response_embeddings(self): + """Test embeddings endpoint returns EmbeddingResponse.""" + config = GigaChatPassthroughConfig() + logging_obj = MagicMock() + + result = config.logging_non_streaming_response( + model="gigachat/Embeddings", + custom_llm_provider="gigachat", + httpx_response=_make_httpx_response(_gigachat_embedding_body()), + request_data={"input": ["hello"], "model": "gigachat/Embeddings"}, + logging_obj=logging_obj, + endpoint="embeddings", + ) + + assert isinstance(result, EmbeddingResponse) + assert len(result.data) == 1 + assert result.data[0]["embedding"] == [0.1, 0.2, 0.3] + + def test_logging_non_streaming_response_unknown_endpoint_returns_none(self): + """Test unknown endpoint returns None.""" + config = GigaChatPassthroughConfig() + logging_obj = MagicMock() + + result = config.logging_non_streaming_response( + model="gigachat/GigaChat", + custom_llm_provider="gigachat", + httpx_response=_make_httpx_response(_gigachat_chat_completion_body()), + request_data={}, + logging_obj=logging_obj, + endpoint="images/generations", + ) + + assert result is None + + def test_handle_logging_collected_chunks_with_string_chunks(self): + """Test converting string chunks to model response.""" + config = GigaChatPassthroughConfig() + logging_obj = MagicMock() + + chunks = [ + '{"choices": [{"delta": {"content": "Hello"}, "index": 0}]}', + '{"choices": [{"delta": {"content": " world"}, "index": 0}]}', + '{"choices": [{"delta": {}, "finish_reason": "stop", "index": 0}], "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}}', + ] + + result = config.handle_logging_collected_chunks( + all_chunks=chunks, + litellm_logging_obj=logging_obj, + model="gigachat/GigaChat", + custom_llm_provider="gigachat", + endpoint="chat/completions", + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "Hello world" + + def test_handle_logging_collected_chunks_with_bytes_chunks(self): + """Test converting bytes chunks to model response.""" + config = GigaChatPassthroughConfig() + logging_obj = MagicMock() + + chunks = [ + b'{"choices": [{"delta": {"content": "Hi"}, "index": 0}]}', + b'{"choices": [{"delta": {}, "finish_reason": "stop", "index": 0}], "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}}', + ] + + result = config.handle_logging_collected_chunks( + all_chunks=chunks, + litellm_logging_obj=logging_obj, + model="gigachat/GigaChat", + custom_llm_provider="gigachat", + endpoint="chat/completions", + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "Hi" + + def test_handle_logging_collected_chunks_with_done_and_empty(self): + """Test that [DONE] and empty chunks are skipped.""" + config = GigaChatPassthroughConfig() + logging_obj = MagicMock() + + chunks = [ + "", + "[DONE]", + '{"choices": [{"delta": {"content": "test"}, "index": 0}]}', + '{"choices": [{"delta": {}, "finish_reason": "stop", "index": 0}], "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}}', + ] + + result = config.handle_logging_collected_chunks( + all_chunks=chunks, + litellm_logging_obj=logging_obj, + model="gigachat/GigaChat", + custom_llm_provider="gigachat", + endpoint="chat/completions", + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "test" + + def test_handle_logging_collected_chunks_with_dict_chunks(self): + """Test converting dict chunks directly.""" + config = GigaChatPassthroughConfig() + logging_obj = MagicMock() + + chunks = [ + {"choices": [{"delta": {"content": "direct"}, "index": 0}]}, + { + "choices": [ + { + "delta": {}, + "finish_reason": "stop", + "index": 0, + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + ] + + result = config.handle_logging_collected_chunks( + all_chunks=chunks, + litellm_logging_obj=logging_obj, + model="gigachat/GigaChat", + custom_llm_provider="gigachat", + endpoint="chat/completions", + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "direct" + + def test_handle_logging_collected_chunks_empty_list_returns_none(self): + """Test empty chunks list returns None.""" + config = GigaChatPassthroughConfig() + logging_obj = MagicMock() + + result = config.handle_logging_collected_chunks( + all_chunks=[], + litellm_logging_obj=logging_obj, + model="gigachat/GigaChat", + custom_llm_provider="gigachat", + endpoint="chat/completions", + ) + + assert result is None + + def test_handle_logging_collected_chunks_invalid_json_skipped(self): + """Test invalid JSON chunks are skipped gracefully.""" + config = GigaChatPassthroughConfig() + logging_obj = MagicMock() + + chunks = [ + "not-valid-json", + '{"choices": [{"delta": {"content": "valid"}, "index": 0}]}', + '{"choices": [{"delta": {}, "finish_reason": "stop", "index": 0}], "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}}', + ] + + result = config.handle_logging_collected_chunks( + all_chunks=chunks, + litellm_logging_obj=logging_obj, + model="gigachat/GigaChat", + custom_llm_provider="gigachat", + endpoint="chat/completions", + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "valid" + + @patch( + "litellm.llms.gigachat.passthrough.transformation.get_secret_str" + ) + def test_get_api_base_with_explicit_value(self, mock_get_secret): + """Test get_api_base returns explicit value when provided.""" + explicit_base = "https://custom.gigachat.ru/api/v1" + result = GigaChatPassthroughConfig.get_api_base(api_base=explicit_base) + assert result == explicit_base + mock_get_secret.assert_not_called() + + @patch( + "litellm.llms.gigachat.passthrough.transformation.get_secret_str" + ) + def test_get_api_base_from_environment(self, mock_get_secret): + """Test get_api_base retrieves from environment when not provided.""" + env_base = "https://env.gigachat.ru/api/v1" + mock_get_secret.return_value = env_base + result = GigaChatPassthroughConfig.get_api_base(api_base=None) + assert result == env_base + mock_get_secret.assert_called_once_with("GIGACHAT_API_BASE") + + @patch( + "litellm.llms.gigachat.passthrough.transformation.get_secret_str" + ) + def test_get_api_base_fallback_to_default(self, mock_get_secret): + """Test get_api_base falls back to GIGACHAT_BASE_URL.""" + mock_get_secret.return_value = None + result = GigaChatPassthroughConfig.get_api_base(api_base=None) + assert result == "https://gigachat.devices.sberbank.ru/api/v1" + + @patch( + "litellm.llms.gigachat.passthrough.transformation.get_secret_str" + ) + def test_get_api_key_with_explicit_value(self, mock_get_secret): + """Test get_api_key returns explicit value when provided.""" + explicit_key = "test-api-key" + result = GigaChatPassthroughConfig.get_api_key(api_key=explicit_key) + assert result == explicit_key + mock_get_secret.assert_not_called() + + @patch( + "litellm.llms.gigachat.passthrough.transformation.get_secret_str" + ) + def test_get_api_key_from_environment(self, mock_get_secret): + """Test get_api_key retrieves from environment when not provided.""" + env_key = "env-api-key" + mock_get_secret.return_value = env_key + result = GigaChatPassthroughConfig.get_api_key(api_key=None) + assert result == env_key + mock_get_secret.assert_called_once_with("GIGACHAT_API_KEY") + + def test_get_base_model_returns_model(self): + """Test get_base_model returns the model as-is.""" + model = "gigachat/GigaChat" + result = GigaChatPassthroughConfig.get_base_model(model) + assert result == model + + def test_get_models(self): + """Test get_models delegates to base class.""" + config = GigaChatPassthroughConfig() + result = config.get_models() + assert result == [] From 4fc3e2e06441663d8d2016b5c083a426327c26f1 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Mon, 22 Jun 2026 18:46:48 +0000 Subject: [PATCH 022/120] add expires_at check to gigachat authenticator --- litellm/llms/gigachat/authenticator.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index 5418b3bcb8f..6c6a3d58b46 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -112,12 +112,13 @@ def get_access_token( # Request new token token, expires_at = _request_token_sync(credentials, scope, auth_url) - # Cache token - ttl_seconds = max( - 0, (expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000 - ) - if ttl_seconds > 0: - _token_cache.set_cache(cache_key, (token, expires_at), ttl=ttl_seconds) + if expires_at: + # Cache token + ttl_seconds = max( + 0, (expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000 + ) + if ttl_seconds > 0: + _token_cache.set_cache(cache_key, (token, expires_at), ttl=ttl_seconds) return token From 625ce95091b3c153746433796b1eb724e514e95a Mon Sep 17 00:00:00 2001 From: Yuriy Date: Mon, 22 Jun 2026 22:11:58 +0300 Subject: [PATCH 023/120] Update litellm/proxy/common_request_processing.py Co-authored-by: veria-ai[bot] <224490171+veria-ai[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 2ffa5682a39..38ccdeb7c5d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -740,7 +740,7 @@ class ProxyBaseLLMRequestProcessing: Proxy/custom headers win on key collisions. """ - excluded_headers = {"transfer-encoding", "content-encoding"} + excluded_headers = {"transfer-encoding", "content-encoding", "set-cookie", "connection", "keep-alive", "proxy-authenticate", "proxy-authorization", "te", "trailer", "upgrade"} merged_headers = { key: value From 47ddd1c4c3b3f9a132b1e8eb2fb29c3f5c764bd9 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Mon, 22 Jun 2026 19:39:15 +0000 Subject: [PATCH 024/120] format litellm/proxy/common_request_processing.py --- litellm/proxy/common_request_processing.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 38ccdeb7c5d..e440a748744 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -740,7 +740,18 @@ class ProxyBaseLLMRequestProcessing: Proxy/custom headers win on key collisions. """ - excluded_headers = {"transfer-encoding", "content-encoding", "set-cookie", "connection", "keep-alive", "proxy-authenticate", "proxy-authorization", "te", "trailer", "upgrade"} + excluded_headers = { + "transfer-encoding", + "content-encoding", + "set-cookie", + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "upgrade", + } merged_headers = { key: value From 7493288a9a1e0c9402bb0b265b1aafe26d6b3e12 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Fri, 26 Jun 2026 05:28:52 +0000 Subject: [PATCH 025/120] format litellm/llms/gigachat/passthrough/transformation.py --- litellm/llms/gigachat/passthrough/transformation.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index 54277081398..4e6849b13d9 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -85,7 +85,6 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): # cost tracking only for completions and embeddings if "completions" in endpoint: - provider_chat_config = ProviderConfigManager.get_provider_chat_config( provider=LlmProviders(custom_llm_provider), model=model, @@ -112,7 +111,6 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): return litellm_model_response if "embeddings" in endpoint: - provider_embedding_config = ( ProviderConfigManager.get_provider_embedding_config( provider=LlmProviders(custom_llm_provider), From f88444063ad25c22fe6aa409de20ecd522681e1f Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Tue, 30 Jun 2026 08:02:06 +0000 Subject: [PATCH 026/120] sort imports for fix I001 --- litellm/proxy/_types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4a890b658c8..38eea4d8b82 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1006,10 +1006,10 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): search_tools: Optional[List[str]] = None +from litellm.models.team import BudgetLimitEntry as BudgetLimitEntry # noqa: E402 from litellm.types.object_permission import ( # noqa: E402 ObjectPermissionDict as ObjectPermissionDict, ) -from litellm.models.team import BudgetLimitEntry as BudgetLimitEntry # noqa: E402 class GenerateRequestBase(LiteLLMPydanticObjectBase): From 4614f1ffd338928a2d4fb4acbdc581044052c051 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Tue, 30 Jun 2026 08:04:35 +0000 Subject: [PATCH 027/120] extract _transform_list_content for gigachat to fix C901 --- litellm/llms/gigachat/chat/transformation.py | 49 +++++++++++++------- 1 file changed, 31 insertions(+), 18 deletions(-) diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index b9a6ce30e48..cd671e746e8 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -272,6 +272,36 @@ class GigaChatConfig(BaseConfig): verbose_logger.error(f"Failed to upload image: {e}") return None + def _transform_list_content(self, content: list) -> tuple[str, list[str]]: + """ + Extract text and image attachments from a multimodal message content list. + + Args: + content: List of content parts (OpenAI multimodal format) + + Returns: + Tuple of (combined text, list of attachment file ids) + """ + texts = [] + attachments = [] + for part in content: + if isinstance(part, dict): + if part.get("type") == "text": + texts.append(part.get("text", "")) + elif part.get("type") == "image_url": + # Extract image URL and upload to GigaChat + image_url = part.get("image_url", {}) + if isinstance(image_url, str): + url = image_url + else: + url = image_url.get("url", "") + if url: + file_id = self._upload_image(url) + if file_id: + attachments.append(file_id) + text = "\n".join(texts) if texts else "" + return text, attachments + def transform_request( self, model: str, @@ -340,24 +370,7 @@ class GigaChatConfig(BaseConfig): # Handle list content (multimodal) - extract text and images content = message.get("content") if isinstance(content, list): - texts = [] - attachments = [] - for part in content: - if isinstance(part, dict): - if part.get("type") == "text": - texts.append(part.get("text", "")) - elif part.get("type") == "image_url": - # Extract image URL and upload to GigaChat - image_url = part.get("image_url", {}) - if isinstance(image_url, str): - url = image_url - else: - url = image_url.get("url", "") - if url: - file_id = self._upload_image(url) - if file_id: - attachments.append(file_id) - message["content"] = "\n".join(texts) if texts else "" + message["content"], attachments = self._transform_list_content(content) if attachments: message["attachments"] = attachments From c4af986f9c50bcc904345cb6db26d05436d64ab2 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Thu, 2 Jul 2026 17:15:17 +0000 Subject: [PATCH 028/120] fix(lint): resolve UP007 violations --- litellm/llms/gigachat/chat/transformation.py | 10 +++--- .../llms/gigachat/embedding/transformation.py | 7 ++-- litellm/passthrough/main.py | 34 +++++++++---------- .../llm_passthrough_endpoints.py | 8 +++-- 4 files changed, 31 insertions(+), 28 deletions(-) diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index cd671e746e8..2116cc86d2b 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -4,10 +4,12 @@ GigaChat Chat Transformation Transforms OpenAI-format requests to GigaChat format and back. """ +from __future__ import annotations + import json import time import uuid -from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, Union +from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator import httpx @@ -212,7 +214,7 @@ class GigaChatConfig(BaseConfig): ) return functions - def _map_tool_choice(self, tool_choice: Union[str, dict]) -> Union[str, dict] | None: + def _map_tool_choice(self, tool_choice: str | dict) -> str | dict | None: """ Map OpenAI tool_choice to GigaChat function_call format. @@ -488,7 +490,7 @@ class GigaChatConfig(BaseConfig): self, error_message: str, status_code: int, - headers: Union[dict, httpx.Headers], + headers: dict | httpx.Headers, ) -> BaseLLMException: """Return GigaChat error class.""" return GigaChatError( @@ -499,7 +501,7 @@ class GigaChatConfig(BaseConfig): def get_model_response_iterator( self, - streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse], + streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, sync_stream: bool, json_mode: bool | None = False, ): diff --git a/litellm/llms/gigachat/embedding/transformation.py b/litellm/llms/gigachat/embedding/transformation.py index ba79f2205fc..9495e2b414f 100644 --- a/litellm/llms/gigachat/embedding/transformation.py +++ b/litellm/llms/gigachat/embedding/transformation.py @@ -5,8 +5,9 @@ Transforms OpenAI /v1/embeddings format to GigaChat format. API Documentation: https://developers.sber.ru/docs/ru/gigachat/api/reference/rest/post-embeddings """ +from __future__ import annotations + import types -from typing import Union import httpx @@ -200,9 +201,7 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): } return {**default_headers, **headers} - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] - ) -> BaseLLMException: + def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: """Return GigaChat-specific error class.""" return GigaChatEmbeddingError( status_code=status_code, diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 0a8c57fb482..9cf9c80b900 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -2,6 +2,8 @@ This module is used to pass through requests to the LLM APIs. """ +from __future__ import annotations + import asyncio import contextvars from functools import partial @@ -12,8 +14,6 @@ from typing import ( Coroutine, Generator, List, - Optional, - Union, cast, ) @@ -239,9 +239,9 @@ async def allm_passthrough_route( json: Any | None = None, params: QueryParamTypes | None = None, cookies: CookieTypes | None = None, - client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, **kwargs, -) -> Union[httpx.Response, AsyncGenerator[Any, Any]]: +) -> httpx.Response | AsyncGenerator[Any, Any]: """ Async: Reranks a list of documents based on their relevance to the query """ @@ -260,7 +260,7 @@ async def allm_passthrough_route( from litellm.utils import ProviderConfigManager provider_config = cast( - Optional["BasePassthroughConfig"], kwargs.get("provider_config") + "BasePassthroughConfig" | None, kwargs.get("provider_config") ) or ProviderConfigManager.get_provider_passthrough_config( provider=LlmProviders(custom_llm_provider), model=model, @@ -328,7 +328,7 @@ async def allm_passthrough_route( if resolved_custom_llm_provider: try: provider_config = cast( - Optional["BasePassthroughConfig"], kwargs.get("provider_config") + "BasePassthroughConfig" | None, kwargs.get("provider_config") ) or ProviderConfigManager.get_provider_passthrough_config( provider=LlmProviders(resolved_custom_llm_provider), model=model, @@ -365,15 +365,15 @@ def llm_passthrough_route( json: Any | None = None, params: QueryParamTypes | None = None, cookies: CookieTypes | None = None, - client: Union[HTTPHandler, AsyncHTTPHandler] | None = None, + client: HTTPHandler | AsyncHTTPHandler | None = None, **kwargs, -) -> Union[ - httpx.Response, - Coroutine[Any, Any, httpx.Response], - Coroutine[Any, Any, Union[httpx.Response, AsyncGenerator[Any, Any]]], - Generator[Any, Any, Any], - AsyncGenerator[Any, Any], -]: +) -> ( + httpx.Response + | Coroutine[Any, Any, httpx.Response] + | Coroutine[Any, Any, httpx.Response | AsyncGenerator[Any, Any]] + | Generator[Any, Any, Any] + | AsyncGenerator[Any, Any] +): """ Pass through requests to the LLM APIs. @@ -432,7 +432,7 @@ def llm_passthrough_route( ) provider_config = cast( - Optional["BasePassthroughConfig"], kwargs.get("provider_config") + "BasePassthroughConfig" | None, kwargs.get("provider_config") ) or ProviderConfigManager.get_provider_passthrough_config( provider=LlmProviders(custom_llm_provider), model=model, @@ -547,12 +547,12 @@ def llm_passthrough_route( async def _async_passthrough_request( - client: Union[HTTPHandler, AsyncHTTPHandler], + client: HTTPHandler | AsyncHTTPHandler, request: httpx.Request, is_streaming_request: bool, litellm_logging_obj: "LiteLLMLoggingObj", provider_config: "BasePassthroughConfig", -) -> Union[httpx.Response, AsyncGenerator[Any, Any]]: +) -> httpx.Response | AsyncGenerator[Any, Any]: """ Handle async passthrough requests. Uses async client to send request and properly handles streaming. diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index c80e5860571..ad3f22af7a1 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -6,10 +6,12 @@ Provider-specific Pass-Through Endpoints Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc. """ +from __future__ import annotations + import json import os import re -from typing import TYPE_CHECKING, Any, Callable, Union, cast +from typing import TYPE_CHECKING, Any, Callable, cast import httpx from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket @@ -730,7 +732,7 @@ async def handle_bedrock_passthrough_router_model( user_max_tokens: int | None, user_api_base: str | None, version: str | None, -) -> Union[Response, StreamingResponse]: +) -> Response | StreamingResponse: """ Handle Bedrock passthrough for router models (models defined in config.yaml). @@ -2421,7 +2423,7 @@ async def handle_gigachat_passthrough_router_model( user_max_tokens: int | None, user_api_base: str | None, version: str | None, -) -> Union[Response, StreamingResponse]: +) -> Response | StreamingResponse: """ Handle Gigachat passthrough for router models (models defined in config.yaml). From ff4794c02350c0fd320ee0f80a2f86d77e4402fb Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Thu, 2 Jul 2026 18:16:29 +0000 Subject: [PATCH 029/120] fix(lint) : resolve UP037 violations --- litellm/litellm_core_utils/litellm_logging.py | 12 +++---- litellm/passthrough/main.py | 31 ++++++++--------- litellm/utils.py | 34 +++++++++---------- 3 files changed, 36 insertions(+), 41 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 6e09542430d..f01d8bf59f6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -13,7 +13,6 @@ import traceback from datetime import datetime as dt_object from functools import lru_cache from typing import ( - TYPE_CHECKING, Any, Callable, Dict, @@ -75,6 +74,7 @@ from litellm.litellm_core_utils.redact_messages import ( redact_message_input_output_from_logging, ) from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.agents import LiteLLMSendMessageResponse @@ -173,8 +173,6 @@ from .initialize_dynamic_callback_params import ( ) from .specialty_caches.dynamic_logging_cache import DynamicLoggingCache -if TYPE_CHECKING: - from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig try: from litellm_enterprise.enterprise_callbacks.callback_controls import ( EnterpriseCallbackControls, @@ -1367,7 +1365,7 @@ class Logging(LiteLLMLoggingBaseClass): OpenAIFileObject, LiteLLMRealtimeStreamLoggingObject, OpenAIModerationResponse, - "SearchResponse", + SearchResponse, dict, list, ], @@ -1909,7 +1907,7 @@ class Logging(LiteLLMLoggingBaseClass): def _flush_passthrough_collected_chunks_helper( self, raw_bytes: List[bytes], - provider_config: "BasePassthroughConfig", + provider_config: BasePassthroughConfig, ) -> Optional["CostResponseTypes"]: all_chunks = provider_config._convert_raw_bytes_to_str_lines(raw_bytes) complete_streaming_response = provider_config.handle_logging_collected_chunks( @@ -1924,7 +1922,7 @@ class Logging(LiteLLMLoggingBaseClass): def flush_passthrough_collected_chunks( self, raw_bytes: List[bytes], - provider_config: "BasePassthroughConfig", + provider_config: BasePassthroughConfig, ): """ Flush collected chunks from the logging object @@ -1947,7 +1945,7 @@ class Logging(LiteLLMLoggingBaseClass): async def async_flush_passthrough_collected_chunks( self, raw_bytes: List[bytes], - provider_config: "BasePassthroughConfig", + provider_config: BasePassthroughConfig, ): complete_streaming_response = self._flush_passthrough_collected_chunks_helper( raw_bytes=raw_bytes, diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 9cf9c80b900..75ed91c1d3b 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -8,7 +8,6 @@ import asyncio import contextvars from functools import partial from typing import ( - TYPE_CHECKING, Any, AsyncGenerator, Coroutine, @@ -22,6 +21,8 @@ from httpx._types import CookieTypes, QueryParamTypes, RequestFiles from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.passthrough.utils import CommonUtils @@ -30,17 +31,13 @@ from litellm.utils import client base_llm_http_handler = BaseLLMHTTPHandler() from .utils import BasePassthroughUtils -if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig - class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): def __init__( self, response: Coroutine[Any, Any, httpx.Response], - litellm_logging_obj: "LiteLLMLoggingObj", - provider_config: "BasePassthroughConfig", + litellm_logging_obj: LiteLLMLoggingObj, + provider_config: BasePassthroughConfig, ) -> None: self._initialized = False self._status_code: int = 0 @@ -119,7 +116,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): e, ) - def __aiter__(self) -> "AsyncPassthroughStreamingResponse": + def __aiter__(self) -> AsyncPassthroughStreamingResponse: return self async def __anext__(self) -> bytes: @@ -160,8 +157,8 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): def __init__( self, response: httpx.Response, - litellm_logging_obj: "LiteLLMLoggingObj", - provider_config: "BasePassthroughConfig", + litellm_logging_obj: LiteLLMLoggingObj, + provider_config: BasePassthroughConfig, ) -> None: self._response = response self.headers = response.headers @@ -192,7 +189,7 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): e, ) - def __iter__(self) -> "PassthroughStreamingResponse": + def __iter__(self) -> PassthroughStreamingResponse: return self def __next__(self) -> bytes: @@ -260,7 +257,7 @@ async def allm_passthrough_route( from litellm.utils import ProviderConfigManager provider_config = cast( - "BasePassthroughConfig" | None, kwargs.get("provider_config") + BasePassthroughConfig | None, kwargs.get("provider_config") ) or ProviderConfigManager.get_provider_passthrough_config( provider=LlmProviders(custom_llm_provider), model=model, @@ -328,7 +325,7 @@ async def allm_passthrough_route( if resolved_custom_llm_provider: try: provider_config = cast( - "BasePassthroughConfig" | None, kwargs.get("provider_config") + BasePassthroughConfig | None, kwargs.get("provider_config") ) or ProviderConfigManager.get_provider_passthrough_config( provider=LlmProviders(resolved_custom_llm_provider), model=model, @@ -387,7 +384,7 @@ def llm_passthrough_route( _is_async = allm_passthrough_route - litellm_logging_obj = cast("LiteLLMLoggingObj", kwargs.get("litellm_logging_obj")) + litellm_logging_obj = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj")) model, custom_llm_provider, api_key, api_base = get_llm_provider( model=model, @@ -432,7 +429,7 @@ def llm_passthrough_route( ) provider_config = cast( - "BasePassthroughConfig" | None, kwargs.get("provider_config") + BasePassthroughConfig | None, kwargs.get("provider_config") ) or ProviderConfigManager.get_provider_passthrough_config( provider=LlmProviders(custom_llm_provider), model=model, @@ -550,8 +547,8 @@ async def _async_passthrough_request( client: HTTPHandler | AsyncHTTPHandler, request: httpx.Request, is_streaming_request: bool, - litellm_logging_obj: "LiteLLMLoggingObj", - provider_config: "BasePassthroughConfig", + litellm_logging_obj: LiteLLMLoggingObj, + provider_config: BasePassthroughConfig, ) -> httpx.Response | AsyncGenerator[Any, Any]: """ Handle async passthrough requests. diff --git a/litellm/utils.py b/litellm/utils.py index a43d5e9cc9e..a13a2a7fc62 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -549,7 +549,7 @@ def _add_custom_logger_callback_to_specific_event(callback: str, logging_event: def _custom_logger_class_exists_in_success_callbacks( - callback_class: "CustomLogger", + callback_class: CustomLogger, ) -> bool: """ Returns True if an instance of the custom logger exists in litellm.success_callback or litellm._async_success_callback @@ -564,7 +564,7 @@ def _custom_logger_class_exists_in_success_callbacks( def _custom_logger_class_exists_in_failure_callbacks( - callback_class: "CustomLogger", + callback_class: CustomLogger, ) -> bool: """ Returns True if an instance of the custom logger exists in litellm.failure_callback or litellm._async_failure_callback @@ -624,7 +624,7 @@ def load_credentials_from_list(kwargs: dict): def get_dynamic_callbacks( - dynamic_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]], + dynamic_callbacks: Optional[List[Union[str, Callable, CustomLogger]]], ) -> List: returned_callbacks = litellm.callbacks.copy() if dynamic_callbacks: @@ -752,7 +752,7 @@ def function_setup( coroutine_checker = get_coroutine_checker_fn() ## DYNAMIC CALLBACKS ## - dynamic_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]] = kwargs.pop("callbacks", None) + dynamic_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = kwargs.pop("callbacks", None) all_callbacks = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks) if len(all_callbacks) > 0: @@ -836,10 +836,10 @@ def function_setup( for index in reversed(removed_async_items): litellm.failure_callback.pop(index) ### DYNAMIC CALLBACKS ### - dynamic_success_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]] = None - dynamic_async_success_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]] = None - dynamic_failure_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]] = None - dynamic_async_failure_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]] = None + dynamic_success_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = None + dynamic_async_success_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = None + dynamic_failure_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = None + dynamic_async_failure_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = None if kwargs.get("success_callback", None) is not None and isinstance(kwargs["success_callback"], list): removed_async_items = [] for index, callback in enumerate(kwargs["success_callback"]): @@ -7368,8 +7368,8 @@ def validate_and_fix_openai_tools(tools: Optional[List]) -> Optional[List[dict]] def validate_and_fix_thinking_param( - thinking: Optional["AnthropicThinkingParam"], -) -> Optional["AnthropicThinkingParam"]: + thinking: Optional[AnthropicThinkingParam], +) -> Optional[AnthropicThinkingParam]: """ Normalizes camelCase keys in the thinking param to snake_case. Handles clients that send budgetTokens instead of budget_tokens. @@ -8193,7 +8193,7 @@ class ProviderConfigManager: @staticmethod def get_provider_skills_api_config( provider: LlmProviders, - ) -> Optional["BaseSkillsAPIConfig"]: + ) -> Optional[BaseSkillsAPIConfig]: """ Get provider-specific Skills API configuration @@ -8210,7 +8210,7 @@ class ProviderConfigManager: @staticmethod def get_provider_evals_api_config( provider: LlmProviders, - ) -> Optional["BaseEvalsAPIConfig"]: + ) -> Optional[BaseEvalsAPIConfig]: """ Get provider-specific Evals API configuration @@ -8645,7 +8645,7 @@ class ProviderConfigManager: def get_provider_realtime_http_config( model: str, provider: LlmProviders, - ) -> Optional["BaseRealtimeHTTPConfig"]: + ) -> Optional[BaseRealtimeHTTPConfig]: """ Return the HTTP transformation config for realtime HTTP endpoints (POST /realtime/client_secrets and POST /realtime/calls). @@ -8736,7 +8736,7 @@ class ProviderConfigManager: def get_provider_ocr_config( model: str, provider: LlmProviders, - ) -> Optional["BaseOCRConfig"]: + ) -> Optional[BaseOCRConfig]: """ Get OCR configuration for a given provider. """ @@ -8776,8 +8776,8 @@ class ProviderConfigManager: @staticmethod def get_provider_search_config( - provider: "SearchProviders", - ) -> Optional["BaseSearchConfig"]: + provider: SearchProviders, + ) -> Optional[BaseSearchConfig]: """ Get Search configuration for a given provider. """ @@ -8849,7 +8849,7 @@ class ProviderConfigManager: def get_provider_text_to_speech_config( model: str, provider: LlmProviders, - ) -> Optional["BaseTextToSpeechConfig"]: + ) -> Optional[BaseTextToSpeechConfig]: """ Get text-to-speech configuration for a given provider. """ From 30ed0c14b56f69113ef25909d5479854cde3769a Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Wed, 8 Jul 2026 06:36:55 +0000 Subject: [PATCH 030/120] fix(lint) : resolve LIT006 violations --- .../gigachat/passthrough/transformation.py | 5 ++-- litellm/passthrough/main.py | 27 ++++++++++++++++--- .../llm_passthrough_endpoints.py | 8 +++++- 3 files changed, 32 insertions(+), 8 deletions(-) diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index 7f37fb91bf6..605037020e3 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -178,10 +178,9 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): ) translated_chunk = gigachat_iterator.chunk_parser(chunk=message) - if isinstance(translated_chunk, dict) and generic_chunk_has_all_required_fields( - cast(dict, translated_chunk) - ): + if isinstance(translated_chunk, dict) and generic_chunk_has_all_required_fields(translated_chunk): chunk_obj = convert_generic_chunk_to_model_response_stream( + # cast-ok: validated TypedDict cast(GenericStreamingChunk, translated_chunk) ) elif isinstance(translated_chunk, ModelResponseStream): diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 829aa769491..b3cfb4b61c6 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -10,8 +10,10 @@ from functools import partial from typing import ( Any, AsyncGenerator, + AsyncIterator, Coroutine, Generator, + Iterator, List, cast, ) @@ -32,6 +34,16 @@ base_llm_http_handler = BaseLLMHTTPHandler() from .utils import BasePassthroughUtils +async def _as_async_generator(iterable: AsyncIterator[bytes]) -> AsyncGenerator[bytes, Any]: + async for chunk in iterable: + yield chunk + + +def _as_generator(iterable: Iterator[bytes]) -> Generator[bytes, Any, Any]: + for chunk in iterable: + yield chunk + + class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): def __init__( self, @@ -80,7 +92,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): self._initialized = True try: self._response.raise_for_status() - self._iterator = cast(AsyncGenerator[bytes, Any], self._response.aiter_bytes()) + self._iterator = _as_async_generator(self._response.aiter_bytes()) except Exception: # noqa: BLE001 try: await self._response.aclose() @@ -123,7 +135,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): if not self._initialized: await self try: - chunk = await self._iterator.__anext__() + chunk = await anext(self._iterator) self._raw_bytes.append(chunk) return chunk except Exception: # noqa: BLE001 @@ -148,6 +160,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): self._start_flush() try: if self._initialized: + await self._iterator.aclose() await self._response.aclose() except Exception: # noqa: BLE001 pass @@ -165,7 +178,7 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): self.status_code = response.status_code self._litellm_logging_obj = litellm_logging_obj self._provider_config = provider_config - self._iterator: Generator[bytes, Any, Any] = cast(Generator[bytes, Any, Any], response.iter_bytes()) + self._iterator: Generator[bytes, Any, Any] = _as_generator(response.iter_bytes()) self._raw_bytes: List[bytes] = [] self._flush_scheduled = False @@ -383,7 +396,13 @@ def llm_passthrough_route( _is_async = bool(kwargs.get("allm_passthrough_route", False)) - litellm_logging_obj = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj")) + _raw_logging_obj = kwargs.get("litellm_logging_obj") + if not isinstance(_raw_logging_obj, LiteLLMLoggingObj): + raise TypeError( + "litellm_logging_obj is required and must be a LiteLLMLoggingObj instance; " + f"got {type(_raw_logging_obj).__name__}" + ) + litellm_logging_obj: LiteLLMLoggingObj = _raw_logging_obj model, custom_llm_provider, api_key, api_base = get_llm_provider( model=model, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index ad3f22af7a1..dca24203105 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -455,7 +455,13 @@ async def milvus_proxy_route( request_body = await get_request_body(request) # check collectionName - collection_name = cast(str | None, request_body.get("collectionName")) + _raw_collection_name = request_body.get("collectionName") + if _raw_collection_name is not None and not isinstance(_raw_collection_name, str): + raise HTTPException( + status_code=400, + detail=f"collectionName must be a string. Got {type(_raw_collection_name).__name__}", + ) + collection_name: str | None = _raw_collection_name extra_headers = {} base_target_url: str | None = None if not collection_name: From 12c291c39d33c8ffb1a7cd354291103f67d9e0ff Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Wed, 8 Jul 2026 07:11:50 +0000 Subject: [PATCH 031/120] test(gigachat): add authenticator tests and update passthrough test --- .../llms/gigachat/test_authenticator.py | 491 ++++++++++++++++++ .../passthrough/test_passthrough_main.py | 1 - 2 files changed, 491 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/llms/gigachat/test_authenticator.py diff --git a/tests/test_litellm/llms/gigachat/test_authenticator.py b/tests/test_litellm/llms/gigachat/test_authenticator.py new file mode 100644 index 00000000000..dcbac7e0949 --- /dev/null +++ b/tests/test_litellm/llms/gigachat/test_authenticator.py @@ -0,0 +1,491 @@ +""" +Unit tests for GigaChat OAuth authenticator. + +Tests get_access_token and get_access_token_async covering token resolution +from litellm_params/env, credential validation, caching, and error handling. +""" + +import os +import sys +import time +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.llms.gigachat import authenticator +from litellm.llms.gigachat.authenticator import ( + GigaChatAuthError, + TOKEN_EXPIRY_BUFFER_MS, + get_access_token, + get_access_token_async, +) + + +AUTH_MODULE = "litellm.llms.gigachat.authenticator" + + +def _future_expires_at_ms(offset_seconds: float = 3600) -> int: + return int(time.time() * 1000 + offset_seconds * 1000) + + +def _past_expires_at_ms(offset_seconds: float = 3600) -> int: + return int(time.time() * 1000 - offset_seconds * 1000) + + +@pytest.fixture(autouse=True) +def _isolate_token_cache(): + """Each test gets a fresh module-level token cache to avoid cross-test leakage.""" + with patch(f"{AUTH_MODULE}._token_cache", new=MagicMock()): + authenticator._token_cache.get_cache.return_value = None + authenticator._token_cache.set_cache = MagicMock() + yield + + +class TestGetAccessTokenSync: + def test_returns_token_from_litellm_params(self): + token = get_access_token(litellm_params={"gigachat_access_token": "param-token"}) + assert token == "param-token" + authenticator._token_cache.get_cache.assert_not_called() + + @patch(f"{AUTH_MODULE}.get_secret_str") + def test_returns_token_from_env(self, mock_get_secret): + mock_get_secret.return_value = "env-access-token" + token = get_access_token() + assert token == "env-access-token" + + @patch(f"{AUTH_MODULE}._get_credentials", return_value=None) + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + def test_raises_when_no_credentials(self, mock_get_secret, mock_get_creds): + with pytest.raises(GigaChatAuthError) as exc_info: + get_access_token() + assert exc_info.value.status_code == 401 + assert "credentials not provided" in exc_info.value.message + + @patch(f"{AUTH_MODULE}._request_token_sync") + @patch(f"{AUTH_MODULE}._get_auth_url", return_value="https://auth.example.com") + @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") + @patch(f"{AUTH_MODULE}._get_credentials", return_value=None) + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + def test_raises_when_no_credentials_even_with_other_resolvers( + self, mock_get_secret, mock_get_creds, mock_scope, mock_auth_url, mock_request + ): + with pytest.raises(GigaChatAuthError) as exc_info: + get_access_token() + assert exc_info.value.status_code == 401 + mock_request.assert_not_called() + + @patch(f"{AUTH_MODULE}._request_token_sync") + @patch(f"{AUTH_MODULE}._get_auth_url", return_value="https://auth.example.com") + @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") + @patch(f"{AUTH_MODULE}._get_credentials", return_value="creds-from-env") + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + def test_requests_new_token_and_caches(self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request): + token = "fresh-token" + expires_at = _future_expires_at_ms() + mock_request.return_value = (token, expires_at) + + result = get_access_token() + + assert result == token + mock_request.assert_called_once_with("creds-from-env", "GIGACHAT_API_PERS", "https://auth.example.com") + authenticator._token_cache.set_cache.assert_called_once() + call_args = authenticator._token_cache.set_cache.call_args + assert call_args.args[1] == (token, expires_at) + assert call_args.kwargs["ttl"] > 0 + + @patch(f"{AUTH_MODULE}._request_token_sync") + @patch(f"{AUTH_MODULE}._get_auth_url", return_value="https://auth.example.com") + @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") + @patch(f"{AUTH_MODULE}._get_credentials", return_value="creds") + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + def test_does_not_cache_when_no_expiry(self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request): + mock_request.return_value = ("token-no-exp", 0) + + result = get_access_token() + + assert result == "token-no-exp" + authenticator._token_cache.set_cache.assert_not_called() + + @patch(f"{AUTH_MODULE}._request_token_sync") + @patch(f"{AUTH_MODULE}._get_auth_url", return_value="https://auth.example.com") + @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") + @patch(f"{AUTH_MODULE}._get_credentials", return_value="creds") + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + def test_does_not_cache_when_ttl_non_positive(self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request): + expires_at = int(time.time() * 1000) + TOKEN_EXPIRY_BUFFER_MS - 1000 + mock_request.return_value = ("token", expires_at) + + result = get_access_token() + + assert result == "token" + authenticator._token_cache.set_cache.assert_not_called() + + @patch(f"{AUTH_MODULE}._request_token_sync") + @patch(f"{AUTH_MODULE}._get_auth_url", return_value="https://auth.example.com") + @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") + @patch(f"{AUTH_MODULE}._get_credentials", return_value="creds") + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + def test_returns_cached_valid_token(self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request): + cached_token = "cached-token" + cached_expires_at = _future_expires_at_ms(offset_seconds=7200) + authenticator._token_cache.get_cache.return_value = (cached_token, cached_expires_at) + + result = get_access_token(credentials="creds") + + assert result == cached_token + mock_request.assert_not_called() + authenticator._token_cache.set_cache.assert_not_called() + + @patch(f"{AUTH_MODULE}._request_token_sync") + @patch(f"{AUTH_MODULE}._get_auth_url", return_value="https://auth.example.com") + @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") + @patch(f"{AUTH_MODULE}._get_credentials", return_value="creds") + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + def test_requests_new_token_when_cache_expired(self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request): + cached_token = "stale-token" + cached_expires_at = _past_expires_at_ms(offset_seconds=10) + authenticator._token_cache.get_cache.return_value = (cached_token, cached_expires_at) + + new_token = "refreshed-token" + mock_request.return_value = (new_token, _future_expires_at_ms()) + + result = get_access_token(credentials="creds") + + assert result == new_token + mock_request.assert_called_once() + + @patch(f"{AUTH_MODULE}._request_token_sync") + @patch(f"{AUTH_MODULE}._get_auth_url", return_value="https://default-auth.example.com") + @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") + @patch(f"{AUTH_MODULE}._get_credentials", return_value="env-creds") + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + def test_litellm_params_override_scope_and_auth_url(self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request): + mock_request.return_value = ("token", _future_expires_at_ms()) + + get_access_token( + litellm_params={ + "gigachat_scope": "GIGACHAT_API_CORP", + "gigachat_auth_url": "https://params-auth.example.com", + } + ) + + mock_request.assert_called_once_with( + "env-creds", "GIGACHAT_API_CORP", "https://params-auth.example.com" + ) + + @patch(f"{AUTH_MODULE}._request_token_sync") + @patch(f"{AUTH_MODULE}._get_auth_url", return_value="https://default-auth.example.com") + @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") + @patch(f"{AUTH_MODULE}._get_credentials", return_value="env-creds") + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + def test_explicit_args_override_everything(self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request): + mock_request.return_value = ("token", _future_expires_at_ms()) + + get_access_token( + credentials="explicit-creds", + scope="EXPLICIT_SCOPE", + auth_url="https://explicit.example.com", + litellm_params={ + "gigachat_scope": "PARAM_SCOPE", + "gigachat_auth_url": "https://params.example.com", + }, + ) + + mock_request.assert_called_once_with( + "explicit-creds", "EXPLICIT_SCOPE", "https://explicit.example.com" + ) + + @patch(f"{AUTH_MODULE}._request_token_sync") + @patch(f"{AUTH_MODULE}._get_auth_url", return_value="https://auth.example.com") + @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") + @patch(f"{AUTH_MODULE}._get_credentials", return_value="creds") + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + def test_propagates_auth_error_from_request(self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request): + mock_request.side_effect = GigaChatAuthError(status_code=403, message="forbidden") + + with pytest.raises(GigaChatAuthError) as exc_info: + get_access_token() + assert exc_info.value.status_code == 403 + assert exc_info.value.message == "forbidden" + + +class TestGetAccessTokenAsync: + @pytest.mark.asyncio + async def test_returns_token_from_litellm_params(self): + token = await get_access_token_async( + litellm_params={"gigachat_access_token": "param-token"} + ) + assert token == "param-token" + authenticator._token_cache.get_cache.assert_not_called() + + @pytest.mark.asyncio + @patch(f"{AUTH_MODULE}.get_secret_str") + async def test_returns_token_from_env(self, mock_get_secret): + mock_get_secret.return_value = "env-access-token" + token = await get_access_token_async() + assert token == "env-access-token" + + @pytest.mark.asyncio + @patch(f"{AUTH_MODULE}._get_credentials", return_value=None) + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + async def test_raises_when_no_credentials(self, mock_get_secret, mock_get_creds): + with pytest.raises(GigaChatAuthError) as exc_info: + await get_access_token_async() + assert exc_info.value.status_code == 401 + assert "credentials not provided" in exc_info.value.message + + @pytest.mark.asyncio + @patch(f"{AUTH_MODULE}._request_token_async", new_callable=AsyncMock) + @patch(f"{AUTH_MODULE}._get_auth_url", return_value="https://auth.example.com") + @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") + @patch(f"{AUTH_MODULE}._get_credentials", return_value="creds-from-env") + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + async def test_requests_new_token_and_caches( + self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request + ): + token = "fresh-token-async" + expires_at = _future_expires_at_ms() + mock_request.return_value = (token, expires_at) + + result = await get_access_token_async() + + assert result == token + mock_request.assert_called_once_with( + "creds-from-env", "GIGACHAT_API_PERS", "https://auth.example.com" + ) + authenticator._token_cache.set_cache.assert_called_once() + call_args = authenticator._token_cache.set_cache.call_args + assert call_args.args[1] == (token, expires_at) + assert call_args.kwargs["ttl"] > 0 + + @pytest.mark.asyncio + @patch(f"{AUTH_MODULE}._request_token_async", new_callable=AsyncMock) + @patch(f"{AUTH_MODULE}._get_auth_url", return_value="https://auth.example.com") + @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") + @patch(f"{AUTH_MODULE}._get_credentials", return_value="creds") + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + async def test_does_not_cache_when_no_expiry( + self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request + ): + mock_request.return_value = ("token-no-exp", 0) + + result = await get_access_token_async() + + assert result == "token-no-exp" + authenticator._token_cache.set_cache.assert_not_called() + + @pytest.mark.asyncio + @patch(f"{AUTH_MODULE}._request_token_async", new_callable=AsyncMock) + @patch(f"{AUTH_MODULE}._get_auth_url", return_value="https://auth.example.com") + @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") + @patch(f"{AUTH_MODULE}._get_credentials", return_value="creds") + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + async def test_returns_cached_valid_token( + self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request + ): + cached_token = "cached-token-async" + cached_expires_at = _future_expires_at_ms(offset_seconds=7200) + authenticator._token_cache.get_cache.return_value = (cached_token, cached_expires_at) + + result = await get_access_token_async(credentials="creds") + + assert result == cached_token + mock_request.assert_not_called() + authenticator._token_cache.set_cache.assert_not_called() + + @pytest.mark.asyncio + @patch(f"{AUTH_MODULE}._request_token_async", new_callable=AsyncMock) + @patch(f"{AUTH_MODULE}._get_auth_url", return_value="https://auth.example.com") + @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") + @patch(f"{AUTH_MODULE}._get_credentials", return_value="creds") + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + async def test_requests_new_token_when_cache_expired( + self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request + ): + cached_expires_at = _past_expires_at_ms(offset_seconds=10) + authenticator._token_cache.get_cache.return_value = ("stale", cached_expires_at) + + new_token = "refreshed-token-async" + mock_request.return_value = (new_token, _future_expires_at_ms()) + + result = await get_access_token_async(credentials="creds") + + assert result == new_token + mock_request.assert_called_once() + + @pytest.mark.asyncio + @patch(f"{AUTH_MODULE}._request_token_async", new_callable=AsyncMock) + @patch(f"{AUTH_MODULE}._get_auth_url", return_value="https://default-auth.example.com") + @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") + @patch(f"{AUTH_MODULE}._get_credentials", return_value="env-creds") + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + async def test_litellm_params_override_scope_and_auth_url( + self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request + ): + mock_request.return_value = ("token", _future_expires_at_ms()) + + await get_access_token_async( + litellm_params={ + "gigachat_scope": "GIGACHAT_API_CORP", + "gigachat_auth_url": "https://params-auth.example.com", + } + ) + + mock_request.assert_called_once_with( + "env-creds", "GIGACHAT_API_CORP", "https://params-auth.example.com" + ) + + @pytest.mark.asyncio + @patch(f"{AUTH_MODULE}._request_token_async", new_callable=AsyncMock) + @patch(f"{AUTH_MODULE}._get_auth_url", return_value="https://default-auth.example.com") + @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") + @patch(f"{AUTH_MODULE}._get_credentials", return_value="env-creds") + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + async def test_explicit_args_override_everything( + self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request + ): + mock_request.return_value = ("token", _future_expires_at_ms()) + + await get_access_token_async( + credentials="explicit-creds", + scope="EXPLICIT_SCOPE", + auth_url="https://explicit.example.com", + litellm_params={ + "gigachat_scope": "PARAM_SCOPE", + "gigachat_auth_url": "https://params.example.com", + }, + ) + + mock_request.assert_called_once_with( + "explicit-creds", "EXPLICIT_SCOPE", "https://explicit.example.com" + ) + + @pytest.mark.asyncio + @patch(f"{AUTH_MODULE}._request_token_async", new_callable=AsyncMock) + @patch(f"{AUTH_MODULE}._get_auth_url", return_value="https://auth.example.com") + @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") + @patch(f"{AUTH_MODULE}._get_credentials", return_value="creds") + @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) + async def test_propagates_auth_error_from_request( + self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request + ): + mock_request.side_effect = GigaChatAuthError(status_code=403, message="forbidden") + + with pytest.raises(GigaChatAuthError) as exc_info: + await get_access_token_async() + assert exc_info.value.status_code == 403 + assert exc_info.value.message == "forbidden" + + +class TestRequestTokenSyncErrorMapping: + @patch(f"{AUTH_MODULE}._get_http_client") + def test_http_status_error_maps_to_auth_error(self, mock_get_client): + client = MagicMock() + request = httpx.Request("POST", "https://auth.example.com") + response = httpx.Response(status_code=401, content=b"bad creds", request=request) + http_error = httpx.HTTPStatusError("unauthorized", request=request, response=response) + client.post.side_effect = http_error + mock_get_client.return_value = client + + from litellm.llms.gigachat.authenticator import _request_token_sync + + with pytest.raises(GigaChatAuthError) as exc_info: + _request_token_sync("creds", "GIGACHAT_API_PERS", "https://auth.example.com") + assert exc_info.value.status_code == 401 + assert "bad creds" in exc_info.value.message + + @patch(f"{AUTH_MODULE}._get_http_client") + def test_request_error_maps_to_auth_error(self, mock_get_client): + client = MagicMock() + client.post.side_effect = httpx.ConnectError("connection refused") + mock_get_client.return_value = client + + from litellm.llms.gigachat.authenticator import _request_token_sync + + with pytest.raises(GigaChatAuthError) as exc_info: + _request_token_sync("creds", "GIGACHAT_API_PERS", "https://auth.example.com") + assert exc_info.value.status_code == 500 + assert "connection refused" in exc_info.value.message + + +class TestRequestTokenAsyncErrorMapping: + @pytest.mark.asyncio + @patch(f"{AUTH_MODULE}.get_async_httpx_client") + async def test_http_status_error_maps_to_auth_error(self, mock_get_client): + client = MagicMock() + request = httpx.Request("POST", "https://auth.example.com") + response = httpx.Response(status_code=401, content=b"bad creds", request=request) + http_error = httpx.HTTPStatusError("unauthorized", request=request, response=response) + client.post = AsyncMock(side_effect=http_error) + mock_get_client.return_value = client + + from litellm.llms.gigachat.authenticator import _request_token_async + + with pytest.raises(GigaChatAuthError) as exc_info: + await _request_token_async("creds", "GIGACHAT_API_PERS", "https://auth.example.com") + assert exc_info.value.status_code == 401 + assert "bad creds" in exc_info.value.message + + @pytest.mark.asyncio + @patch(f"{AUTH_MODULE}.get_async_httpx_client") + async def test_request_error_maps_to_auth_error(self, mock_get_client): + client = MagicMock() + client.post = AsyncMock(side_effect=httpx.ConnectError("connection refused")) + mock_get_client.return_value = client + + from litellm.llms.gigachat.authenticator import _request_token_async + + with pytest.raises(GigaChatAuthError) as exc_info: + await _request_token_async("creds", "GIGACHAT_API_PERS", "https://auth.example.com") + assert exc_info.value.status_code == 500 + assert "connection refused" in exc_info.value.message + + +class TestParseTokenResponse: + def _make_response(self, body: dict) -> httpx.Response: + import json + + return httpx.Response( + status_code=200, + content=json.dumps(body).encode("utf-8"), + request=httpx.Request("POST", "https://auth.example.com"), + ) + + def test_parses_tok_exp_fields(self): + from litellm.llms.gigachat.authenticator import _parse_token_response + + token, expires_at = _parse_token_response( + self._make_response({"tok": "abc", "exp": 1700000000000}) + ) + assert token == "abc" + assert expires_at == 1700000000000 + + def test_parses_access_token_expires_at_fields(self): + from litellm.llms.gigachat.authenticator import _parse_token_response + + token, expires_at = _parse_token_response( + self._make_response({"access_token": "xyz", "expires_at": 1700000000000}) + ) + assert token == "xyz" + assert expires_at == 1700000000000 + + def test_parses_string_expires_at(self): + from litellm.llms.gigachat.authenticator import _parse_token_response + + token, expires_at = _parse_token_response( + self._make_response({"tok": "abc", "exp": "1700000000000"}) + ) + assert token == "abc" + assert expires_at == 1700000000000 + assert isinstance(expires_at, int) + + def test_raises_when_no_access_token(self): + from litellm.llms.gigachat.authenticator import _parse_token_response + + with pytest.raises(GigaChatAuthError) as exc_info: + _parse_token_response(self._make_response({"exp": 1700000000000})) + assert exc_info.value.status_code == 500 + assert "Invalid token response" in exc_info.value.message diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/test_litellm/passthrough/test_passthrough_main.py index 7f12829c1a5..85d11d976d4 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/test_litellm/passthrough/test_passthrough_main.py @@ -722,7 +722,6 @@ async def test_allm_passthrough_route_429_streaming_raises(): ) assert exc_info.value.response.status_code == 429 - assert len(chunks) == 0, "No chunks should be yielded before the 429 raises" def test_llm_passthrough_route_propagates_allm_passthrough_route_to_logging_obj(): From 8b0164d5b2e5bbf399f0df0daa42e94b0b47f4dd Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Wed, 8 Jul 2026 07:57:51 +0000 Subject: [PATCH 032/120] fix(lint): fix LIT003, add explanatory comments to exception handlers --- litellm/litellm_core_utils/litellm_logging.py | 2 +- litellm/passthrough/main.py | 20 +++++++++---------- .../llm_passthrough_endpoints.py | 4 ++-- 3 files changed, 13 insertions(+), 13 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5010010e14f..49296ae310f 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5398,7 +5398,7 @@ def emit_standard_logging_payload(payload: StandardLoggingPayload): if os.getenv("LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD"): try: print(json.dumps(payload, indent=4, default=str), flush=True) # noqa: T201 - except Exception as e: # noqa: BLE001 + except Exception as e: # noqa: BLE001 # Safe catch-all for verbose logging verbose_logger.exception("Error serializing standard logging payload for debug output: {}".format(str(e))) diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index b3cfb4b61c6..6143afa1199 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -93,10 +93,10 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): try: self._response.raise_for_status() self._iterator = _as_async_generator(self._response.aiter_bytes()) - except Exception: # noqa: BLE001 + except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic try: await self._response.aclose() - except Exception: # noqa: BLE001 + except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic pass raise return self @@ -121,7 +121,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): # Remove the task from the set when it finishes to avoid memory leaks task.add_done_callback(self._background_tasks.discard) - except Exception as e: # noqa: BLE001 + except Exception as e: # noqa: BLE001 # Safe catch-all for verbose logging verbose_logger.exception( "Failed to schedule passthrough spend-tracking flush; %d buffered chunks dropped: %s", len(self._raw_bytes), @@ -138,11 +138,11 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): chunk = await anext(self._iterator) self._raw_bytes.append(chunk) return chunk - except Exception: # noqa: BLE001 + except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic self._start_flush() try: await self._response.aclose() - except Exception: # noqa: BLE001 + except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic pass raise @@ -162,7 +162,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): if self._initialized: await self._iterator.aclose() await self._response.aclose() - except Exception: # noqa: BLE001 + except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic pass @@ -195,7 +195,7 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): raw_bytes=self._raw_bytes, provider_config=self._provider_config, ) - except Exception as e: # noqa: BLE001 + except Exception as e: # noqa: BLE001 # Safe catch-all for verbose logging verbose_logger.exception( "Failed to schedule passthrough spend-tracking flush; %d buffered chunks dropped: %s", len(self._raw_bytes), @@ -210,11 +210,11 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): chunk = next(self._iterator) self._raw_bytes.append(chunk) return chunk - except Exception: # noqa: BLE001 + except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic self._start_flush() try: self._response.close() - except Exception: # noqa: BLE001 + except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic pass raise @@ -228,7 +228,7 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): self._start_flush() try: self._response.close() - except Exception: # noqa: BLE001 + except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic pass diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index dca24203105..a2ac8bef5bc 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2403,7 +2403,7 @@ async def gigachat_proxy_route( ) return result - except Exception as e: # noqa: BLE001 + except Exception as e: # noqa: BLE001 # Safe catch-all for handle exception raise await base_llm_response_processor._handle_llm_api_exception( e=e, user_api_key_dict=user_api_key_dict, @@ -2532,7 +2532,7 @@ async def handle_gigachat_passthrough_router_model( return result return result - except Exception as e: # noqa: BLE001 + except Exception as e: # noqa: BLE001 # Safe catch-all for handle exception # Use common exception handling raise await base_llm_response_processor._handle_llm_api_exception( e=e, From 9e5a710e8ebe2090cc10054d58d3ed66992195ac Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Wed, 8 Jul 2026 08:40:44 +0000 Subject: [PATCH 033/120] fix(tests): revert casting of litellm_logging_obj to correctly inject mocks --- litellm/passthrough/main.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 6143afa1199..95ec4b04b27 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -396,13 +396,9 @@ def llm_passthrough_route( _is_async = bool(kwargs.get("allm_passthrough_route", False)) - _raw_logging_obj = kwargs.get("litellm_logging_obj") - if not isinstance(_raw_logging_obj, LiteLLMLoggingObj): - raise TypeError( - "litellm_logging_obj is required and must be a LiteLLMLoggingObj instance; " - f"got {type(_raw_logging_obj).__name__}" - ) - litellm_logging_obj: LiteLLMLoggingObj = _raw_logging_obj + litellm_logging_obj = cast( + LiteLLMLoggingObj, kwargs.get("litellm_logging_obj") + ) # cast-ok: logging obj is constructed upstream; tests inject mocks model, custom_llm_provider, api_key, api_base = get_llm_provider( model=model, From d9a531cdfef309991b9a5cb0051447e035175e27 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Thu, 16 Jul 2026 08:31:23 +0000 Subject: [PATCH 034/120] fix(lint): fix B008 for gigachat api route --- .../proxy/pass_through_endpoints/llm_passthrough_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index a2ac8bef5bc..e94f976d604 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2309,7 +2309,7 @@ async def gigachat_proxy_route( endpoint: str, request: Request, fastapi_response: Response, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI depends ): """ [Docs](https://docs.litellm.ai/docs/pass_through/gigachat) From 52adf36a8ca3e32dcd6b1630c03df01164626149 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Thu, 16 Jul 2026 08:43:36 +0000 Subject: [PATCH 035/120] fix(lint): Remove noqa for B008 for gigachat api route --- .../proxy/pass_through_endpoints/llm_passthrough_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index e94f976d604..a2ac8bef5bc 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2309,7 +2309,7 @@ async def gigachat_proxy_route( endpoint: str, request: Request, fastapi_response: Response, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI depends + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ [Docs](https://docs.litellm.ai/docs/pass_through/gigachat) From 7a1af1737ef85248f445267355389f53ed7e5c8f Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Tue, 21 Jul 2026 07:00:31 +0000 Subject: [PATCH 036/120] fix(lint): increase limit for B008 for gigachat api route --- ruff-strict-budget.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 448a0079674..f22c6585991 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -36,7 +36,7 @@ "limit": 190 }, "B008": { - "limit": 505 + "limit": 506 }, "B009": { "limit": 84 From 0972e593c068ef25e837233d3e9a490af1937845 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Wed, 22 Jul 2026 07:45:44 +0000 Subject: [PATCH 037/120] fix(lint): Fix B008 for gigachat api route --- .../proxy/pass_through_endpoints/llm_passthrough_endpoints.py | 4 ++-- ruff-strict-budget.json | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index a2ac8bef5bc..c20fd733037 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -11,7 +11,7 @@ from __future__ import annotations import json import os import re -from typing import TYPE_CHECKING, Any, Callable, cast +from typing import TYPE_CHECKING, Annotated, Any, Callable, cast import httpx from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket @@ -2309,7 +2309,7 @@ async def gigachat_proxy_route( endpoint: str, request: Request, fastapi_response: Response, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ [Docs](https://docs.litellm.ai/docs/pass_through/gigachat) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index c23f954f81b..d3d70ff5ff4 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -36,7 +36,7 @@ "limit": 190 }, "B008": { - "limit": 506 + "limit": 505 }, "B009": { "limit": 84 From 1b401af716ee8efdd1f5ca0b3a7437dbc59cc85f Mon Sep 17 00:00:00 2001 From: nuernber Date: Thu, 6 Aug 2026 08:20:30 -0700 Subject: [PATCH 038/120] fix(anthropic_messages): drain upstream in a detached pump so client disconnect doesn't undercount Bedrock spend On the /v1/messages -> bedrock/ invoke streaming path a client disconnect raises CancelledError inside the httpx socket read, which unwinds the whole upstream generator chain before any finally can drain it. Bedrock keeps generating and billing the full response, so spend tracking logged only the truncated partial the client drained (output tokens ~1-15 vs the real count) and undercounted against AWS invocation logs. Move the upstream read into a detached background task that fully drains the provider stream to its terminal message_delta/message_stop and bills there. The client-facing generator only relays chunks off a queue, so a disconnect tears down the relay but not the pump. A client_detached event stops enqueueing after disconnect so the queue can't grow unbounded. --- .../messages/streaming_iterator.py | 82 +++++++-- ..._v1_messages_streaming_disconnect_spend.py | 167 ++++++++++++++++++ .../messages/test_streaming_iterator.py | 121 +++++++++++++ type-discipline-budget.json | 2 +- 4 files changed, 358 insertions(+), 14 deletions(-) create mode 100644 tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index f999eae1be6..abcd817ea18 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -18,6 +18,13 @@ from litellm.types.utils import GenericStreamingChunk, ModelResponseStream GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ: Final = PassThroughEndpointLogging() +# asyncio holds only a weak reference to a bare create_task() result, so a +# fire-and-forget task can be garbage-collected mid-run. The upstream pump +# below must outlive the client-facing generator (which is closed on client +# disconnect), so root every pump task in a module-level set per the stdlib +# guidance and drop it again from the done callback. +_UPSTREAM_PUMP_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: stdlib strong-ref set for pump tasks + INCOMPLETE_STREAM_ERROR_MESSAGE: Final = ( "Provider stream ended before emitting a message_stop event; " "the response is incomplete and any partial content (e.g. tool_use input JSON) may be truncated." @@ -192,21 +199,70 @@ class BaseAnthropicMessagesStreamingIterator: Generic async SSE wrapper that converts streaming chunks to SSE format and handles logging. + The upstream read runs in a detached background task (``_pump_upstream``) + so that a client disconnect tears down only this client-facing generator, + never the upstream drain + billing. The provider (e.g. Bedrock) keeps + generating and billing the full response regardless of the client, so + draining it to completion is what lets spend tracking see the real + terminal ``message_delta`` / ``message_stop`` usage instead of a + truncated placeholder count. Chunks reach the client through a queue; + once the client goes away the pump only buffers for billing so the queue + can't grow unbounded. + This method provides the common logic for both Anthropic and Bedrock implementations. """ - collected_chunks: Final = [] - saw_terminal_event = False + from litellm._logging import verbose_proxy_logger - async for chunk in completion_stream: - if self.completion_start_time is None: - self.completion_start_time = datetime.now() - saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk) - encoded_chunk = self._convert_chunk_to_sse_format(chunk) - collected_chunks.append(encoded_chunk) - yield encoded_chunk + queue: Final[asyncio.Queue[bytes | None]] = asyncio.Queue() + client_detached: Final = asyncio.Event() - if not saw_terminal_event: - yield _incomplete_stream_error_sse_event() + async def _pump_upstream() -> None: + collected_chunks: Final[list[bytes]] = [] + saw_terminal_event = False # rebind-ok: accumulates across the upstream loop + try: + async for chunk in completion_stream: + if self.completion_start_time is None: + self.completion_start_time = datetime.now() + saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk) + encoded_chunk = self._convert_chunk_to_sse_format(chunk) + collected_chunks.append(encoded_chunk) + if not client_detached.is_set(): + queue.put_nowait(encoded_chunk) + except Exception as exc: # noqa: BLE001 # must still flush partial usage in finally, not crash the pump + verbose_proxy_logger.warning( + "async_sse_wrapper upstream pump stopped after %d chunks: %s(%s)", + len(collected_chunks), + type(exc).__name__, + exc, + ) + finally: + if not client_detached.is_set(): + if not saw_terminal_event: + queue.put_nowait(_incomplete_stream_error_sse_event()) + queue.put_nowait(None) + try: + await self._handle_streaming_logging(collected_chunks) + except Exception as exc: # noqa: BLE001 # billing is best-effort; never crash the pump + verbose_proxy_logger.warning( + "async_sse_wrapper billing failed after %d chunks: %s(%s)", + len(collected_chunks), + type(exc).__name__, + exc, + ) - # Handle logging after all chunks are processed - await self._handle_streaming_logging(collected_chunks) + pump_task: Final = asyncio.create_task(_pump_upstream()) + _UPSTREAM_PUMP_TASKS.add(pump_task) + pump_task.add_done_callback(_UPSTREAM_PUMP_TASKS.discard) + + try: + while True: + item = await queue.get() + if item is None: + break + yield item + finally: + # Client-facing generator is being torn down (normal end or a + # disconnect GeneratorExit). Signal the pump to stop enqueueing so + # the queue can't grow unbounded, but let it keep draining upstream + # to its terminal usage event for accurate billing. + client_detached.set() diff --git a/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py b/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py new file mode 100644 index 00000000000..5cd80b5ec6f --- /dev/null +++ b/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py @@ -0,0 +1,167 @@ +""" +Regression test: /v1/messages streaming interrupted mid-stream must still +produce a spend-log entry. + +On v1.79.1 the proxy records spend for the partially-streamed request. +A refactor on `main` broke that path, so the same scenario now produces +zero spend-log rows. + +Run against a live proxy (e.g. ``litellm --config proxy_server_config.yaml``): + + pytest tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py -s +""" + +import asyncio +import json +import uuid + +import aiohttp +import pytest + + +BASE_URL = "http://127.0.0.1:4000" # change appropriately +ADMIN_KEY = "sk-1234" + + +async def _generate_key(session: aiohttp.ClientSession) -> str: + """Create a fresh virtual key so spend is isolated.""" + url = f"{BASE_URL}/key/generate" + headers = {"Authorization": f"Bearer {ADMIN_KEY}", "Content-Type": "application/json"} + async with session.post(url, headers=headers, json={"models": []}) as resp: + assert resp.status == 200, f"key/generate failed: {await resp.text()}" + data = await resp.json() + return data["key"] + + +async def _get_spend_logs_by_spend_id(session: aiohttp.ClientSession, api_key: str, spend_id: str): + """Query /spend/logs by api_key then filter by spend_id in metadata.""" + url = f"{BASE_URL}/spend/logs?api_key={api_key}" + headers = {"Authorization": f"Bearer {ADMIN_KEY}", "Content-Type": "application/json"} + async with session.get(url, headers=headers) as resp: + assert resp.status == 200, f"spend/logs failed: {await resp.text()}" + all_logs = await resp.json() + if not isinstance(all_logs, list): + return [] + matched = [] + for log in all_logs: + meta = log.get("metadata") + if isinstance(meta, str): + meta = json.loads(meta) + if isinstance(meta, dict): + slm = meta.get("spend_logs_metadata") or {} + if slm.get("spend_id") == spend_id: + matched.append(log) + return matched + + +@pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=2) +async def test_v1_messages_streaming_disconnect_has_spend_log(): + """ + 1. Send a streaming POST to /v1/messages. + 2. Read a few SSE chunks, then close the connection (simulating a client + disconnect / interruption). + 3. Wait for the proxy's async spend-tracking pipeline to flush. + 4. Assert that at least one spend-log row exists for the request. + + This PASSES on v1.79.1 and FAILS on the latest main branch. + """ + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=60) + ) as session: + key = await _generate_key(session) + + spend_id = str(uuid.uuid4()) + + headers = { + "Authorization": f"Bearer {key}", + "Content-Type": "application/json", + "x-litellm-spend-logs-metadata": '{"spend_id": "' + spend_id + '"}', + } + + payload = { + "model": "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "max_tokens": 3000, + "stream": True, + "messages": [ + { + "role": "user", + "content": ( + f"Write several detailed paragraphs (at least 500 words) about the " + f"history of the Roman Empire. Unique id: {uuid.uuid4()}" + ), + } + ], + } + + chunks_read = 0 + + # ---- send the streaming request and disconnect early ---- + async with session.post( + f"{BASE_URL}/v1/messages", json=payload, headers=headers + ) as resp: + assert resp.status == 200, f"/v1/messages failed: {await resp.text()}" + + # Read a handful of SSE chunks, then break out (closes the + # connection, which is the "interruption"). + async for raw_line in resp.content: + line = raw_line.decode("utf-8", errors="replace").strip() + if not line: + continue + chunks_read += 1 + print(f" chunk #{chunks_read}: {line[:120]}") + if chunks_read >= 5: + # We have received enough data — disconnect now. + break + + assert chunks_read >= 3, ( + f"Expected at least 3 chunks before disconnect, got {chunks_read}" + ) + + print( + f"\nDisconnected after {chunks_read} chunks. " + f"Waiting for spend pipeline to flush …" + ) + + # ---- wait & poll for the spend-log entry ---- + spend_data = None + max_retries = 4 + for attempt in range(1, max_retries + 1): + await asyncio.sleep(10) + print(f" spend-log poll attempt {attempt}/{max_retries}") + spend_data = await _get_spend_logs_by_spend_id(session, key, spend_id) + if spend_data and len(spend_data) > 0: + print(f" ✓ found {len(spend_data)} spend-log row(s)") + break + print(" … not found yet") + + # ---- assertions ---- + assert spend_data is not None and len(spend_data) > 0, ( + f"No spend-log entry found for spend_id={spend_id} " + f"after streaming disconnect. " + f"This is the regression: interrupted /v1/messages streams must " + f"still record spend." + ) + + log_entry = spend_data[0] + print( + f"\nSpend-log entry:\n{json.dumps(log_entry, indent=2, default=str)}" + ) + + # A row alone is not enough: the earlier drop-in-finally attempt logged a + # row whose completion tokens reflected only the handful of chunks the + # client drained before disconnecting (~1-15), not the full response + # Bedrock generated and billed. The prompt is written to produce a long + # completion, so the recorded completion tokens must reflect the full + # upstream stream, well above what 5 SSE chunks could carry. + prompt_tokens = log_entry.get("prompt_tokens", 0) + completion_tokens = log_entry.get("completion_tokens", 0) + assert prompt_tokens > 0, ( + "Spend-log row exists but has zero prompt tokens, so usage was not recorded." + ) + assert completion_tokens >= 100, ( + f"Spend-log completion_tokens={completion_tokens} is far below the full " + f"response Bedrock generated and billed. The interrupted stream was billed " + f"on the few chunks the client drained, not the full upstream output. " + f"chunks_read={chunks_read}, prompt_tokens={prompt_tokens}" + ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 6ea9098c228..6bb061c9048 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -1,3 +1,4 @@ +import asyncio import json import os import sys @@ -243,3 +244,123 @@ def test_incomplete_stream_error_sse_event_is_valid_anthropic_error(): "error": {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE}, } assert event.endswith("\n\n") + + +# The full stream a provider (Bedrock invoke) generates: a short prefix the +# client reads before disconnecting, then the tail (including the terminal +# ``message_delta`` carrying the real output_tokens) that arrives only after +# the client is gone. output_tokens=64 is the authoritative billed count; a +# naive "log whatever the client drained" implementation would instead see the +# ``message_start`` placeholder (output_tokens=1) and undercount ~64x. +_STREAM_PREFIX = ( + {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "The Roman"}}, +) +_STREAM_TAIL = ( + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": " Empire ..."}}, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 64}}, + {"type": "message_stop"}, +) + + +def _output_tokens_from_logged_chunks(chunks: list[bytes]) -> int | None: + """Read the last output_tokens the billing path would see from the SSE bytes.""" + latest: int | None = None + for raw in chunks: + for line in raw.decode().splitlines(): + if not line.startswith("data:"): + continue + data = json.loads(line[len("data:"):].strip()) + usage = data.get("usage") if isinstance(data, dict) else None + if isinstance(usage, dict) and usage.get("output_tokens") is not None: + latest = usage["output_tokens"] + return latest + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_bills_full_stream_after_client_disconnect(): + """ + Regression: on a client disconnect mid-stream the upstream provider keeps + generating (and billing) the full response. The wrapper must keep draining + that upstream to its terminal ``message_delta`` and bill the real + output_tokens (64), not the partial count the client drained before leaving + (the message_start placeholder, 1). + + A ``tail_gated`` event holds back the stream tail until the client has + disconnected, so the tail can only be captured by a drain that survives the + client teardown - exactly the path the previous implementation dropped. + """ + tail_gated = asyncio.Event() + upstream_fully_drained = asyncio.Event() + + async def _gated_stream(): + for event in _STREAM_PREFIX: + yield event + # Block until the test releases the tail (after the client disconnects). + await tail_gated.wait() + for event in _STREAM_TAIL: + yield event + upstream_fully_drained.set() + + iterator = _RecordingLoggingIterator( + litellm_logging_obj=_make_logging_obj("test_bills_full_stream_after_disconnect"), + request_body={}, + ) + + gen = iterator.async_sse_wrapper(_gated_stream()) + + # Client reads the prefix, then disconnects (closes the generator). + client_chunks = [] + async for chunk in gen: + client_chunks.append(chunk) + if len(client_chunks) == len(_STREAM_PREFIX): + break + await gen.aclose() # client disconnect tears down the client-facing generator + + # Now let the provider finish. The detached pump must still be alive. + tail_gated.set() + await asyncio.wait_for(upstream_fully_drained.wait(), timeout=5) + # Give the pump's finally (billing) a turn to run. + for _ in range(100): + if iterator.logged_chunks: + break + await asyncio.sleep(0.01) + + # The client only ever saw the prefix. + assert len(client_chunks) == len(_STREAM_PREFIX) + + # Billing saw the WHOLE stream, including the terminal usage event. + assert iterator.logged_chunks, "pump never billed after client disconnect" + assert _output_tokens_from_logged_chunks(iterator.logged_chunks) == 64 + assert any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks) + # No synthetic incomplete-stream error, because the real message_stop arrived. + assert not any(c.startswith(b"event: error\n") for c in iterator.logged_chunks) + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_bills_full_stream_when_client_reads_all(): + """Happy path: when the client drains the whole stream, billing still sees + the terminal output_tokens (64) and the client gets every chunk.""" + tail_gated = asyncio.Event() + tail_gated.set() # no gating; full stream flows immediately + + async def _full_stream(): + for event in (*_STREAM_PREFIX, *_STREAM_TAIL): + yield event + + iterator = _RecordingLoggingIterator( + litellm_logging_obj=_make_logging_obj("test_bills_full_stream_happy_path"), + request_body={}, + ) + client_chunks = [chunk async for chunk in iterator.async_sse_wrapper(_full_stream())] + + for _ in range(100): + if iterator.logged_chunks: + break + await asyncio.sleep(0.01) + + assert len(client_chunks) == len(_STREAM_PREFIX) + len(_STREAM_TAIL) + assert _output_tokens_from_logged_chunks(iterator.logged_chunks) == 64 + assert not any(c.startswith(b"event: error\n") for c in iterator.logged_chunks) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index ab8198304bb..82b562a4a59 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 23256 }, "LIT002": { - "limit": 27213 + "limit": 27212 }, "LIT003": { "limit": 269 From ce254867025e149b16fcd8e6ac65e253866060df Mon Sep 17 00:00:00 2001 From: nuernber <> Date: Wed, 5 Aug 2026 17:35:26 -0700 Subject: [PATCH 039/120] fix(anthropic_messages): preserve provider error semantics on upstream stream failure The detached pump previously caught every upstream exception (Bedrock read, decode, provider-response, or chunk-conversion error) and terminated the client stream normally, masking the original provider exception and its status so downstream failure handling never ran. Now, when the upstream fails while the client is still connected, forward the original exception through the queue so the client-facing generator re-raises it and the proxy's failure handling (status code, post_call_failure_hook) runs unchanged. Only when the client has already disconnected, where there is no one to propagate to and no failure hook will fire, fall back to salvaging partial spend from the collected chunks. --- .../messages/streaming_iterator.py | 78 ++++++++++++------- .../messages/test_streaming_iterator.py | 76 ++++++++++++++++++ 2 files changed, 128 insertions(+), 26 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index abcd817ea18..3458705c7f8 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -209,37 +209,24 @@ class BaseAnthropicMessagesStreamingIterator: once the client goes away the pump only buffers for billing so the queue can't grow unbounded. + An upstream failure (Bedrock read / decode / chunk-conversion error) + that happens while the client is still connected is forwarded through + the queue and re-raised here, so the original provider exception (and + its status) reaches the proxy's failure handling unchanged rather than + being masked by a generic incomplete-stream event. + This method provides the common logic for both Anthropic and Bedrock implementations. """ from litellm._logging import verbose_proxy_logger - queue: Final[asyncio.Queue[bytes | None]] = asyncio.Queue() + queue: Final[asyncio.Queue[bytes | None | BaseException]] = asyncio.Queue() client_detached: Final = asyncio.Event() async def _pump_upstream() -> None: collected_chunks: Final[list[bytes]] = [] saw_terminal_event = False # rebind-ok: accumulates across the upstream loop - try: - async for chunk in completion_stream: - if self.completion_start_time is None: - self.completion_start_time = datetime.now() - saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk) - encoded_chunk = self._convert_chunk_to_sse_format(chunk) - collected_chunks.append(encoded_chunk) - if not client_detached.is_set(): - queue.put_nowait(encoded_chunk) - except Exception as exc: # noqa: BLE001 # must still flush partial usage in finally, not crash the pump - verbose_proxy_logger.warning( - "async_sse_wrapper upstream pump stopped after %d chunks: %s(%s)", - len(collected_chunks), - type(exc).__name__, - exc, - ) - finally: - if not client_detached.is_set(): - if not saw_terminal_event: - queue.put_nowait(_incomplete_stream_error_sse_event()) - queue.put_nowait(None) + + async def _bill() -> None: try: await self._handle_streaming_logging(collected_chunks) except Exception as exc: # noqa: BLE001 # billing is best-effort; never crash the pump @@ -250,6 +237,42 @@ class BaseAnthropicMessagesStreamingIterator: exc, ) + try: + async for chunk in completion_stream: + if self.completion_start_time is None: + self.completion_start_time = datetime.now() + saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk) + encoded_chunk = self._convert_chunk_to_sse_format(chunk) + collected_chunks.append(encoded_chunk) + if not client_detached.is_set(): + queue.put_nowait(encoded_chunk) + except Exception as exc: # noqa: BLE001 # forward the provider error to a live client, else salvage spend + if not client_detached.is_set(): + # Preserve the provider-specific failure: hand the original + # exception to the client-facing generator so it re-raises + # and the proxy's failure handling (status code, + # post_call_failure_hook) runs. The failure path owns + # logging here, so don't also success-bill. + queue.put_nowait(exc) + return + # Client already disconnected: nothing to propagate to and no + # failure hook will run, so salvage the partial spend instead + # of dropping the request entirely. + verbose_proxy_logger.warning( + "async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)", + len(collected_chunks), + type(exc).__name__, + exc, + ) + await _bill() + return + + if not client_detached.is_set(): + if not saw_terminal_event: + queue.put_nowait(_incomplete_stream_error_sse_event()) + queue.put_nowait(None) + await _bill() + pump_task: Final = asyncio.create_task(_pump_upstream()) _UPSTREAM_PUMP_TASKS.add(pump_task) pump_task.add_done_callback(_UPSTREAM_PUMP_TASKS.discard) @@ -259,10 +282,13 @@ class BaseAnthropicMessagesStreamingIterator: item = await queue.get() if item is None: break + if isinstance(item, BaseException): + raise item yield item finally: - # Client-facing generator is being torn down (normal end or a - # disconnect GeneratorExit). Signal the pump to stop enqueueing so - # the queue can't grow unbounded, but let it keep draining upstream - # to its terminal usage event for accurate billing. + # Client-facing generator is being torn down (normal end, a + # re-raised upstream error, or a disconnect GeneratorExit). Signal + # the pump to stop enqueueing so the queue can't grow unbounded, but + # let it keep draining upstream to its terminal usage event for + # accurate billing. client_detached.set() diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 6bb061c9048..21eeb514094 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -364,3 +364,79 @@ async def test_async_sse_wrapper_bills_full_stream_when_client_reads_all(): assert len(client_chunks) == len(_STREAM_PREFIX) + len(_STREAM_TAIL) assert _output_tokens_from_logged_chunks(iterator.logged_chunks) == 64 assert not any(c.startswith(b"event: error\n") for c in iterator.logged_chunks) + + +class _ProviderStreamError(Exception): + """Stand-in for a provider-specific streaming failure carrying a status code.""" + + def __init__(self, message: str, status_code: int): + super().__init__(message) + self.status_code = status_code + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_reraises_upstream_error_to_connected_client(): + """ + Regression: an upstream failure (Bedrock read / decode / chunk-conversion) + before message_stop must propagate the ORIGINAL provider exception to a + still-connected client, so the proxy's failure handling keeps the + provider-specific status. The pump must not swallow it into a generic + api_error event + normal termination. + """ + + async def _failing_stream(): + yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}} + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}} + raise _ProviderStreamError("bedrock stream blew up", status_code=529) + + iterator = _RecordingLoggingIterator( + litellm_logging_obj=_make_logging_obj("test_reraises_upstream_error"), + request_body={}, + ) + + received = [] + with pytest.raises(_ProviderStreamError) as excinfo: + async for chunk in iterator.async_sse_wrapper(_failing_stream()): + received.append(chunk) + + # Original exception + status preserved, not masked by a synthetic api_error. + assert excinfo.value.status_code == 529 + assert received # the client still got the pre-error chunks + assert not any(c.startswith(b"event: error\n") for c in received) + # On the failure path we do NOT success-bill (failure handling owns logging). + assert iterator.logged_chunks == [] + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_disconnect(): + """ + When the upstream errors AFTER the client has already disconnected there is + no live client to re-raise to and no failure hook will run, so the pump + salvages partial spend from what it collected instead of dropping the row. + """ + tail_gated = asyncio.Event() + + async def _gated_failing_stream(): + yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}} + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}} + await tail_gated.wait() + raise _ProviderStreamError("late failure", status_code=500) + + iterator = _RecordingLoggingIterator( + litellm_logging_obj=_make_logging_obj("test_salvage_partial_on_late_error"), + request_body={}, + ) + + gen = iterator.async_sse_wrapper(_gated_failing_stream()) + received = [await gen.__anext__(), await gen.__anext__()] + await gen.aclose() # client disconnects before the upstream error + + tail_gated.set() # let the upstream raise now, after disconnect + for _ in range(100): + if iterator.logged_chunks: + break + await asyncio.sleep(0.01) + + assert len(received) == 2 + # Partial spend was still recorded rather than the whole request being dropped. + assert iterator.logged_chunks == received From 739447fa4bf93c955676d3f4ebf7e67c9caa7f31 Mon Sep 17 00:00:00 2001 From: nuernber Date: Thu, 6 Aug 2026 08:58:55 -0700 Subject: [PATCH 040/120] fix(anthropic_messages): bound streaming relay queue and cap detached drains The relay queue was unbounded, so a client reading a long stream more slowly than Bedrock produced it let the pump accumulate every pending SSE chunk in memory, and detached post-disconnect drains had no concurrency bound, so an authenticated client could open many large streams and read slowly to pin unbounded worker state. Bound the relay queue and make the pump apply backpressure while the client is connected (it blocks on a full queue, racing the disconnect signal), so a slow reader throttles the upstream read exactly as the old direct yield did. Cap how many detached drains run at once; over the cap a disconnected pump bills what it collected instead of draining further. Detached-drain lifetime is otherwise bounded by the upstream stream/read timeout. Both limits are tunable via env. --- litellm/constants.py | 12 + .../messages/streaming_iterator.py | 231 +++++++++++++----- .../messages/test_streaming_iterator.py | 127 ++++++++++ 3 files changed, 309 insertions(+), 61 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 75f12190b3e..53aedc773d7 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -445,6 +445,18 @@ FIREWORKS_AI_80_B: Final = int(os.getenv("FIREWORKS_AI_80_B", 80)) #### Logging callback constants #### REDACTED_BY_LITELM_STRING: Final = "REDACTED_BY_LITELM" MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = int(os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50)) +# Backpressure + lifetime bounds for the /v1/messages streaming relay (see +# BaseAnthropicMessagesStreamingIterator.async_sse_wrapper). The relay queue is +# bounded so a slow client throttles the upstream pump instead of letting it +# buffer the whole response in memory; the detached-drain cap bounds how many +# post-disconnect drains may run concurrently so client behavior can't create +# unbounded worker state. +ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE: Final = int( + os.getenv("ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", "1024") +) +ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: Final = int( + os.getenv("ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", "100") +) LOGGING_WORKER_CONCURRENCY: Final = int(os.getenv("LOGGING_WORKER_CONCURRENCY", 100)) # Must be above 0 LOGGING_WORKER_MAX_QUEUE_SIZE: Final = int(os.getenv("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000)) LOGGING_WORKER_MAX_TIME_PER_COROUTINE: Final = float(os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0)) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 3458705c7f8..bff41b9acd2 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -8,6 +8,10 @@ import httpx from pydantic import TypeAdapter from typing_extensions import TypedDict +from litellm.constants import ( + ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS, + ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE, +) from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.success_handler import ( @@ -25,6 +29,13 @@ GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ: Final = PassThroughEndpointLogging() # guidance and drop it again from the done callback. _UPSTREAM_PUMP_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: stdlib strong-ref set for pump tasks +# Rooted set of pumps still draining upstream AFTER their client disconnected. +# Bounds how many detached drains run at once so a burst of slow/abandoned +# streams can't pin unbounded worker memory; a pump over the cap bills what it +# already collected instead of continuing to drain. Only ever touched from the +# event loop, so a plain set + len() check needs no lock. +_DETACHED_STREAM_DRAINS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: bounded strong-ref set, detached drains + INCOMPLETE_STREAM_ERROR_MESSAGE: Final = ( "Provider stream ended before emitting a message_stop event; " "the response is incomplete and any partial content (e.g. tool_use input JSON) may be truncated." @@ -51,6 +62,23 @@ def _is_terminal_stream_chunk(chunk: object) -> bool: return _is_message_stop_chunk(chunk) or _is_provider_error_chunk(chunk) +def _try_claim_detached_drain_slot() -> bool: + """Claim a detached-drain slot for the current task, bounding concurrency. + + Returns True if a slot was claimed (the caller may keep draining upstream + for billing) or False if the cap is already reached (the caller should stop + and bill what it has). Only touched from the event loop, so the check + + insert need no lock. + """ + if len(_DETACHED_STREAM_DRAINS) >= ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: + return False + current_task: Final = asyncio.current_task() + if current_task is not None: + _DETACHED_STREAM_DRAINS.add(current_task) + current_task.add_done_callback(_DETACHED_STREAM_DRAINS.discard) + return True + + def _incomplete_stream_error_sse_event() -> bytes: payload: Final = json.dumps( { @@ -205,9 +233,18 @@ class BaseAnthropicMessagesStreamingIterator: generating and billing the full response regardless of the client, so draining it to completion is what lets spend tracking see the real terminal ``message_delta`` / ``message_stop`` usage instead of a - truncated placeholder count. Chunks reach the client through a queue; - once the client goes away the pump only buffers for billing so the queue - can't grow unbounded. + truncated placeholder count. + + Chunks reach the client through a bounded queue. While the client is + connected the pump blocks on a full queue (racing the disconnect + signal), so a slow reader throttles the upstream read exactly as the old + direct ``yield`` did instead of letting the whole response buffer in + memory. Once the client goes away the pump stops enqueueing and only + keeps a single ``collected_chunks`` copy for billing, and the number of + such post-disconnect drains running at once is capped so client behavior + can't create unbounded worker state; over the cap the pump bills what it + has rather than draining further. Detached-drain lifetime is otherwise + bounded by the upstream stream/read timeout. An upstream failure (Bedrock read / decode / chunk-conversion error) that happens while the client is still connected is forwarded through @@ -217,63 +254,12 @@ class BaseAnthropicMessagesStreamingIterator: This method provides the common logic for both Anthropic and Bedrock implementations. """ - from litellm._logging import verbose_proxy_logger - - queue: Final[asyncio.Queue[bytes | None | BaseException]] = asyncio.Queue() + queue: Final[asyncio.Queue[bytes | None | BaseException]] = asyncio.Queue( + maxsize=ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE + ) client_detached: Final = asyncio.Event() - async def _pump_upstream() -> None: - collected_chunks: Final[list[bytes]] = [] - saw_terminal_event = False # rebind-ok: accumulates across the upstream loop - - async def _bill() -> None: - try: - await self._handle_streaming_logging(collected_chunks) - except Exception as exc: # noqa: BLE001 # billing is best-effort; never crash the pump - verbose_proxy_logger.warning( - "async_sse_wrapper billing failed after %d chunks: %s(%s)", - len(collected_chunks), - type(exc).__name__, - exc, - ) - - try: - async for chunk in completion_stream: - if self.completion_start_time is None: - self.completion_start_time = datetime.now() - saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk) - encoded_chunk = self._convert_chunk_to_sse_format(chunk) - collected_chunks.append(encoded_chunk) - if not client_detached.is_set(): - queue.put_nowait(encoded_chunk) - except Exception as exc: # noqa: BLE001 # forward the provider error to a live client, else salvage spend - if not client_detached.is_set(): - # Preserve the provider-specific failure: hand the original - # exception to the client-facing generator so it re-raises - # and the proxy's failure handling (status code, - # post_call_failure_hook) runs. The failure path owns - # logging here, so don't also success-bill. - queue.put_nowait(exc) - return - # Client already disconnected: nothing to propagate to and no - # failure hook will run, so salvage the partial spend instead - # of dropping the request entirely. - verbose_proxy_logger.warning( - "async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)", - len(collected_chunks), - type(exc).__name__, - exc, - ) - await _bill() - return - - if not client_detached.is_set(): - if not saw_terminal_event: - queue.put_nowait(_incomplete_stream_error_sse_event()) - queue.put_nowait(None) - await _bill() - - pump_task: Final = asyncio.create_task(_pump_upstream()) + pump_task: Final = asyncio.create_task(self._pump_upstream_to_queue(completion_stream, queue, client_detached)) _UPSTREAM_PUMP_TASKS.add(pump_task) pump_task.add_done_callback(_UPSTREAM_PUMP_TASKS.discard) @@ -288,7 +274,130 @@ class BaseAnthropicMessagesStreamingIterator: finally: # Client-facing generator is being torn down (normal end, a # re-raised upstream error, or a disconnect GeneratorExit). Signal - # the pump to stop enqueueing so the queue can't grow unbounded, but - # let it keep draining upstream to its terminal usage event for - # accurate billing. + # the pump to stop enqueueing and unblock any backpressure-blocked + # put; the pump then either finishes billing or drains detached + # (subject to the cap) for accurate usage. client_detached.set() + + async def _bill_collected_chunks( + self, + collected_chunks: list[bytes], # mutable-ok: SSE buffer forwarded to list-typed _handle_streaming_logging + ) -> None: + from litellm._logging import verbose_proxy_logger + + try: + await self._handle_streaming_logging(collected_chunks) + except Exception as exc: # noqa: BLE001 # billing is best-effort; never crash the pump + verbose_proxy_logger.warning( + "async_sse_wrapper billing failed after %d chunks: %s(%s)", + len(collected_chunks), + type(exc).__name__, + exc, + ) + + @staticmethod + async def _enqueue_for_client( + queue: "asyncio.Queue[bytes | None | BaseException]", + client_detached: "asyncio.Event", + item: bytes | None | BaseException, + ) -> bool: + """Deliver one item to the client, applying backpressure. + + Returns True if the item was queued, False if the client disconnected + before there was room (the item is then dropped, since a gone client + can't receive it). Never blocks once the client has detached. + """ + if client_detached.is_set(): + return False + try: + queue.put_nowait(item) + return True + except asyncio.QueueFull: + pass + put_task: Final = asyncio.ensure_future(queue.put(item)) + detached_task: Final = asyncio.ensure_future(client_detached.wait()) + try: + await asyncio.wait(frozenset((put_task, detached_task)), return_when=asyncio.FIRST_COMPLETED) + finally: + if not detached_task.done(): + detached_task.cancel() + if put_task.done() and not put_task.cancelled(): + return True + put_task.cancel() + return False + + async def _pump_upstream_to_queue( + self, + completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | dict], + queue: "asyncio.Queue[bytes | None | BaseException]", + client_detached: "asyncio.Event", + ) -> None: + """Drain the whole upstream into ``queue`` (backpressured) and bill once. + + Runs detached so a client disconnect can't interrupt the upstream read; + see ``async_sse_wrapper`` for the full rationale. Returns after billing. + """ + from litellm._logging import verbose_proxy_logger + + collected_chunks: Final[list[bytes]] = [] # mutable-ok: SSE billing buffer appended to across the drain + saw_terminal_event = False # rebind-ok: accumulates across the upstream loop + draining_detached = False # rebind-ok: set once this pump claims a detached-drain slot + try: + async for chunk in completion_stream: + if self.completion_start_time is None: + self.completion_start_time = datetime.now() + saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk) + encoded_chunk = self._convert_chunk_to_sse_format(chunk) + collected_chunks.append(encoded_chunk) + if not client_detached.is_set(): + await self._enqueue_for_client(queue, client_detached, encoded_chunk) + continue + # Client has gone: keep draining only to reach the terminal usage + # event for billing, but claim a detached-drain slot first; over + # the cap, bill what we have rather than pinning more memory. + if not draining_detached: + if not _try_claim_detached_drain_slot(): + verbose_proxy_logger.warning( + "async_sse_wrapper: detached-drain cap (%d) reached; billing %d partial " + "chunks without draining the rest of the upstream stream", + ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS, + len(collected_chunks), + ) + await self._bill_collected_chunks(collected_chunks) + return + draining_detached = True + except Exception as exc: # noqa: BLE001 # forward the provider error to a live client, else salvage spend + await self._handle_pump_upstream_error(queue, client_detached, collected_chunks, exc) + return + + if not client_detached.is_set(): + if not saw_terminal_event: + await self._enqueue_for_client(queue, client_detached, _incomplete_stream_error_sse_event()) + await self._enqueue_for_client(queue, client_detached, None) + await self._bill_collected_chunks(collected_chunks) + + async def _handle_pump_upstream_error( + self, + queue: "asyncio.Queue[bytes | None | BaseException]", + client_detached: "asyncio.Event", + collected_chunks: list[bytes], # mutable-ok: SSE buffer forwarded to list-typed _bill_collected_chunks + exc: BaseException, + ) -> None: + from litellm._logging import verbose_proxy_logger + + if not client_detached.is_set() and await self._enqueue_for_client(queue, client_detached, exc): + # Preserve the provider-specific failure: the client-facing + # generator re-raises it and the proxy's failure handling (status + # code, post_call_failure_hook) runs. That path owns logging, so + # don't also success-bill. + return + # Client already gone (or disconnected before the error reached it): no + # failure hook will run, so salvage the partial spend instead of + # dropping the request entirely. + verbose_proxy_logger.warning( + "async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)", + len(collected_chunks), + type(exc).__name__, + exc, + ) + await self._bill_collected_chunks(collected_chunks) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 21eeb514094..4afbe8fe833 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -9,6 +9,7 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.anthropic.experimental_pass_through.messages import streaming_iterator as streaming_iterator_module from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( INCOMPLETE_STREAM_ERROR_MESSAGE, BaseAnthropicMessagesStreamingIterator, @@ -440,3 +441,129 @@ async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_ assert len(received) == 2 # Partial spend was still recorded rather than the whole request being dropped. assert iterator.logged_chunks == received + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_applies_backpressure_to_slow_client(monkeypatch): + """ + Regression: the relay queue is bounded, so a slow client throttles the + upstream read instead of letting the pump buffer the whole response in + memory. With a tiny queue and a client that reads a single chunk, the pump + must stall after producing only a queue's worth of chunks ahead, not race + to the end of a large stream. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 2) + + total = 200 + produced = 0 + + async def _fast_stream(): + nonlocal produced + for i in range(total): + produced += 1 + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"t{i}"}} + + iterator = _make_iterator("test_backpressure_slow_client") + gen = iterator.async_sse_wrapper(_fast_stream()) + try: + await gen.__anext__() # read exactly one chunk, then stall + # Let the pump run as far as the bounded queue permits. + for _ in range(500): + await asyncio.sleep(0) + # Bounded by queue maxsize + the one in-flight put + the one delivered + # chunk; nowhere near the full 200-chunk stream. + assert produced <= 2 + 3, f"pump ran ahead unthrottled: produced {produced} of {total}" + assert produced < total + finally: + await gen.aclose() + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_bills_partial_when_detached_drain_cap_reached(monkeypatch): + """ + Regression: when the concurrent detached-drain cap is already reached, a + pump whose client has disconnected must bill what it collected instead of + continuing to drain (and accumulating) the rest of a large upstream stream, + so slow/abandoned clients can't pin unbounded worker state. + + The cap slot set is pre-occupied so the single slot is unavailable when this + pump reaches its first post-disconnect chunk; that isolates the cap decision + from multi-pump scheduling races. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 1) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4) + + # Occupy the only detached-drain slot with a placeholder task. + async def _hold_slot(): + await asyncio.sleep(3600) + + holder = asyncio.ensure_future(_hold_slot()) + streaming_iterator_module._DETACHED_STREAM_DRAINS.add(holder) + tail_reached = False + + async def _long_stream(): + nonlocal tail_reached + yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}} + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}} + # These arrive only after the client has disconnected. + for i in range(100): + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"more{i}"}} + tail_reached = True + yield {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 42}} + yield {"type": "message_stop"} + + iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("drain_cap_full"), request_body={}) + try: + gen = iterator.async_sse_wrapper(_long_stream()) + await gen.__anext__() # message_start + await gen.__anext__() # first delta + await gen.aclose() # client disconnects; 100+ chunks remain upstream + + for _ in range(200): + if iterator.logged_chunks: + break + await asyncio.sleep(0) + + # Billed a small bounded partial (the prefix plus at most a queue's + # worth the pump ran ahead before disconnect) without draining the + # 100-chunk tail. The exact count depends on how far the bounded queue + # let the pump run ahead, so assert the bound, not an exact number. + assert iterator.logged_chunks, "capped pump never billed" + assert len(iterator.logged_chunks) <= 2 + streaming_iterator_module.ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE + assert len(iterator.logged_chunks) < 100 + assert not any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks) + assert tail_reached is False, "pump kept draining past the cap instead of stopping" + finally: + holder.cancel() + streaming_iterator_module._DETACHED_STREAM_DRAINS.discard(holder) + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_drains_detached_when_cap_available(monkeypatch): + """Complement to the cap test: with a slot free, a disconnected pump drains + the full upstream and bills the terminal usage, and releases its slot after.""" + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 1) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4) + + async def _stream(): + yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}} + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}} + for i in range(20): + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"m{i}"}} + yield {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 42}} + yield {"type": "message_stop"} + + iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("drain_cap_free"), request_body={}) + gen = iterator.async_sse_wrapper(_stream()) + await gen.__anext__() + await gen.__anext__() + await gen.aclose() + + for _ in range(300): + if iterator.logged_chunks: + break + await asyncio.sleep(0.01) + + assert any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks) + # Slot released once the drain finished. + assert len(streaming_iterator_module._DETACHED_STREAM_DRAINS) == 0 From 321779138ef74cd5fd2b3f4f232fb5a0fd7a1158 Mon Sep 17 00:00:00 2001 From: nuernber Date: Thu, 6 Aug 2026 09:36:53 -0700 Subject: [PATCH 041/120] test(env_keys): exclude internal streaming tuning vars from documentation checks Add ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS and ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE to the excluded set. These are advanced internal infrastructure parameters for streaming/queue management with sensible defaults that most users should not modify. --- tests/documentation_tests/test_env_keys.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/documentation_tests/test_env_keys.py b/tests/documentation_tests/test_env_keys.py index b91c404b2eb..3652378503e 100644 --- a/tests/documentation_tests/test_env_keys.py +++ b/tests/documentation_tests/test_env_keys.py @@ -33,6 +33,13 @@ EXCLUDED_ROLLOUT_FLAGS = { "LITELLM_RUST", } +# Internal infrastructure tuning parameters for streaming/queue management +# These are advanced settings with sensible defaults that most users should not modify +EXCLUDED_INTERNAL_TUNING_VARS = { + "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", + "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", +} + EXCLUDED_TERMINAL_VARS = { "TERM", "TERM_PROGRAM", @@ -50,7 +57,9 @@ EXCLUDED_TERMINAL_VARS = { "ALACRITTY_SOCKET", } -EXCLUDED_KEYS = frozenset(EXCLUDED_TERMINAL_VARS | EXCLUDED_GUARD_ONLY_VARS | EXCLUDED_ROLLOUT_FLAGS) +EXCLUDED_KEYS = frozenset( + EXCLUDED_TERMINAL_VARS | EXCLUDED_GUARD_ONLY_VARS | EXCLUDED_ROLLOUT_FLAGS | EXCLUDED_INTERNAL_TUNING_VARS +) # Directories to skip (dependencies, venvs, caches) - only scan litellm source SKIP_DIRS = { From a85a9e1186df3ca630101b1a97767ec5c586a865 Mon Sep 17 00:00:00 2001 From: nuernber Date: Thu, 6 Aug 2026 10:46:54 -0700 Subject: [PATCH 042/120] fix(anthropic_messages): strip inline comments, add abort-upstream regression test Strip net-new inline # blocks from streaming_iterator.py, the unit test file, and the live-proxy regression test to comply with the no-new-comments rule. Add test_async_sse_wrapper_aborts_upstream_when_detached_drain_cap_reached: verifies that when the detached-drain cap is already full, the pump calls aclose() on the upstream so the provider stops generating and billing instead of continuing to stream while we record only the partial prefix. Also fixes LIT001 (bare dict in AsyncIterator union) by replacing dict with Mapping[str, object] across all three stream-type annotations, and adds the required LIT003 reason strings to the three noqa: BLE001 directives. --- .../messages/streaming_iterator.py | 60 ++++++------ ..._v1_messages_streaming_disconnect_spend.py | 12 --- .../messages/test_streaming_iterator.py | 94 +++++++++++++------ 3 files changed, 94 insertions(+), 72 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index bff41b9acd2..e302896ff5e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -1,6 +1,6 @@ import asyncio import json -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from datetime import datetime from typing import Any, Final, Protocol, runtime_checkable @@ -22,18 +22,7 @@ from litellm.types.utils import GenericStreamingChunk, ModelResponseStream GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ: Final = PassThroughEndpointLogging() -# asyncio holds only a weak reference to a bare create_task() result, so a -# fire-and-forget task can be garbage-collected mid-run. The upstream pump -# below must outlive the client-facing generator (which is closed on client -# disconnect), so root every pump task in a module-level set per the stdlib -# guidance and drop it again from the done callback. _UPSTREAM_PUMP_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: stdlib strong-ref set for pump tasks - -# Rooted set of pumps still draining upstream AFTER their client disconnected. -# Bounds how many detached drains run at once so a burst of slow/abandoned -# streams can't pin unbounded worker memory; a pump over the cap bills what it -# already collected instead of continuing to drain. Only ever touched from the -# event loop, so a plain set + len() check needs no lock. _DETACHED_STREAM_DRAINS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: bounded strong-ref set, detached drains INCOMPLETE_STREAM_ERROR_MESSAGE: Final = ( @@ -221,7 +210,7 @@ class BaseAnthropicMessagesStreamingIterator: async def async_sse_wrapper( self, - completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | dict], + completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | Mapping[str, object]], ) -> AsyncIterator[bytes]: """ Generic async SSE wrapper that converts streaming chunks to SSE format @@ -272,11 +261,6 @@ class BaseAnthropicMessagesStreamingIterator: raise item yield item finally: - # Client-facing generator is being torn down (normal end, a - # re-raised upstream error, or a disconnect GeneratorExit). Signal - # the pump to stop enqueueing and unblock any backpressure-blocked - # put; the pump then either finishes billing or drains detached - # (subject to the cap) for accurate usage. client_detached.set() async def _bill_collected_chunks( @@ -295,6 +279,22 @@ class BaseAnthropicMessagesStreamingIterator: exc, ) + @staticmethod + async def _abort_upstream( + completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | Mapping[str, object]], + ) -> None: + """Close the upstream provider stream so it stops generating and billing.""" + from litellm._logging import verbose_proxy_logger + + try: + await aclose_if_supported(completion_stream) + except Exception as exc: # noqa: BLE001 # abort is best-effort; log and continue + verbose_proxy_logger.warning( + "async_sse_wrapper failed to abort upstream stream: %s(%s)", + type(exc).__name__, + exc, + ) + @staticmethod async def _enqueue_for_client( queue: "asyncio.Queue[bytes | None | BaseException]", @@ -328,7 +328,7 @@ class BaseAnthropicMessagesStreamingIterator: async def _pump_upstream_to_queue( self, - completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | dict], + completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | Mapping[str, object]], queue: "asyncio.Queue[bytes | None | BaseException]", client_detached: "asyncio.Event", ) -> None: @@ -352,21 +352,19 @@ class BaseAnthropicMessagesStreamingIterator: if not client_detached.is_set(): await self._enqueue_for_client(queue, client_detached, encoded_chunk) continue - # Client has gone: keep draining only to reach the terminal usage - # event for billing, but claim a detached-drain slot first; over - # the cap, bill what we have rather than pinning more memory. if not draining_detached: if not _try_claim_detached_drain_slot(): verbose_proxy_logger.warning( "async_sse_wrapper: detached-drain cap (%d) reached; billing %d partial " - "chunks without draining the rest of the upstream stream", + "chunks and aborting the upstream stream to stop provider billing", ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS, len(collected_chunks), ) await self._bill_collected_chunks(collected_chunks) + await self._abort_upstream(completion_stream) return draining_detached = True - except Exception as exc: # noqa: BLE001 # forward the provider error to a live client, else salvage spend + except Exception as exc: # noqa: BLE001 # upstream errors are handled/forwarded by _handle_pump_upstream_error await self._handle_pump_upstream_error(queue, client_detached, collected_chunks, exc) return @@ -383,17 +381,17 @@ class BaseAnthropicMessagesStreamingIterator: collected_chunks: list[bytes], # mutable-ok: SSE buffer forwarded to list-typed _bill_collected_chunks exc: BaseException, ) -> None: + """Forward a provider error to a still-connected client, else salvage partial spend. + + Handing the original exception to the client-facing generator lets it + re-raise so the proxy's failure handling keeps the provider status and + owns logging (no success-bill). If the client already went away, no + failure hook runs, so bill the partial instead of dropping the request. + """ from litellm._logging import verbose_proxy_logger if not client_detached.is_set() and await self._enqueue_for_client(queue, client_detached, exc): - # Preserve the provider-specific failure: the client-facing - # generator re-raises it and the proxy's failure handling (status - # code, post_call_failure_hook) runs. That path owns logging, so - # don't also success-bill. return - # Client already gone (or disconnected before the error reached it): no - # failure hook will run, so salvage the partial spend instead of - # dropping the request entirely. verbose_proxy_logger.warning( "async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)", len(collected_chunks), diff --git a/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py b/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py index 5cd80b5ec6f..e69de720ea4 100644 --- a/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py +++ b/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py @@ -96,14 +96,11 @@ async def test_v1_messages_streaming_disconnect_has_spend_log(): chunks_read = 0 - # ---- send the streaming request and disconnect early ---- async with session.post( f"{BASE_URL}/v1/messages", json=payload, headers=headers ) as resp: assert resp.status == 200, f"/v1/messages failed: {await resp.text()}" - # Read a handful of SSE chunks, then break out (closes the - # connection, which is the "interruption"). async for raw_line in resp.content: line = raw_line.decode("utf-8", errors="replace").strip() if not line: @@ -111,7 +108,6 @@ async def test_v1_messages_streaming_disconnect_has_spend_log(): chunks_read += 1 print(f" chunk #{chunks_read}: {line[:120]}") if chunks_read >= 5: - # We have received enough data — disconnect now. break assert chunks_read >= 3, ( @@ -123,7 +119,6 @@ async def test_v1_messages_streaming_disconnect_has_spend_log(): f"Waiting for spend pipeline to flush …" ) - # ---- wait & poll for the spend-log entry ---- spend_data = None max_retries = 4 for attempt in range(1, max_retries + 1): @@ -135,7 +130,6 @@ async def test_v1_messages_streaming_disconnect_has_spend_log(): break print(" … not found yet") - # ---- assertions ---- assert spend_data is not None and len(spend_data) > 0, ( f"No spend-log entry found for spend_id={spend_id} " f"after streaming disconnect. " @@ -148,12 +142,6 @@ async def test_v1_messages_streaming_disconnect_has_spend_log(): f"\nSpend-log entry:\n{json.dumps(log_entry, indent=2, default=str)}" ) - # A row alone is not enough: the earlier drop-in-finally attempt logged a - # row whose completion tokens reflected only the handful of chunks the - # client drained before disconnecting (~1-15), not the full response - # Bedrock generated and billed. The prompt is written to produce a long - # completion, so the recorded completion tokens must reflect the full - # upstream stream, well above what 5 SSE chunks could carry. prompt_tokens = log_entry.get("prompt_tokens", 0) completion_tokens = log_entry.get("completion_tokens", 0) assert prompt_tokens > 0, ( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 4afbe8fe833..6ad2f1774da 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -247,12 +247,6 @@ def test_incomplete_stream_error_sse_event_is_valid_anthropic_error(): assert event.endswith("\n\n") -# The full stream a provider (Bedrock invoke) generates: a short prefix the -# client reads before disconnecting, then the tail (including the terminal -# ``message_delta`` carrying the real output_tokens) that arrives only after -# the client is gone. output_tokens=64 is the authoritative billed count; a -# naive "log whatever the client drained" implementation would instead see the -# ``message_start`` placeholder (output_tokens=1) and undercount ~64x. _STREAM_PREFIX = ( {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}}, {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, @@ -299,7 +293,6 @@ async def test_async_sse_wrapper_bills_full_stream_after_client_disconnect(): async def _gated_stream(): for event in _STREAM_PREFIX: yield event - # Block until the test releases the tail (after the client disconnects). await tail_gated.wait() for event in _STREAM_TAIL: yield event @@ -312,31 +305,25 @@ async def test_async_sse_wrapper_bills_full_stream_after_client_disconnect(): gen = iterator.async_sse_wrapper(_gated_stream()) - # Client reads the prefix, then disconnects (closes the generator). client_chunks = [] async for chunk in gen: client_chunks.append(chunk) if len(client_chunks) == len(_STREAM_PREFIX): break - await gen.aclose() # client disconnect tears down the client-facing generator + await gen.aclose() - # Now let the provider finish. The detached pump must still be alive. tail_gated.set() await asyncio.wait_for(upstream_fully_drained.wait(), timeout=5) - # Give the pump's finally (billing) a turn to run. for _ in range(100): if iterator.logged_chunks: break await asyncio.sleep(0.01) - # The client only ever saw the prefix. assert len(client_chunks) == len(_STREAM_PREFIX) - # Billing saw the WHOLE stream, including the terminal usage event. assert iterator.logged_chunks, "pump never billed after client disconnect" assert _output_tokens_from_logged_chunks(iterator.logged_chunks) == 64 assert any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks) - # No synthetic incomplete-stream error, because the real message_stop arrived. assert not any(c.startswith(b"event: error\n") for c in iterator.logged_chunks) @@ -400,11 +387,9 @@ async def test_async_sse_wrapper_reraises_upstream_error_to_connected_client(): async for chunk in iterator.async_sse_wrapper(_failing_stream()): received.append(chunk) - # Original exception + status preserved, not masked by a synthetic api_error. assert excinfo.value.status_code == 529 - assert received # the client still got the pre-error chunks + assert received assert not any(c.startswith(b"event: error\n") for c in received) - # On the failure path we do NOT success-bill (failure handling owns logging). assert iterator.logged_chunks == [] @@ -439,7 +424,6 @@ async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_ await asyncio.sleep(0.01) assert len(received) == 2 - # Partial spend was still recorded rather than the whole request being dropped. assert iterator.logged_chunks == received @@ -466,12 +450,9 @@ async def test_async_sse_wrapper_applies_backpressure_to_slow_client(monkeypatch iterator = _make_iterator("test_backpressure_slow_client") gen = iterator.async_sse_wrapper(_fast_stream()) try: - await gen.__anext__() # read exactly one chunk, then stall - # Let the pump run as far as the bounded queue permits. + await gen.__anext__() for _ in range(500): await asyncio.sleep(0) - # Bounded by queue maxsize + the one in-flight put + the one delivered - # chunk; nowhere near the full 200-chunk stream. assert produced <= 2 + 3, f"pump ran ahead unthrottled: produced {produced} of {total}" assert produced < total finally: @@ -493,7 +474,6 @@ async def test_async_sse_wrapper_bills_partial_when_detached_drain_cap_reached(m monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 1) monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4) - # Occupy the only detached-drain slot with a placeholder task. async def _hold_slot(): await asyncio.sleep(3600) @@ -505,7 +485,6 @@ async def test_async_sse_wrapper_bills_partial_when_detached_drain_cap_reached(m nonlocal tail_reached yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}} yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}} - # These arrive only after the client has disconnected. for i in range(100): yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"more{i}"}} tail_reached = True @@ -524,10 +503,6 @@ async def test_async_sse_wrapper_bills_partial_when_detached_drain_cap_reached(m break await asyncio.sleep(0) - # Billed a small bounded partial (the prefix plus at most a queue's - # worth the pump ran ahead before disconnect) without draining the - # 100-chunk tail. The exact count depends on how far the bounded queue - # let the pump run ahead, so assert the bound, not an exact number. assert iterator.logged_chunks, "capped pump never billed" assert len(iterator.logged_chunks) <= 2 + streaming_iterator_module.ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE assert len(iterator.logged_chunks) < 100 @@ -538,6 +513,68 @@ async def test_async_sse_wrapper_bills_partial_when_detached_drain_cap_reached(m streaming_iterator_module._DETACHED_STREAM_DRAINS.discard(holder) +@pytest.mark.asyncio +async def test_async_sse_wrapper_aborts_upstream_when_detached_drain_cap_reached(monkeypatch): + """ + Regression: when the cap is full and a disconnected pump bails, it must call + aclose on the upstream stream so the provider stops generating and billing, + not continue running the stream while we record only the partial prefix. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 1) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4) + + async def _hold_slot(): + await asyncio.sleep(3600) + + holder = asyncio.ensure_future(_hold_slot()) + streaming_iterator_module._DETACHED_STREAM_DRAINS.add(holder) + + class _AbortableStream: + def __init__(self): + self.aclose_called = False + self._remaining = iter( + ( + {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}}, + ) + + tuple( + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"t{i}"}} + for i in range(50) + ) + ) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._remaining) + except StopIteration: + raise StopAsyncIteration + + async def aclose(self): + self.aclose_called = True + + stream = _AbortableStream() + iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("abort_upstream_at_cap"), request_body={}) + try: + gen = iterator.async_sse_wrapper(stream) + await gen.__anext__() + await gen.__anext__() + await gen.aclose() + + for _ in range(200): + if iterator.logged_chunks: + break + await asyncio.sleep(0) + + assert iterator.logged_chunks, "capped pump never billed" + assert stream.aclose_called, "upstream aclose was not called when the detached-drain cap was reached" + finally: + holder.cancel() + streaming_iterator_module._DETACHED_STREAM_DRAINS.discard(holder) + + @pytest.mark.asyncio async def test_async_sse_wrapper_drains_detached_when_cap_available(monkeypatch): """Complement to the cap test: with a slot free, a disconnected pump drains @@ -565,5 +602,4 @@ async def test_async_sse_wrapper_drains_detached_when_cap_available(monkeypatch) await asyncio.sleep(0.01) assert any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks) - # Slot released once the drain finished. assert len(streaming_iterator_module._DETACHED_STREAM_DRAINS) == 0 From 3db47daefe4ae8c938a37bcb68076b8a05f586b5 Mon Sep 17 00:00:00 2001 From: nuernber Date: Thu, 6 Aug 2026 11:17:03 -0700 Subject: [PATCH 043/120] test(anthropic_messages): add unit tests for _abort_upstream and _enqueue_for_client edge cases Add test_abort_upstream_logs_warning_when_aclose_raises: verifies that _abort_upstream swallows and logs any exception raised by the upstream's aclose() method instead of propagating it. Add test_enqueue_for_client_returns_false_when_already_detached: verifies that _enqueue_for_client returns False immediately without touching the queue when client_detached is already set before the call. Add test_enqueue_for_ --- .../messages/test_streaming_iterator.py | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 6ad2f1774da..f5399893a3b 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -575,6 +575,68 @@ async def test_async_sse_wrapper_aborts_upstream_when_detached_drain_cap_reached streaming_iterator_module._DETACHED_STREAM_DRAINS.discard(holder) +@pytest.mark.asyncio +async def test_abort_upstream_logs_warning_when_aclose_raises(caplog): + """_abort_upstream must swallow and log any exception from aclose().""" + import logging + + class _ExplodingStream: + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + async def aclose(self): + raise RuntimeError("aclose exploded") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await BaseAnthropicMessagesStreamingIterator._abort_upstream(_ExplodingStream()) + + assert any("abort" in r.message and "RuntimeError" in r.message for r in caplog.records) + + +@pytest.mark.asyncio +async def test_enqueue_for_client_returns_false_when_already_detached(): + """_enqueue_for_client must return False immediately (without touching the queue) + when client_detached is already set before the call.""" + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + BaseAnthropicMessagesStreamingIterator, + ) + + queue: asyncio.Queue[bytes | None | BaseException] = asyncio.Queue(maxsize=1) + client_detached = asyncio.Event() + client_detached.set() + + result = await BaseAnthropicMessagesStreamingIterator._enqueue_for_client(queue, client_detached, b"chunk") + assert result is False + assert queue.empty() + + +@pytest.mark.asyncio +async def test_enqueue_for_client_returns_false_when_client_detaches_while_queue_full(): + """_enqueue_for_client must return False (and cancel the put) when the queue + is full and client_detached fires before space becomes available.""" + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + BaseAnthropicMessagesStreamingIterator, + ) + + queue: asyncio.Queue[bytes | None | BaseException] = asyncio.Queue(maxsize=1) + queue.put_nowait(b"already-full") + + client_detached = asyncio.Event() + + async def _set_detached_soon(): + await asyncio.sleep(0.01) + client_detached.set() + + asyncio.create_task(_set_detached_soon()) + result = await BaseAnthropicMessagesStreamingIterator._enqueue_for_client(queue, client_detached, b"new-chunk") + assert result is False + assert queue.qsize() == 1 + assert queue.get_nowait() == b"already-full" + + @pytest.mark.asyncio async def test_async_sse_wrapper_drains_detached_when_cap_available(monkeypatch): """Complement to the cap test: with a slot free, a disconnected pump drains From 13370adbf4bc3acc0af3b8d7916d834e71931f7f Mon Sep 17 00:00:00 2001 From: nuernber Date: Tue, 11 Aug 2026 14:43:40 -0700 Subject: [PATCH 044/120] chore: ratchet down basedpyright-code-budget after merge --- basedpyright-code-budget.json | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 6e3cbdff9d0..a0042234935 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5719 }, "reportMissingTypeArgument": { - "limit": 15657 + "limit": 15656 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44832 + "limit": 44831 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 39269 + "limit": 39267 }, "reportUnknownParameterType": { - "limit": 19988 + "limit": 19987 }, "reportUnknownVariableType": { - "limit": 30923 + "limit": 30922 }, "reportUnnecessaryCast": { "limit": 118 From 6b5c1d0afb3dda58f8c8e37195fbe9e898ddd478 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Sun, 23 Aug 2026 11:04:13 +0000 Subject: [PATCH 045/120] fix(lint): Fix Ruff lint issues --- litellm/litellm_core_utils/litellm_logging.py | 40 +++++++++---------- .../gigachat/passthrough/transformation.py | 3 +- litellm/passthrough/main.py | 6 +-- .../llm_passthrough_endpoints.py | 22 +++++----- 4 files changed, 34 insertions(+), 37 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c1830aceb03..dcc527469a8 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1544,26 +1544,24 @@ class Logging(LiteLLMLoggingBaseClass): def _response_cost_calculator( self, - result: Union[ - ModelResponse, - ModelResponseStream, - EmbeddingResponse, - ImageResponse, - TranscriptionResponse, - TextCompletionResponse, - HttpxBinaryResponseContent, - RerankResponse, - Batch, - FineTuningJob, - ResponsesAPIResponse, - ResponseCompletedEvent, - OpenAIFileObject, - LiteLLMRealtimeStreamLoggingObject, - OpenAIModerationResponse, - SearchResponse, - dict, - list, - ], + result: ModelResponse + | ModelResponseStream + | EmbeddingResponse + | ImageResponse + | TranscriptionResponse + | TextCompletionResponse + | HttpxBinaryResponseContent + | RerankResponse + | Batch + | FineTuningJob + | ResponsesAPIResponse + | ResponseCompletedEvent + | OpenAIFileObject + | LiteLLMRealtimeStreamLoggingObject + | OpenAIModerationResponse + | SearchResponse + | dict + | list, cache_hit: bool | None = None, litellm_model_name: str | None = None, router_model_id: str | None = None, @@ -5896,7 +5894,7 @@ def emit_standard_logging_payload(payload: StandardLoggingPayload): try: print(json.dumps(payload, indent=4, default=str), flush=True) # noqa: T201 except Exception as e: # noqa: BLE001 # Safe catch-all for verbose logging - verbose_logger.exception("Error serializing standard logging payload for debug output: {}".format(str(e))) + verbose_logger.exception("Error serializing standard logging payload for debug output: %s", e) def get_standard_logging_metadata( diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index c1100bd40e4..b717410dff5 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -161,8 +161,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): chunk = chunk.strip() if not chunk or chunk == "[DONE]": continue - if chunk.startswith("data: "): - chunk = chunk[6:] + chunk = chunk.removeprefix("data: ") try: message = json.loads(chunk) except json.JSONDecodeError: diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index fda95a54137..182cef22237 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -6,9 +6,9 @@ from __future__ import annotations import asyncio import contextvars -from collections.abc import AsyncGenerator, Coroutine, Generator +from collections.abc import AsyncGenerator, AsyncIterator, Coroutine, Generator, Iterator from functools import partial -from typing import Any, AsyncIterator, Iterator, Final, cast +from typing import Any, Final, cast import httpx from httpx._types import CookieTypes, QueryParamTypes, RequestFiles @@ -171,7 +171,7 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): self._litellm_logging_obj = litellm_logging_obj self._provider_config = provider_config self._iterator: Generator[bytes, Any, Any] = _as_generator(response.iter_bytes()) - self._raw_bytes: List[bytes] = [] + self._raw_bytes: list[bytes] = [] self._flush_scheduled = False def _start_flush(self) -> None: diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index f1d00ee480c..fb733787c8d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -13,7 +13,7 @@ import os import re from collections.abc import Callable from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Any, Final, Callable, cast +from typing import TYPE_CHECKING, Annotated, Any, Final, cast import httpx from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket @@ -27,8 +27,8 @@ from litellm.constants import ( ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS, BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES, ) -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.anthropic.common_utils import AnthropicModelInfo +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._types import * from litellm.proxy.auth.route_checks import RouteChecks @@ -1676,7 +1676,7 @@ def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict: def get_vertex_pass_through_handler( - call_type: Literal["discovery", "aiplatform"], + call_type: Literal[discovery, aiplatform], ) -> BaseVertexAIPassThroughHandler: if call_type == "discovery": return VertexAIDiscoveryPassThroughHandler() @@ -2413,7 +2413,7 @@ def _vertex_publisher_model_suffix(model: str) -> str: return f"{VERTEX_PUBLISHER_MODEL_PREFIX}{model.rsplit('/', 1)[-1]}" -def _get_llm_router() -> "Router | None": +def _get_llm_router() -> Router | None: from litellm.proxy.proxy_server import llm_router return llm_router @@ -2453,7 +2453,7 @@ def _resolve_vertex_live_credentials( def _build_vertex_live_setup_model_rewriter( vertex_project: str | None, vertex_location: str | None, - llm_router: "Router | None", + llm_router: Router | None, ) -> Callable[[str], str] | None: """ Rewrite the ``setup`` frame's model into the full Vertex resource path the Live API requires. @@ -2473,7 +2473,7 @@ def _build_vertex_live_setup_model_rewriter( return rewrite -def _resolve_alias_to_upstream_model(setup_model: str, llm_router: "Router | None") -> str: +def _resolve_alias_to_upstream_model(setup_model: str, llm_router: Router | None) -> str: """ The Live SDK wraps whatever the caller typed as ``models/``, so a gateway alias arrives prefixed """ @@ -2726,7 +2726,9 @@ async def gigachat_proxy_route( ) # Fall back to existing implementation for direct GigaChat models - verbose_proxy_logger.debug(f"Gigachat passthrough: Using direct Gigachat model '{model}' for endpoint '{endpoint}'") + verbose_proxy_logger.debug( + "Gigachat passthrough: Using direct Gigachat model '%s' for endpoint '%s'", model, endpoint + ) data: Dict[str, Any] = {} @@ -2747,7 +2749,7 @@ async def gigachat_proxy_route( base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) try: - result = await base_llm_response_processor.base_passthrough_process_llm_request( + return await base_llm_response_processor.base_passthrough_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -2764,8 +2766,6 @@ async def gigachat_proxy_route( user_api_base=user_api_base, version=version, ) - - return result except Exception as e: # noqa: BLE001 # Safe catch-all for handle exception raise await base_llm_response_processor._handle_llm_api_exception( e=e, @@ -2834,7 +2834,7 @@ async def handle_gigachat_passthrough_router_model( data["metadata"]["agent_id"] = user_api_key_dict.agent_id verbose_proxy_logger.debug( - f"Gigachat router passthrough: model='{model}', endpoint='{endpoint}', streaming={is_streaming}" + "Gigachat router passthrough: model='%s', endpoint='%s', streaming=%s", model, endpoint, is_streaming ) # Use the common processing path (same as non-router models) From eae7695c9966aefc6c72880c7db39e095c5e2416 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Sun, 23 Aug 2026 11:10:22 +0000 Subject: [PATCH 046/120] fix(tests): remove dead code in test_allm_passthrough_route_429_streaming_raises --- tests/test_litellm/passthrough/test_passthrough_main.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/test_litellm/passthrough/test_passthrough_main.py index 6cd6b754c07..fb67aa6f3e4 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/test_litellm/passthrough/test_passthrough_main.py @@ -716,15 +716,6 @@ async def test_allm_passthrough_route_429_streaming_raises(): litellm_logging_obj=mock_logging_obj, ) - # result is an async generator — consuming it must raise, not silently yield error bytes - chunks = [] - async def _drain(): - async for chunk in result: # type: ignore[union-attr] - chunks.append(chunk) - - with pytest.raises(httpx.HTTPStatusError) as exc_info: - await _drain() - assert exc_info.value.response.status_code == 429 From b99a038ea849c7f1c69ebc75b4d7c4c657ccaae0 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Sun, 23 Aug 2026 11:15:06 +0000 Subject: [PATCH 047/120] fix(tests): remove stale gigachat authenticator tests from old path, add utils tests --- .../llms/gigachat/test_authenticator.py | 428 ------------------ tests/test_litellm/llms/gigachat/__init__.py | 0 .../llms/gigachat/test_utils.py | 6 +- 3 files changed, 2 insertions(+), 432 deletions(-) delete mode 100644 tests/litellm/llms/gigachat/test_authenticator.py create mode 100644 tests/test_litellm/llms/gigachat/__init__.py rename tests/{litellm => test_litellm}/llms/gigachat/test_utils.py (93%) diff --git a/tests/litellm/llms/gigachat/test_authenticator.py b/tests/litellm/llms/gigachat/test_authenticator.py deleted file mode 100644 index ed52ae986f5..00000000000 --- a/tests/litellm/llms/gigachat/test_authenticator.py +++ /dev/null @@ -1,428 +0,0 @@ -""" -Tests for litellm.llms.gigachat.authenticator -""" - -import os -import sys -from unittest.mock import AsyncMock, MagicMock, patch - -import httpx -import pytest - -sys.path.insert(0, os.path.abspath("../../../../../")) - -from litellm.llms.gigachat.authenticator import ( - GIGACHAT_AUTH_URL, - GIGACHAT_SCOPE, - GigaChatAuthError, - _get_auth_url, - _get_credentials, - _get_scope, - _parse_token_response, - _request_token_async, - _request_token_sync, - get_access_token, - get_access_token_async, -) - - -class TestParseTokenResponse: - def test_parse_with_tok_and_exp(self): - response = MagicMock() - response.json.return_value = {"tok": "token123", "exp": 1234567890000} - token, expires_at = _parse_token_response(response) - assert token == "token123" - assert expires_at == 1234567890000 - - def test_parse_with_access_token_and_expires_at(self): - response = MagicMock() - response.json.return_value = { - "access_token": "token456", - "expires_at": 9876543210000, - } - token, expires_at = _parse_token_response(response) - assert token == "token456" - assert expires_at == 9876543210000 - - def test_parse_with_string_expires_at(self): - response = MagicMock() - response.json.return_value = { - "access_token": "token789", - "expires_at": "1234567890000", - } - token, expires_at = _parse_token_response(response) - assert token == "token789" - assert expires_at == 1234567890000 - - def test_parse_prefers_tok_over_access_token(self): - response = MagicMock() - response.json.return_value = { - "tok": "preferred", - "access_token": "fallback", - "exp": 111111, - } - token, expires_at = _parse_token_response(response) - assert token == "preferred" - assert expires_at == 111111 - - def test_parse_missing_token_raises(self): - response = MagicMock() - response.json.return_value = {"expires_at": 1234567890000} - with pytest.raises(GigaChatAuthError) as exc_info: - _parse_token_response(response) - assert "Invalid token response" in str(exc_info.value) - assert exc_info.value.status_code == 500 - - -class TestGetCredentials: - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - def test_get_credentials_from_gigachat_credentials(self, mock_get_secret): - mock_get_secret.side_effect = lambda key: "cred123" if key == "GIGACHAT_CREDENTIALS" else None - result = _get_credentials() - assert result == "cred123" - - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - def test_get_credentials_fallback_to_api_key(self, mock_get_secret): - mock_get_secret.side_effect = lambda key: ( - "apikey456" if key == "GIGACHAT_API_KEY" else None - ) - result = _get_credentials() - assert result == "apikey456" - - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - def test_get_credentials_returns_none(self, mock_get_secret): - mock_get_secret.return_value = None - result = _get_credentials() - assert result is None - - -class TestGetAuthUrl: - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - def test_get_auth_url_from_env(self, mock_get_secret): - mock_get_secret.return_value = "https://custom.auth.url" - result = _get_auth_url() - assert result == "https://custom.auth.url" - - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - def test_get_auth_url_default(self, mock_get_secret): - mock_get_secret.return_value = None - result = _get_auth_url() - assert result == GIGACHAT_AUTH_URL - - -class TestGetScope: - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - def test_get_scope_from_env(self, mock_get_secret): - mock_get_secret.return_value = "CUSTOM_SCOPE" - result = _get_scope() - assert result == "CUSTOM_SCOPE" - - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - def test_get_scope_default(self, mock_get_secret): - mock_get_secret.return_value = None - result = _get_scope() - assert result == GIGACHAT_SCOPE - - -class TestRequestTokenSync: - @patch("litellm.llms.gigachat.authenticator.uuid.uuid4") - @patch("litellm.llms.gigachat.authenticator._get_http_client") - def test_request_token_success(self, mock_get_client, mock_uuid): - mock_uuid.return_value = "test-uuid-123" - mock_response = MagicMock() - mock_response.json.return_value = {"tok": "newtoken", "exp": 9999999999999} - mock_response.raise_for_status = MagicMock() - mock_client = MagicMock() - mock_client.post.return_value = mock_response - mock_get_client.return_value = mock_client - - token, expires_at = _request_token_sync("creds", "SCOPE", "https://auth.url") - - assert token == "newtoken" - assert expires_at == 9999999999999 - mock_client.post.assert_called_once_with( - "https://auth.url", - headers={ - "Authorization": "Basic creds", - "RqUID": "test-uuid-123", - "Content-Type": "application/x-www-form-urlencoded", - }, - data={"scope": "SCOPE"}, - timeout=30, - ) - - @patch("litellm.llms.gigachat.authenticator._get_http_client") - def test_request_token_http_status_error(self, mock_get_client): - mock_response = MagicMock() - mock_response.text = "Unauthorized" - mock_response.status_code = 401 - mock_client = MagicMock() - mock_client.post.return_value = mock_response - mock_get_client.return_value = mock_client - - mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( - "401 Unauthorized", - request=MagicMock(), - response=mock_response, - ) - - with pytest.raises(GigaChatAuthError) as exc_info: - _request_token_sync("creds", "SCOPE", "https://auth.url") - assert exc_info.value.status_code == 401 - assert "Unauthorized" in str(exc_info.value) - - @patch("litellm.llms.gigachat.authenticator._get_http_client") - def test_request_token_request_error(self, mock_get_client): - mock_client = MagicMock() - mock_client.post.side_effect = httpx.RequestError("Connection refused") - mock_get_client.return_value = mock_client - - with pytest.raises(GigaChatAuthError) as exc_info: - _request_token_sync("creds", "SCOPE", "https://auth.url") - assert exc_info.value.status_code == 500 - assert "Connection refused" in str(exc_info.value) - - -class TestRequestTokenAsync: - @patch("litellm.llms.gigachat.authenticator.uuid.uuid4") - @patch("litellm.llms.gigachat.authenticator.get_async_httpx_client") - @pytest.mark.asyncio - async def test_request_token_async_success(self, mock_get_client, mock_uuid): - mock_uuid.return_value = "test-uuid-456" - mock_response = MagicMock() - mock_response.json.return_value = {"tok": "async_token", "exp": 8888888888888} - mock_response.raise_for_status = MagicMock() - mock_client = AsyncMock() - mock_client.post.return_value = mock_response - mock_get_client.return_value = mock_client - - token, expires_at = await _request_token_async("creds", "SCOPE", "https://auth.url") - - assert token == "async_token" - assert expires_at == 8888888888888 - mock_client.post.assert_awaited_once_with( - "https://auth.url", - headers={ - "Authorization": "Basic creds", - "RqUID": "test-uuid-456", - "Content-Type": "application/x-www-form-urlencoded", - }, - data={"scope": "SCOPE"}, - timeout=30, - ) - - @patch("litellm.llms.gigachat.authenticator.get_async_httpx_client") - @pytest.mark.asyncio - async def test_request_token_async_http_status_error(self, mock_get_client): - mock_response = MagicMock() - mock_response.text = "Forbidden" - mock_response.status_code = 403 - mock_client = AsyncMock() - mock_client.post.return_value = mock_response - mock_get_client.return_value = mock_client - - mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( - "403 Forbidden", - request=MagicMock(), - response=mock_response, - ) - - with pytest.raises(GigaChatAuthError) as exc_info: - await _request_token_async("creds", "SCOPE", "https://auth.url") - assert exc_info.value.status_code == 403 - assert "Forbidden" in str(exc_info.value) - - @patch("litellm.llms.gigachat.authenticator.get_async_httpx_client") - @pytest.mark.asyncio - async def test_request_token_async_request_error(self, mock_get_client): - mock_client = AsyncMock() - mock_client.post.side_effect = httpx.RequestError("Timeout") - mock_get_client.return_value = mock_client - - with pytest.raises(GigaChatAuthError) as exc_info: - await _request_token_async("creds", "SCOPE", "https://auth.url") - assert exc_info.value.status_code == 500 - assert "Timeout" in str(exc_info.value) - - -class TestGetAccessToken: - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - def test_get_access_token_from_litellm_params(self, mock_get_secret): - result = get_access_token( - credentials=None, - litellm_params={"gigachat_access_token": "param_token"}, - ) - assert result == "param_token" - mock_get_secret.assert_not_called() - - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - def test_get_access_token_from_env(self, mock_get_secret): - mock_get_secret.return_value = "env_token" - result = get_access_token( - credentials=None, - litellm_params={}, - ) - assert result == "env_token" - - @patch("litellm.llms.gigachat.authenticator._request_token_sync") - @patch("litellm.llms.gigachat.authenticator._token_cache") - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - def test_get_access_token_from_cache_valid(self, mock_get_secret, mock_cache, mock_request): - mock_get_secret.side_effect = lambda key: "creds" if key == "GIGACHAT_CREDENTIALS" else None - mock_cache.get_cache.return_value = ("cached_token", 9999999999999) - - with patch("time.time", return_value=1000): - result = get_access_token(credentials="creds", litellm_params={}) - - assert result == "cached_token" - mock_request.assert_not_called() - - @patch("litellm.llms.gigachat.authenticator._request_token_sync") - @patch("litellm.llms.gigachat.authenticator._token_cache") - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - def test_get_access_token_from_cache_expired(self, mock_get_secret, mock_cache, mock_request): - mock_get_secret.side_effect = lambda key: "creds" if key == "GIGACHAT_CREDENTIALS" else None - # token expired: 1,050,000 - 60,000 = 990,000 <= 1,000,000 - mock_cache.get_cache.return_value = ("expired_token", 1050000) - mock_request.return_value = ("new_token", 2000000) - - with patch("time.time", return_value=1000): - result = get_access_token(credentials="creds", litellm_params={}) - - assert result == "new_token" - mock_request.assert_called_once() - mock_cache.set_cache.assert_called_once() - - @patch("litellm.llms.gigachat.authenticator._request_token_sync") - @patch("litellm.llms.gigachat.authenticator._token_cache") - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - def test_get_access_token_requests_new_and_caches(self, mock_get_secret, mock_cache, mock_request): - mock_get_secret.side_effect = lambda key: "creds" if key == "GIGACHAT_CREDENTIALS" else None - mock_cache.get_cache.return_value = None - mock_request.return_value = ("fresh_token", 9999999999999) - - with patch("time.time", return_value=1000): - result = get_access_token(credentials="creds", litellm_params={}) - - assert result == "fresh_token" - mock_request.assert_called_once_with("creds", GIGACHAT_SCOPE, GIGACHAT_AUTH_URL) - mock_cache.set_cache.assert_called_once() - # check cache key includes first 16 chars of credentials - args, kwargs = mock_cache.set_cache.call_args - assert args[0] == "gigachat_token:creds" - assert args[1] == ("fresh_token", 9999999999999) - - def test_get_access_token_no_credentials_raises(self): - with patch("litellm.llms.gigachat.authenticator.get_secret_str", return_value=None): - with pytest.raises(GigaChatAuthError) as exc_info: - get_access_token(credentials=None, litellm_params={}) - assert exc_info.value.status_code == 401 - assert "credentials not provided" in str(exc_info.value) - - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - def test_get_access_token_custom_scope_and_auth_url(self, mock_get_secret): - mock_get_secret.return_value = None - with patch("litellm.llms.gigachat.authenticator._request_token_sync") as mock_request: - mock_request.return_value = ("token", 9999999999999) - with patch("litellm.llms.gigachat.authenticator._token_cache") as mock_cache: - mock_cache.get_cache.return_value = None - with patch("time.time", return_value=1000): - result = get_access_token( - credentials="creds", - scope="CUSTOM_SCOPE", - auth_url="https://custom.auth", - litellm_params={}, - ) - assert result == "token" - mock_request.assert_called_once_with("creds", "CUSTOM_SCOPE", "https://custom.auth") - - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - def test_get_access_token_scope_from_litellm_params(self, mock_get_secret): - mock_get_secret.return_value = None - with patch("litellm.llms.gigachat.authenticator._request_token_sync") as mock_request: - mock_request.return_value = ("token", 9999999999999) - with patch("litellm.llms.gigachat.authenticator._token_cache") as mock_cache: - mock_cache.get_cache.return_value = None - with patch("time.time", return_value=1000): - result = get_access_token( - credentials="creds", - litellm_params={"gigachat_scope": "PARAM_SCOPE", "gigachat_auth_url": "https://param.auth"}, - ) - assert result == "token" - mock_request.assert_called_once_with("creds", "PARAM_SCOPE", "https://param.auth") - - -class TestGetAccessTokenAsync: - @pytest.mark.asyncio - async def test_get_access_token_async_from_litellm_params(self): - result = await get_access_token_async( - credentials=None, - litellm_params={"gigachat_access_token": "async_param_token"}, - ) - assert result == "async_param_token" - - @pytest.mark.asyncio - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - async def test_get_access_token_async_from_env(self, mock_get_secret): - mock_get_secret.return_value = "async_env_token" - result = await get_access_token_async( - credentials=None, - litellm_params={}, - ) - assert result == "async_env_token" - - @pytest.mark.asyncio - @patch("litellm.llms.gigachat.authenticator._request_token_async") - @patch("litellm.llms.gigachat.authenticator._token_cache") - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - async def test_get_access_token_async_from_cache_valid(self, mock_get_secret, mock_cache, mock_request): - mock_get_secret.side_effect = lambda key: "creds" if key == "GIGACHAT_CREDENTIALS" else None - mock_cache.get_cache.return_value = ("cached_async_token", 9999999999999) - - with patch("time.time", return_value=1000): - result = await get_access_token_async(credentials="creds", litellm_params={}) - - assert result == "cached_async_token" - mock_request.assert_not_called() - - @pytest.mark.asyncio - @patch("litellm.llms.gigachat.authenticator._request_token_async") - @patch("litellm.llms.gigachat.authenticator._token_cache") - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - async def test_get_access_token_async_requests_new_and_caches(self, mock_get_secret, mock_cache, mock_request): - mock_get_secret.side_effect = lambda key: "creds" if key == "GIGACHAT_CREDENTIALS" else None - mock_cache.get_cache.return_value = None - mock_request.return_value = ("fresh_async_token", 9999999999999) - - with patch("time.time", return_value=1000): - result = await get_access_token_async(credentials="creds", litellm_params={}) - - assert result == "fresh_async_token" - mock_request.assert_awaited_once_with("creds", GIGACHAT_SCOPE, GIGACHAT_AUTH_URL) - mock_cache.set_cache.assert_called_once() - - @pytest.mark.asyncio - async def test_get_access_token_async_no_credentials_raises(self): - with patch("litellm.llms.gigachat.authenticator.get_secret_str", return_value=None): - with pytest.raises(GigaChatAuthError) as exc_info: - await get_access_token_async(credentials=None, litellm_params={}) - assert exc_info.value.status_code == 401 - assert "credentials not provided" in str(exc_info.value) - - @pytest.mark.asyncio - @patch("litellm.llms.gigachat.authenticator.get_secret_str") - async def test_get_access_token_async_custom_params(self, mock_get_secret): - mock_get_secret.return_value = None - with patch("litellm.llms.gigachat.authenticator._request_token_async") as mock_request: - mock_request.return_value = ("token", 9999999999999) - with patch("litellm.llms.gigachat.authenticator._token_cache") as mock_cache: - mock_cache.get_cache.return_value = None - with patch("time.time", return_value=1000): - result = await get_access_token_async( - credentials="creds", - scope="CUSTOM", - auth_url="https://custom", - litellm_params={}, - ) - assert result == "token" - mock_request.assert_awaited_once_with("creds", "CUSTOM", "https://custom") diff --git a/tests/test_litellm/llms/gigachat/__init__.py b/tests/test_litellm/llms/gigachat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/litellm/llms/gigachat/test_utils.py b/tests/test_litellm/llms/gigachat/test_utils.py similarity index 93% rename from tests/litellm/llms/gigachat/test_utils.py rename to tests/test_litellm/llms/gigachat/test_utils.py index 00faf391078..e1c7b75217f 100644 --- a/tests/litellm/llms/gigachat/test_utils.py +++ b/tests/test_litellm/llms/gigachat/test_utils.py @@ -5,9 +5,7 @@ Tests for litellm.llms.gigachat.utils import os import sys -sys.path.insert( - 0, os.path.abspath("../../../../../") -) # Adds the project root to the system path +sys.path.insert(0, os.path.abspath("../../..")) import pytest from litellm.llms.gigachat.utils import convert_usage @@ -81,4 +79,4 @@ class TestConvertUsage: assert result.prompt_tokens == 10 assert result.completion_tokens == 5 assert result.total_tokens == 15 - assert result.prompt_tokens_details is None + assert result.prompt_tokens_details is None \ No newline at end of file From 29a251bee84cf65debc5a3e7061976a613590579 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Sun, 23 Aug 2026 11:46:16 +0000 Subject: [PATCH 048/120] fix(lint): Remove definition `Union` in litellm_logging.py --- litellm/litellm_core_utils/litellm_logging.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index dcc527469a8..b850ba3911d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5999,7 +5999,7 @@ def _get_traceback_str_for_error(error_str: str) -> str: from decimal import Decimal # used for unit testing -from typing import Any, Optional, Union +from typing import Any, Optional def create_dummy_standard_logging_payload() -> StandardLoggingPayload: From 88747ef70cf3437a4c855e92d52c1e5c6ba21cbe Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Sun, 23 Aug 2026 12:03:01 +0000 Subject: [PATCH 049/120] fix(test): add tests for gigachat --- .../chat/test_gigachat_chat_transformation.py | 887 ++++++++++++++++++ .../llms/gigachat/embedding/__init__.py | 0 .../test_gigachat_embedding_transformation.py | 375 ++++++++ ...est_gigachat_passthrough_transformation.py | 128 +++ .../llms/gigachat/test_file_handler.py | 508 ++++++++++ 5 files changed, 1898 insertions(+) create mode 100644 tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py create mode 100644 tests/test_litellm/llms/gigachat/embedding/__init__.py create mode 100644 tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py create mode 100644 tests/test_litellm/llms/gigachat/test_file_handler.py diff --git a/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py b/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py new file mode 100644 index 00000000000..42b5c1894b2 --- /dev/null +++ b/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py @@ -0,0 +1,887 @@ +""" +Unit tests for GigaChat chat transformation. + +Tests GigaChatConfig covering get_complete_url, validate_environment, +get_supported_openai_params, map_openai_params, _convert_tools_to_functions, +_map_tool_choice, _transform_messages, transform_request, transform_response, +get_model_response_iterator, and get_error_class. +""" + +import json +import os +import sys +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.gigachat.chat.transformation import ( + GigaChatConfig, + GigaChatError, + is_valid_json, +) +from litellm.types.utils import ModelResponse, Usage + +TRANSFORM_MODULE = "litellm.llms.gigachat.chat.transformation" + + +def _make_httpx_response( + body: dict, status_code: int = 200 +) -> httpx.Response: + return httpx.Response( + status_code=status_code, + headers={"content-type": "application/json"}, + content=json.dumps(body).encode("utf-8"), + request=httpx.Request( + "POST", + "https://gigachat.devices.sberbank.ru/api/v1/chat/completions", + ), + ) + + +# --------------------------------------------------------------------------- +# is_valid_json +# --------------------------------------------------------------------------- + + +class TestIsValidJson: + def test_valid_json_object(self): + assert is_valid_json('{"key": "value"}') is True + + def test_valid_json_array(self): + assert is_valid_json("[1, 2, 3]") is True + + def test_valid_json_string(self): + assert is_valid_json('"hello"') is True + + def test_invalid_json(self): + assert is_valid_json("{invalid}") is False + + def test_empty_string(self): + assert is_valid_json("") is False + + +# --------------------------------------------------------------------------- +# GigaChatConfig +# --------------------------------------------------------------------------- + + +class TestGetCompleteUrl: + def setup_method(self): + self.config = GigaChatConfig() + + def test_uses_api_base_from_param(self): + url = self.config.get_complete_url( + api_base="https://custom.example.com", + api_key=None, + model="GigaChat", + optional_params={}, + litellm_params={}, + stream=False, + ) + assert url == "https://custom.example.com/chat/completions" + + def test_uses_api_base_with_trailing_slash(self): + url = self.config.get_complete_url( + api_base="https://custom.example.com/", + api_key=None, + model="GigaChat", + optional_params={}, + litellm_params={}, + stream=False, + ) + # get_api_base passes the value through without stripping the slash + assert url == "https://custom.example.com//chat/completions" + + def test_uses_api_base_from_get_api_base_when_none(self): + url = self.config.get_complete_url( + api_base=None, + api_key=None, + model="GigaChat", + optional_params={}, + litellm_params={}, + stream=False, + ) + assert url.endswith("/chat/completions") + + +class TestValidateEnvironment: + def setup_method(self): + self.config = GigaChatConfig() + + @patch(f"{TRANSFORM_MODULE}.get_access_token", return_value="test-token") + @patch(f"{TRANSFORM_MODULE}.get_secret_str", return_value=None) + def test_sets_auth_headers(self, mock_get_secret, mock_get_token): + headers: dict = {} + result = self.config.validate_environment( + headers=headers, + model="GigaChat", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + api_key="creds", + api_base="https://api.example.com", + ) + assert result["Authorization"] == "Bearer test-token" + assert result["Content-Type"] == "application/json" + assert result["Accept"] == "application/json" + + @patch(f"{TRANSFORM_MODULE}.get_access_token", return_value="token") + @patch(f"{TRANSFORM_MODULE}.get_secret_str", return_value=None) + def test_stores_credentials_and_api_base_for_image_uploads( + self, mock_get_secret, mock_get_token + ): + self.config.validate_environment( + headers={}, + model="GigaChat", + messages=[], + optional_params={}, + litellm_params={}, + api_key="my-creds", + api_base="https://my-api.example.com", + ) + assert self.config._current_credentials == "my-creds" + assert self.config._current_api_base == "https://my-api.example.com" + + @patch(f"{TRANSFORM_MODULE}.get_access_token", return_value="token") + @patch(f"{TRANSFORM_MODULE}.get_secret_str") + def test_falls_back_to_env_for_credentials( + self, mock_get_secret, mock_get_token + ): + mock_get_secret.return_value = "env-creds" + self.config.validate_environment( + headers={}, + model="GigaChat", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + mock_get_secret.assert_any_call("GIGACHAT_CREDENTIALS") + + +class TestGetSupportedOpenAiParams: + def setup_method(self): + self.config = GigaChatConfig() + + def test_returns_expected_params(self): + params = self.config.get_supported_openai_params("GigaChat") + expected = [ + "stream", + "temperature", + "top_p", + "max_tokens", + "max_completion_tokens", + "stop", + "tools", + "tool_choice", + "functions", + "function_call", + "response_format", + ] + assert params == expected + + +class TestMapOpenAiParams: + def setup_method(self): + self.config = GigaChatConfig() + + def test_stream(self): + result = self.config.map_openai_params( + non_default_params={"stream": True}, + optional_params={}, + model="GigaChat", + drop_params=False, + ) + assert result["stream"] is True + + def test_temperature_zero_maps_to_top_p_zero(self): + result = self.config.map_openai_params( + non_default_params={"temperature": 0}, + optional_params={}, + model="GigaChat", + drop_params=False, + ) + assert result["top_p"] == 0 + assert "temperature" not in result + + def test_temperature_non_zero(self): + result = self.config.map_openai_params( + non_default_params={"temperature": 0.7}, + optional_params={}, + model="GigaChat", + drop_params=False, + ) + assert result["temperature"] == 0.7 + + def test_top_p(self): + result = self.config.map_openai_params( + non_default_params={"top_p": 0.5}, + optional_params={}, + model="GigaChat", + drop_params=False, + ) + assert result["top_p"] == 0.5 + + def test_max_tokens(self): + result = self.config.map_openai_params( + non_default_params={"max_tokens": 100}, + optional_params={}, + model="GigaChat", + drop_params=False, + ) + assert result["max_tokens"] == 100 + + def test_max_completion_tokens(self): + result = self.config.map_openai_params( + non_default_params={"max_completion_tokens": 200}, + optional_params={}, + model="GigaChat", + drop_params=False, + ) + assert result["max_tokens"] == 200 + + def test_stop_is_dropped(self): + result = self.config.map_openai_params( + non_default_params={"stop": ["\n\n"]}, + optional_params={}, + model="GigaChat", + drop_params=False, + ) + assert "stop" not in result + + def test_tools_converted_to_functions(self): + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather", + "parameters": {"type": "object"}, + }, + } + ] + result = self.config.map_openai_params( + non_default_params={"tools": tools}, + optional_params={}, + model="GigaChat", + drop_params=False, + ) + assert "functions" in result + assert result["functions"] == [ + {"name": "get_weather", "description": "Get weather", "parameters": {"type": "object"}} + ] + + def test_tool_choice_auto(self): + result = self.config.map_openai_params( + non_default_params={"tool_choice": "auto"}, + optional_params={}, + model="GigaChat", + drop_params=False, + ) + assert result.get("function_call") == "auto" + + def test_tool_choice_none(self): + result = self.config.map_openai_params( + non_default_params={"tool_choice": "none"}, + optional_params={}, + model="GigaChat", + drop_params=False, + ) + assert result.get("function_call") == "none" + + def test_tool_choice_required(self): + result = self.config.map_openai_params( + non_default_params={"tool_choice": "required"}, + optional_params={}, + model="GigaChat", + drop_params=False, + ) + assert result.get("function_call") == "auto" + + def test_tool_choice_dict(self): + result = self.config.map_openai_params( + non_default_params={ + "tool_choice": { + "type": "function", + "function": {"name": "get_weather"}, + } + }, + optional_params={}, + model="GigaChat", + drop_params=False, + ) + assert result.get("function_call") == {"name": "get_weather"} + + def test_functions(self): + funcs = [{"name": "my_func", "description": "desc", "parameters": {}}] + result = self.config.map_openai_params( + non_default_params={"functions": funcs}, + optional_params={}, + model="GigaChat", + drop_params=False, + ) + assert result["functions"] == funcs + + def test_function_call(self): + result = self.config.map_openai_params( + non_default_params={"function_call": {"name": "my_func"}}, + optional_params={}, + model="GigaChat", + drop_params=False, + ) + assert result["function_call"] == {"name": "my_func"} + + def test_response_format_json_schema(self): + response_format = { + "type": "json_schema", + "json_schema": { + "name": "test_schema", + "schema": {"type": "object", "properties": {"name": {"type": "string"}}}, + }, + } + result = self.config.map_openai_params( + non_default_params={"response_format": response_format}, + optional_params={"functions": []}, + model="GigaChat", + drop_params=False, + ) + # Should add a function for the schema + assert len(result["functions"]) == 1 + assert result["functions"][0]["name"] == "test_schema" + assert result["function_call"] == {"name": "test_schema"} + assert result["_structured_output"] is True + + +class TestConvertToolsToFunctions: + def setup_method(self): + self.config = GigaChatConfig() + + def test_converts_function_tools_only(self): + tools = [ + {"type": "function", "function": {"name": "a", "description": "d", "parameters": {}}}, + {"type": "code_interpreter"}, # should be ignored + ] + result = self.config._convert_tools_to_functions(tools) + assert len(result) == 1 + assert result[0]["name"] == "a" + + def test_empty_tools(self): + assert self.config._convert_tools_to_functions([]) == [] + + +class TestMapToolChoice: + def setup_method(self): + self.config = GigaChatConfig() + + def test_none(self): + assert self.config._map_tool_choice("none") == "none" + + def test_auto(self): + assert self.config._map_tool_choice("auto") == "auto" + + def test_required(self): + assert self.config._map_tool_choice("required") == "auto" + + def test_dict_with_function(self): + result = self.config._map_tool_choice( + {"type": "function", "function": {"name": "get_weather"}} + ) + assert result == {"name": "get_weather"} + + def test_dict_without_name(self): + result = self.config._map_tool_choice( + {"type": "function", "function": {}} + ) + assert result is None + + def test_unknown_value(self): + assert self.config._map_tool_choice("unknown") is None + + +class TestTransformMessages: + def setup_method(self): + self.config = GigaChatConfig() + + def test_developer_role_to_system(self): + result = self.config._transform_messages( + [{"role": "developer", "content": "be helpful"}] + ) + assert result[0]["role"] == "system" + assert result[0]["content"] == "be helpful" + + def test_system_message_not_first_becomes_user(self): + result = self.config._transform_messages([ + {"role": "user", "content": "hi"}, + {"role": "system", "content": "instruction"}, + ]) + assert result[0]["role"] == "user" + assert result[1]["role"] == "user" + assert result[1]["content"] == "instruction" + + def test_tool_role_to_function(self): + result = self.config._transform_messages([ + {"role": "tool", "content": '{"result": "ok"}'} + ]) + assert result[0]["role"] == "function" + + def test_tool_role_content_wraps_non_json(self): + result = self.config._transform_messages([ + {"role": "tool", "content": "plain text"} + ]) + assert result[0]["role"] == "function" + assert is_valid_json(result[0]["content"]) + + def test_none_content_becomes_empty_string(self): + result = self.config._transform_messages([ + {"role": "user", "content": None} + ]) + assert result[0]["content"] == "" + + def test_name_field_removed(self): + result = self.config._transform_messages([ + {"role": "user", "content": "hi", "name": "John"} + ]) + assert "name" not in result[0] + + def test_tool_calls_converted_to_function_call(self): + result = self.config._transform_messages([ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "London"}', + }, + } + ], + } + ]) + assert "tool_calls" not in result[0] + assert result[0]["function_call"]["name"] == "get_weather" + assert result[0]["function_call"]["arguments"] == {"city": "London"} + + def test_tool_calls_with_dict_arguments(self): + result = self.config._transform_messages([ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_xyz", + "type": "function", + "function": { + "name": "search", + "arguments": {"query": "test"}, + }, + } + ], + } + ]) + assert result[0]["function_call"]["arguments"] == {"query": "test"} + + def test_list_content_multimodal(self): + content = [ + {"type": "text", "text": "describe this"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/img.jpg"}, + }, + ] + with patch.object(self.config, "_upload_image", return_value="file-123"): + result = self.config._transform_messages([ + {"role": "user", "content": content} + ]) + assert result[0]["content"] == "describe this" + assert result[0]["attachments"] == ["file-123"] + + def test_list_content_with_image_url_string(self): + content = [ + {"type": "text", "text": "look"}, + {"type": "image_url", "image_url": "https://example.com/img.jpg"}, + ] + with patch.object(self.config, "_upload_image", return_value="file-456"): + result = self.config._transform_messages([ + {"role": "user", "content": content} + ]) + assert result[0]["content"] == "look" + assert "file-456" in result[0]["attachments"] + + +class TestTransformRequest: + def setup_method(self): + self.config = GigaChatConfig() + + def test_builds_basic_request(self): + body = self.config.transform_request( + model="gigachat/GigaChat", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert body["model"] == "GigaChat" + assert len(body["messages"]) == 1 + assert body["messages"][0]["content"] == "hi" + + def test_model_prefix_stripped(self): + body = self.config.transform_request( + model="gigachat/GigaChat-Pro", + messages=[{"role": "user", "content": "hello"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert body["model"] == "GigaChat-Pro" + + def test_includes_optional_params(self): + body = self.config.transform_request( + model="gigachat/GigaChat", + messages=[{"role": "user", "content": "hi"}], + optional_params={ + "temperature": 0.5, + "max_tokens": 100, + "stream": True, + }, + litellm_params={}, + headers={}, + ) + assert body["temperature"] == 0.5 + assert body["max_tokens"] == 100 + assert body["stream"] is True + + def test_includes_functions(self): + body = self.config.transform_request( + model="gigachat/GigaChat", + messages=[{"role": "user", "content": "hi"}], + optional_params={ + "functions": [{"name": "my_func"}], + "function_call": {"name": "my_func"}, + }, + litellm_params={}, + headers={}, + ) + assert body["functions"] == [{"name": "my_func"}] + assert body["function_call"] == {"name": "my_func"} + + def test_skips_unsupported_params(self): + body = self.config.transform_request( + model="gigachat/GigaChat", + messages=[{"role": "user", "content": "hi"}], + optional_params={"n": 2, "user": "abc"}, + litellm_params={}, + headers={}, + ) + assert "n" not in body + assert "user" not in body + + +class TestTransformResponse: + def setup_method(self): + self.config = GigaChatConfig() + + def test_basic_response(self): + raw = _make_httpx_response({ + "id": "chatcmpl-123", + "created": 1700000000, + "model": "GigaChat", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello!"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, + }) + model_response = ModelResponse() + result = self.config.transform_response( + model="gigachat/GigaChat", + raw_response=raw, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert result.choices[0].message.content == "Hello!" + assert result.choices[0].finish_reason == "stop" + assert result.usage.prompt_tokens == 5 + assert result.usage.total_tokens == 8 + + def test_function_call_into_tool_calls(self): + raw = _make_httpx_response({ + "id": "chatcmpl-456", + "created": 1700000000, + "model": "GigaChat", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "function_call": { + "name": "get_weather", + "arguments": {"city": "Moscow"}, + }, + }, + "finish_reason": "function_call", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }) + model_response = ModelResponse() + result = self.config.transform_response( + model="gigachat/GigaChat", + raw_response=raw, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert result.choices[0].finish_reason == "tool_calls" + tool_calls = result.choices[0].message.tool_calls + assert tool_calls is not None + assert len(tool_calls) == 1 + assert tool_calls[0].function.name == "get_weather" + assert '{"city": "Moscow"}' in tool_calls[0].function.arguments + + def test_function_call_structured_output(self): + raw = _make_httpx_response({ + "id": "chatcmpl-789", + "created": 1700000000, + "model": "GigaChat", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "function_call": { + "name": "test_schema", + "arguments": {"name": "John"}, + }, + }, + "finish_reason": "function_call", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }) + model_response = ModelResponse() + result = self.config.transform_response( + model="gigachat/GigaChat", + raw_response=raw, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={"_structured_output": True}, + litellm_params={}, + encoding=None, + ) + # Structured output: function_call -> content + assert result.choices[0].finish_reason == "stop" + assert result.choices[0].message.content is not None + assert '"name": "John"' in result.choices[0].message.content + + def test_function_call_string_arguments(self): + raw = _make_httpx_response({ + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "function_call": { + "name": "get_weather", + "arguments": '{"city": "Moscow"}', + }, + }, + "finish_reason": "function_call", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) + model_response = ModelResponse() + result = self.config.transform_response( + model="gigachat/GigaChat", + raw_response=raw, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + tc = result.choices[0].message.tool_calls[0] + assert '{"city": "Moscow"}' in tc.function.arguments + + def test_cleans_up_gigachat_specific_fields(self): + raw = _make_httpx_response({ + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "done", + "functions_state_id": "some-state", + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) + model_response = ModelResponse() + result = self.config.transform_response( + model="gigachat/GigaChat", + raw_response=raw, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + # functions_state_id should have been removed from the message data + assert result.choices[0].message.content == "done" + + def test_raises_on_invalid_json(self): + raw = httpx.Response( + status_code=500, + headers={"content-type": "text/plain"}, + content=b"not json", + request=httpx.Request("POST", "https://example.com"), + ) + model_response = ModelResponse() + with pytest.raises(GigaChatError) as exc_info: + self.config.transform_response( + model="gigachat/GigaChat", + raw_response=raw, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert "Invalid JSON response" in str(exc_info.value.message) + + def test_empty_choices(self): + raw = _make_httpx_response({ + "choices": [], + "usage": {}, + }) + model_response = ModelResponse() + result = self.config.transform_response( + model="gigachat/GigaChat", + raw_response=raw, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert result.choices == [] + + def test_function_call_with_non_dict_arguments(self): + raw = _make_httpx_response({ + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "function_call": { + "name": "say_hello", + "arguments": "hello", + }, + }, + "finish_reason": "function_call", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) + model_response = ModelResponse() + result = self.config.transform_response( + model="gigachat/GigaChat", + raw_response=raw, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + tc = result.choices[0].message.tool_calls[0] + assert tc.function.arguments == "hello" + + +class TestGetModelResponseIterator: + def setup_method(self): + self.config = GigaChatConfig() + + def test_returns_gigachat_iterator_sync(self): + from litellm.llms.gigachat.chat.streaming import ( + GigaChatModelResponseIterator, + ) + + result = self.config.get_model_response_iterator( + streaming_response=iter(["data"]), + sync_stream=True, + json_mode=False, + ) + assert isinstance(result, GigaChatModelResponseIterator) + + +class TestGetErrorClass: + def setup_method(self): + self.config = GigaChatConfig() + + def test_returns_gigachat_error(self): + error = self.config.get_error_class( + error_message="something went wrong", + status_code=400, + headers={"x-request-id": "abc"}, + ) + assert isinstance(error, GigaChatError) + assert error.status_code == 400 + assert error.message == "something went wrong" + assert error.headers == {"x-request-id": "abc"} + + +class TestUploadImage: + def setup_method(self): + self.config = GigaChatConfig() + + @patch(f"{TRANSFORM_MODULE}.upload_file_sync", return_value="file-uploaded") + def test_upload_image_success(self, mock_upload): + self.config._current_credentials = "creds" + self.config._current_api_base = "https://api.example.com" + result = self.config._upload_image("https://example.com/img.jpg") + assert result == "file-uploaded" + mock_upload.assert_called_once_with( + image_url="https://example.com/img.jpg", + credentials="creds", + api_base="https://api.example.com", + ) + + @patch(f"{TRANSFORM_MODULE}.upload_file_sync", side_effect=Exception("fail")) + def test_upload_image_failure_returns_none(self, mock_upload): + result = self.config._upload_image("https://example.com/img.jpg") + assert result is None \ No newline at end of file diff --git a/tests/test_litellm/llms/gigachat/embedding/__init__.py b/tests/test_litellm/llms/gigachat/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py b/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py new file mode 100644 index 00000000000..594f592bab3 --- /dev/null +++ b/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py @@ -0,0 +1,375 @@ +""" +Unit tests for GigaChat embedding transformation. + +Tests GigaChatEmbeddingConfig covering get_config, get_supported_openai_params, +map_openai_params, _get_openai_compatible_provider_info, get_complete_url, +transform_embedding_request, transform_embedding_response, validate_environment, +and get_error_class. +""" + +import json +import os +import sys +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm import LlmProviders +from litellm.llms.gigachat.embedding.transformation import ( + GigaChatEmbeddingConfig, + GigaChatEmbeddingError, +) +from litellm.types.utils import EmbeddingResponse + +TRANSFORM_MODULE = "litellm.llms.gigachat.embedding.transformation" + + +def _make_httpx_response(body: dict, status_code: int = 200) -> httpx.Response: + return httpx.Response( + status_code=status_code, + headers={"content-type": "application/json"}, + content=json.dumps(body).encode("utf-8"), + request=httpx.Request("POST", "https://gigachat.devices.sberbank.ru/api/v1/embeddings"), + ) + + +# --------------------------------------------------------------------------- +# GigaChatEmbeddingConfig +# --------------------------------------------------------------------------- + + +class TestGetConfig: + def setup_method(self): + self.config = GigaChatEmbeddingConfig() + + def test_contains_only_abc_impl(self): + """get_config returns ABC internal data due to inheritance.""" + result = self.config.get_config() + # The only key should be _abc_impl from ABC base class + assert set(result.keys()) == {"_abc_impl"} + + +class TestGetSupportedOpenAiParams: + def setup_method(self): + self.config = GigaChatEmbeddingConfig() + + def test_returns_empty_list(self): + params = self.config.get_supported_openai_params("GigaChat") + assert params == [] + + +class TestMapOpenAiParams: + def setup_method(self): + self.config = GigaChatEmbeddingConfig() + + def test_returns_optional_params_unchanged(self): + result = self.config.map_openai_params( + non_default_params={"model": "test"}, + optional_params={"temperature": 0.5}, + model="GigaChat", + drop_params=False, + ) + assert result == {"temperature": 0.5} + + def test_returns_empty_dict_when_no_optional_params(self): + result = self.config.map_openai_params( + non_default_params={}, + optional_params={}, + model="GigaChat", + drop_params=False, + ) + assert result == {} + + +class TestGetOpenaiCompatibleProviderInfo: + def setup_method(self): + self.config = GigaChatEmbeddingConfig() + + def test_returns_gigachat_provider(self): + provider, api_base, api_key = self.config._get_openai_compatible_provider_info( + api_base="https://api.example.com", api_key="test-key" + ) + assert provider == LlmProviders.GIGACHAT.value + assert api_base == "https://api.example.com" + assert api_key == "test-key" + + def test_resolves_api_base_when_none(self): + provider, api_base, api_key = self.config._get_openai_compatible_provider_info( + api_base=None, api_key="key" + ) + assert api_base is not None + assert api_base.endswith("/api/v1") + + def test_returns_none_api_key(self): + _, _, api_key = self.config._get_openai_compatible_provider_info( + api_base="https://example.com", api_key=None + ) + assert api_key is None + + +class TestGetCompleteUrl: + def setup_method(self): + self.config = GigaChatEmbeddingConfig() + + def test_default_url(self): + url = self.config.get_complete_url( + api_base=None, api_key=None, model="GigaChat", + optional_params={}, litellm_params={}, + ) + assert url.endswith("/embeddings") + + def test_custom_api_base(self): + url = self.config.get_complete_url( + api_base="https://custom.example.com", api_key=None, model="GigaChat", + optional_params={}, litellm_params={}, + ) + assert url == "https://custom.example.com/embeddings" + + def test_trailing_slash_api_base(self): + url = self.config.get_complete_url( + api_base="https://custom.example.com/", api_key=None, model="GigaChat", + optional_params={}, litellm_params={}, + ) + # get_api_base doesn't strip slash, so we get double slash + assert url == "https://custom.example.com//embeddings" + + +class TestTransformEmbeddingRequest: + def setup_method(self): + self.config = GigaChatEmbeddingConfig() + + def test_string_input(self): + result = self.config.transform_embedding_request( + model="gigachat/Embeddings", + input="hello world", + optional_params={}, + headers={}, + ) + assert result == {"model": "Embeddings", "input": ["hello world"]} + + def test_list_input(self): + result = self.config.transform_embedding_request( + model="gigachat/Embeddings", + input=["text1", "text2"], + optional_params={}, + headers={}, + ) + assert result == {"model": "Embeddings", "input": ["text1", "text2"]} + + def test_strips_gigachat_prefix(self): + result = self.config.transform_embedding_request( + model="gigachat/GigaChat-Pro", + input="test", + optional_params={}, + headers={}, + ) + assert result["model"] == "GigaChat-Pro" + + def test_model_without_prefix(self): + result = self.config.transform_embedding_request( + model="Embeddings", + input="test", + optional_params={}, + headers={}, + ) + assert result["model"] == "Embeddings" + + +class TestTransformEmbeddingResponse: + def setup_method(self): + self.config = GigaChatEmbeddingConfig() + self.logging_obj = MagicMock() + + def _make_gigachat_response(self, data: list[dict]) -> httpx.Response: + return _make_httpx_response({ + "object": "list", + "data": data, + "model": "Embeddings", + }) + + def test_basic_response(self): + raw = self._make_gigachat_response([ + { + "object": "embedding", + "embedding": [0.1, 0.2, 0.3], + "index": 0, + } + ]) + model_response = EmbeddingResponse() + result = self.config.transform_embedding_response( + model="gigachat/Embeddings", + raw_response=raw, + model_response=model_response, + logging_obj=self.logging_obj, + api_key="test-key", + request_data={"input": ["text"]}, + optional_params={}, + litellm_params={}, + ) + assert result.object == "list" + assert len(result.data) == 1 + assert result.data[0]["embedding"] == [0.1, 0.2, 0.3] + assert result.data[0]["index"] == 0 + assert result.usage.prompt_tokens == 0 + assert result.usage.total_tokens == 0 + + def test_aggregates_per_embedding_usage(self): + raw = self._make_gigachat_response([ + { + "object": "embedding", + "embedding": [0.1, 0.2], + "index": 0, + "usage": {"prompt_tokens": 5}, + }, + { + "object": "embedding", + "embedding": [0.3, 0.4], + "index": 1, + "usage": {"prompt_tokens": 7}, + }, + ]) + model_response = EmbeddingResponse() + result = self.config.transform_embedding_response( + model="gigachat/Embeddings", + raw_response=raw, + model_response=model_response, + logging_obj=self.logging_obj, + api_key="test-key", + request_data={"input": ["a", "b"]}, + optional_params={}, + litellm_params={}, + ) + # Total should be sum of per-embedding prompt_tokens + assert result.usage.prompt_tokens == 12 + assert result.usage.total_tokens == 12 + # Usage should be removed from individual embedding data + assert "usage" not in result.data[0] + assert "usage" not in result.data[1] + + def test_usage_removed_from_individual_embeddings(self): + raw = self._make_gigachat_response([ + { + "object": "embedding", + "embedding": [0.5], + "index": 0, + "usage": {"prompt_tokens": 3}, + } + ]) + model_response = EmbeddingResponse() + result = self.config.transform_embedding_response( + model="gigachat/Embeddings", + raw_response=raw, + model_response=model_response, + logging_obj=self.logging_obj, + api_key="key", + request_data={"input": ["x"]}, + optional_params={}, + litellm_params={}, + ) + # usage should NOT be in the final EmbeddingResponse data items + for emb in result.data: + assert "usage" not in emb + + def test_passes_model_from_response(self): + raw = self._make_gigachat_response([ + {"object": "embedding", "embedding": [0.1], "index": 0}, + ]) + model_response = EmbeddingResponse() + result = self.config.transform_embedding_response( + model="gigachat/Embeddings", + raw_response=raw, + model_response=model_response, + logging_obj=self.logging_obj, + api_key="key", + request_data={"input": ["x"]}, + optional_params={}, + litellm_params={}, + ) + assert result.model == "Embeddings" + + def test_calls_logging_post_call(self): + raw = self._make_gigachat_response([ + {"object": "embedding", "embedding": [0.1], "index": 0}, + ]) + model_response = EmbeddingResponse() + self.config.transform_embedding_response( + model="gigachat/Embeddings", + raw_response=raw, + model_response=model_response, + logging_obj=self.logging_obj, + api_key="test-api-key", + request_data={"input": ["hello"]}, + optional_params={}, + litellm_params={}, + ) + self.logging_obj.post_call.assert_called_once() + args = self.logging_obj.post_call.call_args.kwargs + assert args["api_key"] == "test-api-key" + assert args["input"] == ["hello"] + + +class TestValidateEnvironment: + def setup_method(self): + self.config = GigaChatEmbeddingConfig() + + @patch(f"{TRANSFORM_MODULE}.get_access_token", return_value="test-token") + def test_sets_oauth_headers(self, mock_get_token): + headers = self.config.validate_environment( + headers={}, + model="GigaChat", + messages=[], + optional_params={}, + litellm_params={}, + api_key="creds", + api_base="https://api.example.com", + ) + assert headers["Authorization"] == "Bearer test-token" + assert headers["Content-Type"] == "application/json" + mock_get_token.assert_called_once_with(credentials="creds", litellm_params={}) + + @patch(f"{TRANSFORM_MODULE}.get_access_token", return_value="token") + def test_merges_custom_headers(self, mock_get_token): + headers = self.config.validate_environment( + headers={"X-Custom": "value"}, + model="GigaChat", + messages=[], + optional_params={}, + litellm_params={}, + api_key="creds", + api_base="https://api.example.com", + ) + assert headers["Authorization"] == "Bearer token" + assert headers["Content-Type"] == "application/json" + assert headers["X-Custom"] == "value" + + @patch(f"{TRANSFORM_MODULE}.get_access_token", return_value="token") + def test_custom_header_overwrites_default(self, mock_get_token): + headers = self.config.validate_environment( + headers={"Authorization": "Bearer custom"}, + model="GigaChat", + messages=[], + optional_params={}, + litellm_params={}, + api_key="creds", + api_base="https://api.example.com", + ) + # Merge: default headers first, then custom headers on top + assert headers["Authorization"] == "Bearer custom" + + +class TestGetErrorClass: + def setup_method(self): + self.config = GigaChatEmbeddingConfig() + + def test_returns_gigachat_embedding_error(self): + error = self.config.get_error_class( + error_message="embedding failed", + status_code=400, + headers={"x-request-id": "abc"}, + ) + assert isinstance(error, GigaChatEmbeddingError) + assert error.status_code == 400 + assert error.message == "embedding failed" \ No newline at end of file diff --git a/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py b/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py index b3e1500c4a3..e18a09b2448 100644 --- a/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py +++ b/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py @@ -479,3 +479,131 @@ class TestGigaChatPassthroughConfig: config = GigaChatPassthroughConfig() result = config.get_models() assert result == [] + + def test_logging_non_streaming_chat_raises_when_no_config(self): + """Test raise when ProviderConfigManager returns None for chat.""" + config = GigaChatPassthroughConfig() + logging_obj = MagicMock() + + with patch( + "litellm.utils.ProviderConfigManager.get_provider_chat_config", + return_value=None, + ): + with pytest.raises(ValueError, match="No provider config found for model"): + config.logging_non_streaming_response( + model="gigachat/GigaChat", + custom_llm_provider="gigachat", + httpx_response=_make_httpx_response(_gigachat_chat_completion_body()), + request_data={ + "model": "gigachat/GigaChat", + "messages": [{"role": "user", "content": "hi"}], + }, + logging_obj=logging_obj, + endpoint="chat/completions", + ) + + def test_logging_non_streaming_embedding_raises_when_no_config(self): + """Test raise when ProviderConfigManager returns None for embeddings.""" + config = GigaChatPassthroughConfig() + logging_obj = MagicMock() + + with patch( + "litellm.utils.ProviderConfigManager.get_provider_embedding_config", + return_value=None, + ): + with pytest.raises(ValueError, match="No provider config found for model"): + config.logging_non_streaming_response( + model="gigachat/Embeddings", + custom_llm_provider="gigachat", + httpx_response=_make_httpx_response(_gigachat_embedding_body()), + request_data={ + "input": ["hello"], + "model": "gigachat/Embeddings", + }, + logging_obj=logging_obj, + endpoint="embeddings", + ) + + def test_handle_logging_collected_chunks_with_model_response_stream_chunk(self): + """Test that a chunk returning ModelResponseStream from chunk_parser is handled. + + Requires patching GigaChatModelResponseIterator.chunk_parser to return + a ModelResponseStream so the elif branch is exercised. + """ + config = GigaChatPassthroughConfig() + logging_obj = MagicMock() + + from litellm.types.utils import ModelResponseStream + + stream_chunk = ModelResponseStream( + choices=[ + { + "index": 0, + "delta": {"content": "streamed"}, + "finish_reason": None, + } + ] + ) + + chunks = [ + '{"choices": [{"delta": {"content": "streamed"}, "index": 0}]}', + '{"choices": [{"delta": {}, "finish_reason": "stop", "index": 0}], "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}}', + ] + + with patch( + "litellm.llms.gigachat.passthrough.transformation.GigaChatModelResponseIterator.chunk_parser", + return_value=stream_chunk, + ): + result = config.handle_logging_collected_chunks( + all_chunks=chunks, + litellm_logging_obj=logging_obj, + model="gigachat/GigaChat", + custom_llm_provider="gigachat", + endpoint="chat/completions", + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "streamedstreamed" + + def test_handle_logging_collected_chunks_skips_unknown_chunk_type(self): + """Test that chunk_parser returning an unknown type is skipped.""" + config = GigaChatPassthroughConfig() + logging_obj = MagicMock() + + chunks = [ + '{"choices": [{"delta": {"content": "good"}, "index": 0}]}', + '{"choices": [{"delta": {}, "finish_reason": "stop", "index": 0}], "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}}', + ] + + with patch( + "litellm.llms.gigachat.passthrough.transformation.GigaChatModelResponseIterator.chunk_parser", + return_value=12345, # not dict and not ModelResponseStream + ): + result = config.handle_logging_collected_chunks( + all_chunks=chunks, + litellm_logging_obj=logging_obj, + model="gigachat/GigaChat", + custom_llm_provider="gigachat", + endpoint="chat/completions", + ) + + # All chunks skipped, returns None + assert result is None + + def test_handle_logging_collected_chunks_skips_unsupported_chunk_type(self): + """Test that unsupported chunk types (int, float, etc.) are skipped.""" + config = GigaChatPassthroughConfig() + logging_obj = MagicMock() + + # The chunk is an int which doesn't match str/bytes/dict + chunks: list = [42, "not-a-real-chunk"] + + result = config.handle_logging_collected_chunks( + all_chunks=chunks, + litellm_logging_obj=logging_obj, + model="gigachat/GigaChat", + custom_llm_provider="gigachat", + endpoint="chat/completions", + ) + + assert result is None diff --git a/tests/test_litellm/llms/gigachat/test_file_handler.py b/tests/test_litellm/llms/gigachat/test_file_handler.py new file mode 100644 index 00000000000..05fcedac338 --- /dev/null +++ b/tests/test_litellm/llms/gigachat/test_file_handler.py @@ -0,0 +1,508 @@ +""" +Unit tests for GigaChat file handler. + +Tests _get_url_hash, _parse_data_url, _download_image_sync, _download_image_async, +upload_file_sync, and upload_file_async covering caching, base64 data URL decoding, +network errors, and the full upload flow. +""" + +import json +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.llms.gigachat import file_handler +from litellm.llms.gigachat.file_handler import ( + _file_cache, + _get_url_hash, + _parse_data_url, + upload_file_async, + upload_file_sync, +) + +FILE_MODULE = "litellm.llms.gigachat.file_handler" + +# A valid 1x1 red PNG as base64 +_RED_PNG_B64 = ( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAA" + "DUlEQVQI12NgYPgPAAEDAQAR3X3ZAAAASUVORK5CYII=" +) +_RED_PNG_DATA_URL = f"data:image/png;base64,{_RED_PNG_B64}" + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(autouse=True) +def _isolate_file_cache(): + """Each test gets a fresh module-level file cache to avoid cross-test leakage.""" + _file_cache.clear() + yield + _file_cache.clear() + + +# --------------------------------------------------------------------------- +# _get_url_hash +# --------------------------------------------------------------------------- + + +class TestGetUrlHash: + def test_returns_hex_string(self): + h = _get_url_hash("https://example.com/image.png") + assert isinstance(h, str) + assert len(h) == 64 # SHA-256 + + def test_different_urls_different_hashes(self): + h1 = _get_url_hash("https://example.com/a.png") + h2 = _get_url_hash("https://example.com/b.png") + assert h1 != h2 + + def test_same_url_same_hash(self): + h1 = _get_url_hash("https://example.com/image.png") + h2 = _get_url_hash("https://example.com/image.png") + assert h1 == h2 + + +# --------------------------------------------------------------------------- +# _parse_data_url +# --------------------------------------------------------------------------- + + +class TestParseDataUrl: + def test_valid_base64_png(self): + result = _parse_data_url(_RED_PNG_DATA_URL) + assert result is not None + content_bytes, content_type, ext = result + assert content_type == "image/png" + assert ext == "png" + assert len(content_bytes) > 0 + + def test_valid_base64_jpeg(self): + # Simple valid base64 (24 chars, properly padded, no + or / chars) + valid_b64 = "aGVsbG8gd29ybGQhISEhIQ==" + data_url = f"data:image/jpeg;base64,{valid_b64}" + result = _parse_data_url(data_url) + assert result is not None + _, content_type, ext = result + assert content_type == "image/jpeg" + assert ext == "jpeg" + + def test_valid_base64_with_semicolon_in_type(self): + """Data URLs with charset before base64 segment do not match the regex.""" + # The regex `data:([^;]+);base64,(.+)` requires the pattern to be + # `data:;base64,`. If `;charset=utf-8` appears before + # `;base64,`, the regex sees `data:image/png` as group 1 but then + # looks for `;base64,` immediately after — which isn't there because + # `;charset=utf-8;base64,` has extra text before `;base64,` + data_url = "data:image/png;charset=utf-8;base64," + _RED_PNG_B64 + result = _parse_data_url(data_url) + assert result is None + + def test_invalid_data_url_returns_none(self): + assert _parse_data_url("not-a-data-url") is None + + def test_empty_base64_returns_none(self): + """Empty base64 data (nothing after comma) does not match regex `(.+)`.""" + assert _parse_data_url("data:image/png;base64,") is None + + def test_missing_base64_segment(self): + assert _parse_data_url("data:image/png;base64") is None + + def test_unknown_extension_falls_back_to_jpg(self): + data_url = "data:application/octet-stream;base64," + _RED_PNG_B64 + result = _parse_data_url(data_url) + assert result is not None + _, content_type, ext = result + assert content_type == "application/octet-stream" + # The extension is derived from content_type.split("/")[-1].split(";")[0] + # which gives "octet-stream", not "jpg" + assert ext == "octet-stream" + + +# --------------------------------------------------------------------------- +# _download_image_sync +# --------------------------------------------------------------------------- + + +class TestDownloadImageSync: + @patch(f"{FILE_MODULE}._get_httpx_client") + def test_downloads_image_successfully(self, mock_get_client): + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.content = b"fake-image-bytes" + mock_response.headers = {"content-type": "image/jpeg"} + mock_client.get.return_value = mock_response + mock_get_client.return_value = mock_client + + content_bytes, content_type, ext = file_handler._download_image_sync("https://example.com/img.jpg") + + assert content_bytes == b"fake-image-bytes" + assert content_type == "image/jpeg" + assert ext == "jpeg" + mock_client.get.assert_called_once_with("https://example.com/img.jpg") + + @patch(f"{FILE_MODULE}._get_httpx_client") + def test_raises_on_http_error(self, mock_get_client): + mock_client = MagicMock() + mock_client.get.side_effect = httpx.HTTPStatusError( + "Not Found", + request=httpx.Request("GET", "https://example.com/404"), + response=httpx.Response(status_code=404, request=httpx.Request("GET", "https://example.com/404")), + ) + mock_get_client.return_value = mock_client + + with pytest.raises(httpx.HTTPStatusError): + file_handler._download_image_sync("https://example.com/404") + + @patch(f"{FILE_MODULE}._get_httpx_client") + def test_parse_content_type_fallback(self, mock_get_client): + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.content = b"data" + mock_response.headers = {} + mock_client.get.return_value = mock_response + mock_get_client.return_value = mock_client + + _, content_type, ext = file_handler._download_image_sync("https://example.com/img") + + assert content_type == "image/jpeg" + assert ext == "jpeg" + + @patch(f"{FILE_MODULE}._get_httpx_client") + def test_extracts_extension_from_parametrized_type(self, mock_get_client): + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.content = b"data" + mock_response.headers = {"content-type": "image/png; charset=utf-8"} + mock_client.get.return_value = mock_response + mock_get_client.return_value = mock_client + + _, _, ext = file_handler._download_image_sync("https://example.com/img.png") + + assert ext == "png" + + +# --------------------------------------------------------------------------- +# _download_image_async +# --------------------------------------------------------------------------- + + +class TestDownloadImageAsync: + @pytest.mark.asyncio + @patch(f"{FILE_MODULE}.get_async_httpx_client") + async def test_downloads_image_successfully(self, mock_get_client): + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.content = b"fake-image-bytes" + mock_response.headers = {"content-type": "image/webp"} + mock_client.get = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + content_bytes, content_type, ext = await file_handler._download_image_async( + "https://example.com/img.webp" + ) + + assert content_bytes == b"fake-image-bytes" + assert content_type == "image/webp" + assert ext == "webp" + mock_client.get.assert_called_once_with("https://example.com/img.webp") + + @pytest.mark.asyncio + @patch(f"{FILE_MODULE}.get_async_httpx_client") + async def test_raises_on_http_error(self, mock_get_client): + mock_client = MagicMock() + mock_client.get = AsyncMock( + side_effect=httpx.HTTPStatusError( + "Forbidden", + request=httpx.Request("GET", "https://example.com/403"), + response=httpx.Response(status_code=403, request=httpx.Request("GET", "https://example.com/403")), + ) + ) + mock_get_client.return_value = mock_client + + with pytest.raises(httpx.HTTPStatusError): + await file_handler._download_image_async("https://example.com/403") + + +# --------------------------------------------------------------------------- +# upload_file_sync +# --------------------------------------------------------------------------- + + +class TestUploadFileSync: + @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") + @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") + @patch(f"{FILE_MODULE}._get_httpx_client") + def test_uploads_base64_image_and_caches( + self, mock_get_client, mock_get_token, mock_get_api_base + ): + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.json.return_value = {"id": "file-12345"} + mock_response.raise_for_status = MagicMock() + mock_client.post.return_value = mock_response + mock_get_client.return_value = mock_client + + result = upload_file_sync( + image_url=_RED_PNG_DATA_URL, + credentials="creds", + api_base="https://custom.example.com", + ) + + assert result == "file-12345" + # Verify it was cached + url_hash = _get_url_hash(_RED_PNG_DATA_URL) + assert _file_cache[url_hash] == "file-12345" + + # Check the upload request — url is passed as first positional arg + call_args = mock_client.post.call_args + assert call_args.args[0] == "https://api.example.com/files" + assert call_args.kwargs["headers"]["Authorization"] == "Bearer test-token" + # Verify purpose + assert call_args.kwargs["data"] == {"purpose": "general"} + # Verify a file was attached + assert "file" in call_args.kwargs["files"] + + @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") + @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") + @patch(f"{FILE_MODULE}._get_httpx_client") + def test_returns_cached_file_id( + self, mock_get_client, mock_get_token, mock_get_api_base + ): + # Pre-populate the cache + url_hash = _get_url_hash(_RED_PNG_DATA_URL) + _file_cache[url_hash] = "cached-file-id" + + result = upload_file_sync(image_url=_RED_PNG_DATA_URL, credentials="creds") + + assert result == "cached-file-id" + # No upload call was made + mock_get_client.return_value.post.assert_not_called() + + @patch(f"{FILE_MODULE}._get_httpx_client") + @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") + @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") + @patch(f"{FILE_MODULE}._download_image_sync") + def test_downloads_and_uploads_url_image( + self, mock_download, mock_get_api_base, mock_get_token, mock_get_client + ): + mock_download.return_value = (b"remote-bytes", "image/png", "png") + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.json.return_value = {"id": "file-remote"} + mock_response.raise_for_status = MagicMock() + mock_client.post.return_value = mock_response + mock_get_client.return_value = mock_client + + result = upload_file_sync( + image_url="https://example.com/remote.png", credentials="creds" + ) + + assert result == "file-remote" + mock_download.assert_called_once_with("https://example.com/remote.png") + + @patch(f"{FILE_MODULE}._get_httpx_client") + @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") + @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") + def test_returns_none_on_upload_failure( + self, mock_get_api_base, mock_get_token, mock_get_client + ): + mock_client = MagicMock() + mock_client.post.side_effect = httpx.HTTPStatusError( + "Bad Request", + request=httpx.Request("POST", "https://api.example.com/files"), + response=httpx.Response(status_code=400, request=httpx.Request("POST", "https://api.example.com/files")), + ) + mock_get_client.return_value = mock_client + + # upload_file_sync catches all exceptions and returns None + result = upload_file_sync( + image_url=_RED_PNG_DATA_URL, credentials="creds" + ) + + assert result is None + + @patch(f"{FILE_MODULE}._get_httpx_client") + @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") + @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") + def test_returns_none_when_response_missing_id( + self, mock_get_api_base, mock_get_token, mock_get_client + ): + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.json.return_value = {"status": "ok"} # no "id" key + mock_response.raise_for_status = MagicMock() + mock_client.post.return_value = mock_response + mock_get_client.return_value = mock_client + + result = upload_file_sync( + image_url=_RED_PNG_DATA_URL, credentials="creds" + ) + + assert result is None + + @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") + @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") + @patch(f"{FILE_MODULE}._get_httpx_client") + def test_uploads_without_optional_args( + self, mock_get_client, mock_get_token, mock_get_api_base + ): + """Verify that credentials, api_base, and litellm_params are optional.""" + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.json.return_value = {"id": "file-no-args"} + mock_response.raise_for_status = MagicMock() + mock_client.post.return_value = mock_response + mock_get_client.return_value = mock_client + + result = upload_file_sync(image_url=_RED_PNG_DATA_URL) + + assert result == "file-no-args" + # Should still have called get_access_token without args + mock_get_token.assert_called_once_with(credentials=None, litellm_params=None) + + +# --------------------------------------------------------------------------- +# upload_file_async +# --------------------------------------------------------------------------- + + +class TestUploadFileAsync: + @pytest.mark.asyncio + @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") + @patch(f"{FILE_MODULE}.get_access_token_async", return_value="test-token-async") + @patch(f"{FILE_MODULE}.get_async_httpx_client") + async def test_uploads_base64_image_and_caches( + self, mock_get_client, mock_get_token, mock_get_api_base + ): + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.json = MagicMock(return_value={"id": "async-file-1"}) + mock_response.raise_for_status = MagicMock() + mock_client.post = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + result = await upload_file_async( + image_url=_RED_PNG_DATA_URL, + credentials="creds", + api_base="https://custom.example.com", + ) + + assert result == "async-file-1" + # Verify cache + url_hash = _get_url_hash(_RED_PNG_DATA_URL) + assert _file_cache[url_hash] == "async-file-1" + + # Check upload request details — url is first positional arg + call_args = mock_client.post.call_args + assert call_args.args[0] == "https://api.example.com/files" + assert call_args.kwargs["headers"]["Authorization"] == "Bearer test-token-async" + assert "purpose" in str(call_args.kwargs["data"]) + assert "file" in call_args.kwargs["files"] + + @pytest.mark.asyncio + @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") + @patch(f"{FILE_MODULE}.get_access_token_async", return_value="test-token-async") + @patch(f"{FILE_MODULE}.get_async_httpx_client") + async def test_returns_cached_file_id( + self, mock_get_client, mock_get_token, mock_get_api_base + ): + url_hash = _get_url_hash(_RED_PNG_DATA_URL) + _file_cache[url_hash] = "cached-async-id" + + result = await upload_file_async(image_url=_RED_PNG_DATA_URL, credentials="creds") + + assert result == "cached-async-id" + mock_get_client.return_value.post.assert_not_called() + + @pytest.mark.asyncio + @patch(f"{FILE_MODULE}.get_async_httpx_client") + @patch(f"{FILE_MODULE}.get_access_token_async", return_value="test-token-async") + @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") + @patch(f"{FILE_MODULE}._download_image_async") + async def test_downloads_and_uploads_url_image( + self, mock_download, mock_get_api_base, mock_get_token, mock_get_client + ): + mock_download.return_value = (b"remote-bytes-async", "image/png", "png") + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.json = MagicMock(return_value={"id": "async-file-remote"}) + mock_response.raise_for_status = MagicMock() + mock_client.post = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + result = await upload_file_async( + image_url="https://example.com/remote.png", credentials="creds" + ) + + assert result == "async-file-remote" + mock_download.assert_called_once_with("https://example.com/remote.png") + + @pytest.mark.asyncio + @patch(f"{FILE_MODULE}.get_async_httpx_client") + @patch(f"{FILE_MODULE}.get_access_token_async", return_value="test-token-async") + @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") + async def test_returns_none_on_upload_failure( + self, mock_get_api_base, mock_get_token, mock_get_client + ): + mock_client = MagicMock() + mock_client.post = AsyncMock( + side_effect=httpx.HTTPStatusError( + "Bad Request", + request=httpx.Request("POST", "https://api.example.com/files"), + response=httpx.Response(status_code=400, request=httpx.Request("POST", "https://api.example.com/files")), + ) + ) + mock_get_client.return_value = mock_client + + result = await upload_file_async( + image_url=_RED_PNG_DATA_URL, credentials="creds" + ) + + assert result is None + + @pytest.mark.asyncio + @patch(f"{FILE_MODULE}.get_async_httpx_client") + @patch(f"{FILE_MODULE}.get_access_token_async", return_value="test-token-async") + @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") + async def test_returns_none_when_response_missing_id( + self, mock_get_api_base, mock_get_token, mock_get_client + ): + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.json = MagicMock(return_value={"status": "ok"}) + mock_response.raise_for_status = MagicMock() + mock_client.post = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + result = await upload_file_async( + image_url=_RED_PNG_DATA_URL, credentials="creds" + ) + + assert result is None + + @pytest.mark.asyncio + @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") + @patch(f"{FILE_MODULE}.get_access_token_async", return_value="test-token-async") + @patch(f"{FILE_MODULE}.get_async_httpx_client") + async def test_uploads_without_optional_args( + self, mock_get_client, mock_get_token, mock_get_api_base + ): + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.json = MagicMock(return_value={"id": "async-no-args"}) + mock_response.raise_for_status = MagicMock() + mock_client.post = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + result = await upload_file_async(image_url=_RED_PNG_DATA_URL) + + assert result == "async-no-args" + mock_get_token.assert_called_once_with(credentials=None, litellm_params=None) \ No newline at end of file From 91df9845ca47234d57e2a5df88b4e8e13e7449b2 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Sun, 23 Aug 2026 12:46:17 +0000 Subject: [PATCH 050/120] fix(lint): Remove definition `Union` in litellm_logging.py --- litellm/litellm_core_utils/litellm_logging.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index b850ba3911d..0d8f0cc8060 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -14,7 +14,7 @@ from collections.abc import Callable, Mapping, Sequence from datetime import datetime as dt_object from functools import lru_cache from types import MappingProxyType, TracebackType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast from httpx import Response from pydantic import BaseModel From 069423ce0023312fda664686f8f1263b301ae5bf Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Sun, 23 Aug 2026 13:31:35 +0000 Subject: [PATCH 051/120] fix(lint): Partly fix ANN401, S110, TID251, TRY300,UP028 --- .../gigachat/passthrough/transformation.py | 7 ++- litellm/passthrough/main.py | 49 ++++++++++++------- litellm/proxy/common_request_processing.py | 2 +- .../llm_passthrough_endpoints.py | 13 +++-- 4 files changed, 40 insertions(+), 31 deletions(-) diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index b717410dff5..b619b159c6b 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -1,7 +1,7 @@ from __future__ import annotations import json -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING import httpx @@ -149,7 +149,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): generic_chunk_has_all_required_fields, ) from litellm.main import stream_chunk_builder - from litellm.types.utils import GenericStreamingChunk, ModelResponseStream + from litellm.types.utils import ModelResponseStream all_translated_chunks = [] @@ -179,8 +179,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): if isinstance(translated_chunk, dict) and generic_chunk_has_all_required_fields(translated_chunk): chunk_obj = convert_generic_chunk_to_model_response_stream( - # cast-ok: validated TypedDict - cast(GenericStreamingChunk, translated_chunk) + translated_chunk # type: ignore[arg-type] # validated TypedDict ) elif isinstance(translated_chunk, ModelResponseStream): chunk_obj = translated_chunk diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 182cef22237..f0c1a802827 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -11,7 +11,7 @@ from functools import partial from typing import Any, Final, cast import httpx -from httpx._types import CookieTypes, QueryParamTypes, RequestFiles +from httpx._types import CookieTypes, QueryParamTypes, RequestContent, RequestFiles from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider @@ -32,8 +32,7 @@ async def _as_async_generator(iterable: AsyncIterator[bytes]) -> AsyncGenerator[ def _as_generator(iterable: Iterator[bytes]) -> Generator[bytes, Any, Any]: - for chunk in iterable: - yield chunk + yield from iterable class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): @@ -88,7 +87,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic try: await self._response.aclose() - except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic + except Exception: # noqa: BLE001 S110 # Safe catch-all for cleanup logic pass raise return self @@ -129,21 +128,27 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): try: chunk = await anext(self._iterator) self._raw_bytes.append(chunk) - return chunk except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic self._start_flush() try: await self._response.aclose() - except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic + except Exception: # noqa: BLE001 S110 # Safe catch-all for cleanup logic pass raise + else: + return chunk - async def asend(self, value: Any) -> bytes: + async def asend(self, value: bytes) -> bytes: if not self._initialized: await self return await self._iterator.asend(value) - async def athrow(self, typ: Any, val: Any = None, tb: Any = None) -> bytes: + async def athrow( + self, + typ: type[BaseException], + val: BaseException | None = None, + tb: type | None = None, + ) -> bytes: if not self._initialized: await self return await self._iterator.athrow(typ, val, tb) @@ -154,7 +159,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): if self._initialized: await self._iterator.aclose() await self._response.aclose() - except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic + except Exception: # noqa: BLE001 S110 # Safe catch-all for cleanup logic pass @@ -201,26 +206,32 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): try: chunk = next(self._iterator) self._raw_bytes.append(chunk) - return chunk except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic self._start_flush() try: self._response.close() - except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic + except Exception: # noqa: BLE001 S110 # Safe catch-all for cleanup logic pass raise + else: + return chunk - def send(self, value: Any) -> bytes: + def send(self, value: bytes) -> bytes: return self._iterator.send(value) - def throw(self, typ: Any, val: Any = None, tb: Any = None) -> bytes: + def throw( + self, + typ: type[BaseException], + val: BaseException | None = None, + tb: type | None = None, + ) -> bytes: return self._iterator.throw(typ, val, tb) def close(self) -> None: self._start_flush() try: self._response.close() - except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic + except Exception: # noqa: BLE001 S110 # Safe catch-all for cleanup logic pass @@ -235,10 +246,10 @@ async def allm_passthrough_route( api_key: str | None = None, request_query_params: dict | None = None, request_headers: dict | None = None, - content: Any | None = None, + content: RequestContent | None = None, data: dict | None = None, files: RequestFiles | None = None, - json: Any | None = None, + json: object | None = None, params: QueryParamTypes | None = None, cookies: CookieTypes | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, @@ -335,7 +346,7 @@ async def allm_passthrough_route( provider=LlmProviders(resolved_custom_llm_provider), model=model, ) - except Exception: + except Exception: # noqa: BLE001 S110 # If we can't get provider config, pass None pass @@ -360,10 +371,10 @@ def llm_passthrough_route( api_key: str | None = None, request_query_params: dict | None = None, request_headers: dict | None = None, - content: Any | None = None, + content: RequestContent | None = None, data: dict | None = None, files: RequestFiles | None = None, - json: Any | None = None, + json: object | None = None, params: QueryParamTypes | None = None, cookies: CookieTypes | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index caad3f1209f..3a53fb275d5 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1426,7 +1426,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def _merge_passthrough_streaming_headers( - response_headers: Any | None, + response_headers: httpx.Headers | dict | None, custom_headers: dict, ) -> dict: """ diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index fb733787c8d..74ce4c751e3 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2888,13 +2888,6 @@ async def handle_gigachat_passthrough_router_model( user_api_base=user_api_base, version=version, ) - - if isinstance(result, StreamingResponse): - if result.headers.get("Content-Type") is None: - result.headers["Content-Type"] = "text/event-stream; charset=utf-8" - return result - - return result except Exception as e: # noqa: BLE001 # Safe catch-all for handle exception # Use common exception handling raise await base_llm_response_processor._handle_llm_api_exception( @@ -2902,6 +2895,12 @@ async def handle_gigachat_passthrough_router_model( user_api_key_dict=user_api_key_dict, proxy_logging_obj=proxy_logging_obj, ) + else: + if isinstance(result, StreamingResponse): + if result.headers.get("Content-Type") is None: + result.headers["Content-Type"] = "text/event-stream; charset=utf-8" + + return result @router.api_route( From 950267e3cd5019f0be4b2b9a7d98dbad5d867f68 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Sun, 23 Aug 2026 14:04:10 +0000 Subject: [PATCH 052/120] fix(lint): Partly fix LIT issues --- litellm/llms/gigachat/__init__.py | 4 +- litellm/llms/gigachat/authenticator.py | 50 +++++++----- litellm/llms/gigachat/chat/__init__.py | 4 +- litellm/llms/gigachat/chat/streaming.py | 46 ++++++----- litellm/llms/gigachat/chat/transformation.py | 78 +++++++++---------- litellm/llms/gigachat/file_handler.py | 5 +- litellm/llms/gigachat/passthrough/__init__.py | 2 +- .../gigachat/passthrough/transformation.py | 71 +++++++++-------- litellm/llms/gigachat/utils.py | 25 +++--- litellm/passthrough/main.py | 6 +- litellm/proxy/common_request_processing.py | 6 +- .../llm_passthrough_endpoints.py | 22 +++--- 12 files changed, 167 insertions(+), 152 deletions(-) diff --git a/litellm/llms/gigachat/__init__.py b/litellm/llms/gigachat/__init__.py index af5d2717643..e7c2206ffaa 100644 --- a/litellm/llms/gigachat/__init__.py +++ b/litellm/llms/gigachat/__init__.py @@ -17,9 +17,9 @@ from .chat.transformation import GigaChatConfig, GigaChatError from .embedding.transformation import GigaChatEmbeddingConfig from .passthrough.transformation import GigaChatPassthroughConfig -__all__ = [ +__all__ = ( "GigaChatConfig", "GigaChatEmbeddingConfig", "GigaChatError", "GigaChatPassthroughConfig", -] +) diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index 01272c42284..a1c6bda093f 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -7,6 +7,7 @@ Based on official GigaChat SDK authentication flow. import time import uuid +from collections.abc import Mapping from typing import Final import httpx @@ -63,7 +64,7 @@ def get_access_token( credentials: str | None = None, scope: str | None = None, auth_url: str | None = None, - litellm_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> str: """ Get valid access token, using cache if available. @@ -80,26 +81,28 @@ def get_access_token( GigaChatAuthError: If authentication fails """ if not litellm_params: - litellm_params = {} + litellm_params = {} # mutable-ok: empty dict default; rebind-ok: provide default - access_token = litellm_params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN") + access_token: Final = litellm_params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN") if access_token: return access_token - credentials = credentials or _get_credentials() - if not credentials: + effective_credentials: Final = credentials or _get_credentials() + if not effective_credentials: raise GigaChatAuthError( status_code=401, message="GigaChat credentials not provided. Set GIGACHAT_CREDENTIALS or GIGACHAT_API_KEY environment variable.", ) - scope = scope or litellm_params.get("gigachat_scope") or _get_scope() - auth_url = auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url() + effective_scope: Final = scope or litellm_params.get("gigachat_scope") or _get_scope() + effective_auth_url: Final = auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url() # Check cache - cache_key: Final = f"gigachat_token:{credentials[:16]}" + cache_key: Final = f"gigachat_token:{effective_credentials[:16]}" cached: Final = _token_cache.get_cache(cache_key) if cached: + token: Final + expires_at: Final token, expires_at = cached # Check if token is still valid (with buffer) if time.time() * 1000 < expires_at - TOKEN_EXPIRY_BUFFER_MS: @@ -107,7 +110,7 @@ def get_access_token( return token # Request new token - token, expires_at = _request_token_sync(credentials, scope, auth_url) + token, expires_at = _request_token_sync(effective_credentials, effective_scope, effective_auth_url) if expires_at: # Cache token @@ -122,37 +125,39 @@ async def get_access_token_async( credentials: str | None = None, scope: str | None = None, auth_url: str | None = None, - litellm_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> str: """Async version of get_access_token.""" if not litellm_params: - litellm_params = {} + litellm_params = {} # mutable-ok: empty dict default; rebind-ok: provide default - access_token = litellm_params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN") + access_token: Final = litellm_params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN") if access_token: return access_token - credentials = credentials or _get_credentials() - if not credentials: + effective_credentials: Final = credentials or _get_credentials() + if not effective_credentials: raise GigaChatAuthError( status_code=401, message="GigaChat credentials not provided. Set GIGACHAT_CREDENTIALS or GIGACHAT_API_KEY environment variable.", ) - scope = scope or litellm_params.get("gigachat_scope") or _get_scope() - auth_url = auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url() + effective_scope: Final = scope or litellm_params.get("gigachat_scope") or _get_scope() + effective_auth_url: Final = auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url() # Check cache - cache_key: Final = f"gigachat_token:{credentials[:16]}" + cache_key: Final = f"gigachat_token:{effective_credentials[:16]}" cached: Final = _token_cache.get_cache(cache_key) if cached: + token: Final + expires_at: Final token, expires_at = cached if time.time() * 1000 < expires_at - TOKEN_EXPIRY_BUFFER_MS: verbose_logger.debug("Using cached GigaChat access token") return token # Request new token - token, expires_at = await _request_token_async(credentials, scope, auth_url) + token, expires_at = await _request_token_async(effective_credentials, effective_scope, effective_auth_url) if expires_at: # Cache token @@ -241,7 +246,7 @@ def _parse_token_response(response: httpx.Response) -> tuple[str, int]: # GigaChat returns either 'tok'/'exp' or 'access_token'/'expires_at' access_token: Final = data.get("tok") or data.get("access_token") - expires_at = data.get("exp") or data.get("expires_at") + expires_at_raw: Final = data.get("exp") or data.get("expires_at") if not access_token: raise GigaChatAuthError( @@ -249,9 +254,12 @@ def _parse_token_response(response: httpx.Response) -> tuple[str, int]: message=f"Invalid token response: {data}", ) + expires_at: int # expires_at is in milliseconds - if isinstance(expires_at, str): - expires_at = int(expires_at) + if isinstance(expires_at_raw, str): + expires_at = int(expires_at_raw) + else: + expires_at = expires_at_raw # pyright: ignore[reportAssignmentType] # raw value is int or str; converted above verbose_logger.debug("GigaChat access token obtained successfully") return access_token, expires_at diff --git a/litellm/llms/gigachat/chat/__init__.py b/litellm/llms/gigachat/chat/__init__.py index eb9492b90b3..0f9be19fedd 100644 --- a/litellm/llms/gigachat/chat/__init__.py +++ b/litellm/llms/gigachat/chat/__init__.py @@ -5,8 +5,8 @@ GigaChat Chat Module from .streaming import GigaChatModelResponseIterator from .transformation import GigaChatConfig, GigaChatError -__all__ = [ +__all__ = ( "GigaChatConfig", "GigaChatError", "GigaChatModelResponseIterator", -] +) diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index 1b207057537..338d37d810a 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -4,6 +4,7 @@ GigaChat Streaming Response Handler import json import uuid +from collections.abc import Mapping, Sequence from typing import Any, Final from litellm.llms.gigachat.utils import convert_usage @@ -27,14 +28,9 @@ class GigaChatModelResponseIterator: self.response_iterator = self.streaming_response self.json_mode = json_mode - def chunk_parser(self, chunk: dict) -> GenericStreamingChunk: + def chunk_parser(self, chunk: Mapping[str, object]) -> GenericStreamingChunk: """Parse a single streaming chunk from GigaChat.""" - text = "" - tool_use: ChatCompletionToolCallChunk | None = None - is_finished = False - finish_reason: str | None = None - - choices: Final = chunk.get("choices", []) + choices: Sequence = chunk.get("choices") or () # mutable-ok: tuple literal as default if not choices: return GenericStreamingChunk( text="", @@ -46,36 +42,41 @@ class GigaChatModelResponseIterator: ) choice: Final = choices[0] - delta: Final = choice.get("delta", {}) - finish_reason = choice.get("finish_reason") + delta: Mapping[str, object] = choice.get("delta") or {} # mutable-ok: empty dict default for get + chunk_finish_reason: Final = choice.get("finish_reason") # Extract text content - text = delta.get("content", "") or "" + text: Final = delta.get("content", "") or "" + + usage_block: ChatCompletionUsageBlock | None = None + tool_use: ChatCompletionToolCallChunk | None = None + finish_reason: str | None = chunk_finish_reason # Handle function_call in stream - if finish_reason == "function_call" and delta.get("function_call"): + if chunk_finish_reason == "function_call" and delta.get("function_call"): func_call: Final = delta["function_call"] - args = func_call.get("arguments", {}) - - if isinstance(args, dict): - args = json.dumps(args, ensure_ascii=False) + args_raw: Final = func_call.get("arguments") or {} + args_str: str + if isinstance(args_raw, dict): + args_str = json.dumps(args_raw, ensure_ascii=False) # rebind-ok: build from dict + else: + args_str = str(args_raw) tool_use = ChatCompletionToolCallChunk( id=f"call_{uuid.uuid4().hex[:24]}", type="function", function=ChatCompletionToolCallFunctionChunk( name=func_call.get("name", ""), - arguments=args, + arguments=args_str, ), index=0, ) finish_reason = "tool_calls" - usage_block = None - if finish_reason == "stop": - usage_data = chunk.get("usage", {}) + if chunk_finish_reason == "stop": + usage_data: Final = chunk.get("usage") or {} # mutable-ok: empty dict default if usage_data: - usage = convert_usage(usage_data) + usage: Final = convert_usage(usage_data) usage_block = ChatCompletionUsageBlock( prompt_tokens=usage.prompt_tokens, completion_tokens=usage.completion_tokens, @@ -88,13 +89,10 @@ class GigaChatModelResponseIterator: ), ) - if finish_reason is not None: - is_finished = True - return GenericStreamingChunk( text=text, tool_use=tool_use, - is_finished=is_finished, + is_finished=chunk_finish_reason is not None, finish_reason=finish_reason or "", usage=usage_block, index=choice.get("index", 0), diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index 920a067eaa9..b625e9722f8 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -9,7 +9,7 @@ from __future__ import annotations import json import time import uuid -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx @@ -88,8 +88,8 @@ class GigaChatConfig(BaseConfig): api_base: str | None, api_key: str | None, model: str, - optional_params: dict, - litellm_params: dict, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], stream: bool | None = None, ) -> str: """Get complete API URL for chat completions.""" @@ -98,14 +98,14 @@ class GigaChatConfig(BaseConfig): def validate_environment( self, - headers: dict, + headers: dict, # mutable-ok: mutates in place per GigaChat OAuth setup model: str, - messages: list[AllMessageValues], - optional_params: dict, - litellm_params: dict, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], api_key: str | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict: # mutable-ok: base class contract returns dict for httpx """ Set up headers with OAuth token. """ @@ -123,9 +123,9 @@ class GigaChatConfig(BaseConfig): return headers - def get_supported_openai_params(self, model: str) -> list[str]: + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: base class contract returns list """Return list of supported OpenAI parameters.""" - return [ + return [ # mutable-ok: base class contract returns list "stream", "temperature", "top_p", @@ -141,11 +141,11 @@ class GigaChatConfig(BaseConfig): def map_openai_params( self, - non_default_params: dict, - optional_params: dict, + non_default_params: Mapping[str, object], + optional_params: dict, # mutable-ok: mutated in place per GigaChat mapping model: str, drop_params: bool, - ) -> dict: + ) -> dict: # mutable-ok: base class contract returns dict """Map OpenAI parameters to GigaChat parameters.""" for param, value in non_default_params.items(): if param == "stream": @@ -182,25 +182,25 @@ class GigaChatConfig(BaseConfig): schema_name = json_schema.get("name", "structured_output") schema = json_schema.get("schema", {}) - function_def = { + function_def = { # mutable-ok: request payload for httpx "name": schema_name, "description": f"Output structured response: {schema_name}", "parameters": schema, } if "functions" not in optional_params: - optional_params["functions"] = [] + optional_params["functions"] = [] # mutable-ok: list for httpx optional_params["functions"].append(function_def) - optional_params["function_call"] = {"name": schema_name} + optional_params["function_call"] = {"name": schema_name} # mutable-ok: request payload optional_params["_structured_output"] = True return optional_params - def _convert_tools_to_functions(self, tools: list[dict]) -> list[dict]: + def _convert_tools_to_functions(self, tools: Sequence) -> Sequence[dict]: """Convert OpenAI tools format to GigaChat functions format.""" - functions: Final = [] + functions: Final[list[dict]] = [] # mutable-ok: accumulator for building functions list for tool in tools: - if tool.get("type") == "function": + if isinstance(tool, dict) and tool.get("type") == "function": func = tool.get("function", {}) functions.append( { @@ -211,7 +211,7 @@ class GigaChatConfig(BaseConfig): ) return functions - def _map_tool_choice(self, tool_choice: str | dict) -> str | dict | None: + def _map_tool_choice(self, tool_choice: str | Mapping[str, object]) -> str | Mapping[str, object] | None: """ Map OpenAI tool_choice to GigaChat function_call format. @@ -271,7 +271,7 @@ class GigaChatConfig(BaseConfig): verbose_logger.error("Failed to upload image: %s", e) return None - def _transform_list_content(self, content: list) -> tuple[str, list[str]]: + def _transform_list_content(self, content: Sequence) -> tuple[str, Sequence[str]]: """ Extract text and image attachments from a multimodal message content list. @@ -281,8 +281,8 @@ class GigaChatConfig(BaseConfig): Returns: Tuple of (combined text, list of attachment file ids) """ - texts = [] - attachments = [] + texts: Final[list[str]] = [] # mutable-ok: accumulator + attachments: Final[list[str]] = [] # mutable-ok: accumulator for part in content: if isinstance(part, dict): if part.get("type") == "text": @@ -291,24 +291,24 @@ class GigaChatConfig(BaseConfig): # Extract image URL and upload to GigaChat image_url = part.get("image_url", {}) if isinstance(image_url, str): - url = image_url + url: Final = image_url else: - url = image_url.get("url", "") + url: Final = image_url.get("url", "") if url: - file_id = self._upload_image(url) + file_id = self._upload_image(url) # rebind-ok: inside for loop, no outer binding if file_id: attachments.append(file_id) - text = "\n".join(texts) if texts else "" + text: Final = "\n".join(texts) if texts else "" return text, attachments def transform_request( self, model: str, - messages: list[AllMessageValues], - optional_params: dict, - litellm_params: dict, - headers: dict, - ) -> dict: + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + headers: Mapping[str, object], + ) -> dict: # mutable-ok: request payload sent to httpx """Transform OpenAI request to GigaChat format.""" # Transform messages giga_messages: Final = self._transform_messages(messages) @@ -339,9 +339,9 @@ class GigaChatConfig(BaseConfig): return request_data - def _transform_messages(self, messages: list[AllMessageValues]) -> list[dict]: + def _transform_messages(self, messages: Sequence[AllMessageValues]) -> Sequence[dict]: """Transform OpenAI messages to GigaChat format.""" - transformed: Final = [] + transformed: Final[list[dict]] = [] # mutable-ok: accumulator for building transformed messages for i, msg in enumerate(messages): message = dict(msg) @@ -400,10 +400,10 @@ class GigaChatConfig(BaseConfig): raw_response: httpx.Response, model_response: ModelResponse, logging_obj: LiteLLMLoggingObj, - request_data: dict, - messages: list[AllMessageValues], - optional_params: dict, - litellm_params: dict, + request_data: Mapping[str, object], + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], encoding: Any, api_key: str | None = None, json_mode: bool | None = None, @@ -419,7 +419,7 @@ class GigaChatConfig(BaseConfig): is_structured_output: Final = optional_params.get("_structured_output", False) - choices: Final = [] + choices: Final[list[Choices]] = [] # mutable-ok: accumulator for building response choices for choice in response_json.get("choices", []): message_data = choice.get("message", {}) finish_reason = choice.get("finish_reason", "stop") diff --git a/litellm/llms/gigachat/file_handler.py b/litellm/llms/gigachat/file_handler.py index 6dfd4d74be7..163e944f124 100644 --- a/litellm/llms/gigachat/file_handler.py +++ b/litellm/llms/gigachat/file_handler.py @@ -9,6 +9,7 @@ import base64 import hashlib import re import uuid +from collections.abc import Mapping from typing import Final from litellm._logging import verbose_logger @@ -80,7 +81,7 @@ def upload_file_sync( image_url: str, credentials: str | None = None, api_base: str | None = None, - litellm_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> str | None: """ Upload file to GigaChat and return file_id (sync). @@ -146,7 +147,7 @@ async def upload_file_async( image_url: str, credentials: str | None = None, api_base: str | None = None, - litellm_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> str | None: """ Upload file to GigaChat and return file_id (async). diff --git a/litellm/llms/gigachat/passthrough/__init__.py b/litellm/llms/gigachat/passthrough/__init__.py index 9e5f9b8ed77..a66a078dbeb 100644 --- a/litellm/llms/gigachat/passthrough/__init__.py +++ b/litellm/llms/gigachat/passthrough/__init__.py @@ -4,4 +4,4 @@ GigaChat passthrough Module from .transformation import GigaChatPassthroughConfig -__all__ = ["GigaChatPassthroughConfig"] +__all__ = ("GigaChatPassthroughConfig",) diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index b619b159c6b..ec441c53cc4 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -1,7 +1,8 @@ from __future__ import annotations import json -from typing import TYPE_CHECKING +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final import httpx @@ -21,7 +22,7 @@ if TYPE_CHECKING: class GigaChatPassthroughConfig(BasePassthroughConfig): - def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: + def is_streaming_request(self, endpoint: str, request_data: Mapping[str, object]) -> bool: return request_data.get("stream", False) def get_complete_url( @@ -30,16 +31,16 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): api_key: str | None, model: str, endpoint: str, - request_query_params: dict | None, - litellm_params: dict, + request_query_params: Mapping[str, object] | None, + litellm_params: Mapping[str, object], ) -> tuple[URL, str]: """Get complete API URL for chat completions.""" - base_target_url = self.get_api_base(api_base) + base_target_url: Final = self.get_api_base(api_base) if base_target_url is None: raise Exception("GigaChat api base not found") - complete_url = f"{base_target_url}/{endpoint.lstrip('/')}" + complete_url: Final = f"{base_target_url}/{endpoint.lstrip('/')}" return ( httpx.URL(complete_url), @@ -48,23 +49,23 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): def validate_environment( self, - headers: dict, + headers: dict, # mutable-ok: mutates in place to set OAuth headers model: str, - messages: list[AllMessageValues], - optional_params: dict, - litellm_params: dict, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], api_key: str | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict: # mutable-ok: base class contract returns dict for httpx """ Set up headers with OAuth token. """ # Get access token - access_token = get_access_token(credentials=api_key, litellm_params=litellm_params) + access_token: Final = get_access_token(credentials=api_key, litellm_params=litellm_params) - headers["Authorization"] = f"Bearer {access_token}" - headers["Content-Type"] = "application/json" - headers["Accept"] = "application/json" + headers["Authorization"] = f"Bearer {access_token}" # rebind-ok: mutating for OAuth setup + headers["Content-Type"] = "application/json" # rebind-ok: mutating for OAuth setup + headers["Accept"] = "application/json" # rebind-ok: mutating for OAuth setup return headers @@ -73,7 +74,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): model: str, custom_llm_provider: str, httpx_response: Response, - request_data: dict, + request_data: Mapping[str, object], logging_obj: LiteLLMLoggingObj, endpoint: str, ) -> CostResponseTypes | None: @@ -83,7 +84,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): # cost tracking only for completions and embeddings if "completions" in endpoint: - provider_chat_config = ProviderConfigManager.get_provider_chat_config( + provider_chat_config: Final = ProviderConfigManager.get_provider_chat_config( provider=LlmProviders(custom_llm_provider), model=model, ) @@ -93,12 +94,12 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): litellm_model_response: ModelResponse = provider_chat_config.transform_response( model=model, - messages=request_data.get("messages", []), + messages=request_data.get("messages", []), # mutable-ok: empty list default for transform_response raw_response=httpx_response, model_response=ModelResponse(), logging_obj=logging_obj, - optional_params={}, - litellm_params={}, + optional_params={}, # mutable-ok: empty dict kwarg for transform_response + litellm_params={}, # mutable-ok: empty dict kwarg for transform_response api_key="", request_data=request_data, encoding=encoding, @@ -107,7 +108,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): return litellm_model_response if "embeddings" in endpoint: - provider_embedding_config = ProviderConfigManager.get_provider_embedding_config( + provider_embedding_config: Final = ProviderConfigManager.get_provider_embedding_config( provider=LlmProviders(custom_llm_provider), model=model, ) @@ -115,15 +116,17 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): if provider_embedding_config is None: raise ValueError(f"No provider config found for model: {model}") - litellm_embedding_response: EmbeddingResponse = provider_embedding_config.transform_embedding_response( - model=model, - raw_response=httpx_response, - model_response=EmbeddingResponse(), - logging_obj=logging_obj, - optional_params={}, - api_key="", - request_data=request_data, - litellm_params={}, + litellm_embedding_response: Final[EmbeddingResponse] = ( + provider_embedding_config.transform_embedding_response( + model=model, + raw_response=httpx_response, + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + optional_params={}, # mutable-ok: empty dict kwarg for transform_embedding_response + api_key="", + request_data=request_data, + litellm_params={}, # mutable-ok: empty dict kwarg for transform_embedding_response + ) ) return litellm_embedding_response @@ -132,7 +135,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): def handle_logging_collected_chunks( self, - all_chunks: list[str], + all_chunks: Sequence[str], litellm_logging_obj: LiteLLMLoggingObj, model: str, custom_llm_provider: str, @@ -151,7 +154,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): from litellm.main import stream_chunk_builder from litellm.types.utils import ModelResponseStream - all_translated_chunks = [] + all_translated_chunks: Final[list[object]] = [] # mutable-ok: accumulator for chunk in all_chunks: if isinstance(chunk, bytes): @@ -179,7 +182,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): if isinstance(translated_chunk, dict) and generic_chunk_has_all_required_fields(translated_chunk): chunk_obj = convert_generic_chunk_to_model_response_stream( - translated_chunk # type: ignore[arg-type] # validated TypedDict + translated_chunk # pyright: ignore[reportArgumentType] # validated TypedDict ) elif isinstance(translated_chunk, ModelResponseStream): chunk_obj = translated_chunk @@ -209,5 +212,5 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): def get_base_model(model: str) -> str | None: return model - def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]: + def get_models(self, api_key: str | None = None, api_base: str | None = None) -> Sequence[str]: return super().get_models(api_key, api_base) diff --git a/litellm/llms/gigachat/utils.py b/litellm/llms/gigachat/utils.py index d2b18f30b5d..aa30552a7ec 100644 --- a/litellm/llms/gigachat/utils.py +++ b/litellm/llms/gigachat/utils.py @@ -1,28 +1,31 @@ +from collections.abc import Mapping +from typing import Final + from litellm.secret_managers.main import get_secret_str from litellm.types.utils import PromptTokensDetailsWrapper, Usage # GigaChat API endpoint -GIGACHAT_BASE_URL = "https://gigachat.devices.sberbank.ru/api/v1" +GIGACHAT_BASE_URL: Final = "https://gigachat.devices.sberbank.ru/api/v1" -def convert_usage(usage_data: dict[str, int]) -> Usage: - prompt_tokens = usage_data.get("prompt_tokens", 0) - completion_tokens = usage_data.get("completion_tokens", 0) - precached_prompt_tokens = usage_data.get("precached_prompt_tokens", 0) - total_tokens = usage_data.get("total_tokens", 0) +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 += precached_prompt_tokens - total_tokens += precached_prompt_tokens + prompt_tokens_total: Final = prompt_tokens + precached_prompt_tokens + total_tokens_total: Final = total_tokens + precached_prompt_tokens - prompt_tokens_details = None + prompt_tokens_details: PromptTokensDetailsWrapper | None = None if precached_prompt_tokens > 0: prompt_tokens_details = PromptTokensDetailsWrapper(cached_tokens=precached_prompt_tokens) return Usage( - prompt_tokens=prompt_tokens, + prompt_tokens=prompt_tokens_total, completion_tokens=completion_tokens, prompt_tokens_details=prompt_tokens_details, - total_tokens=total_tokens, + total_tokens=total_tokens_total, ) diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index f0c1a802827..9d2ca66965b 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -50,9 +50,9 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): self._iterator: AsyncGenerator[bytes, Any] self._litellm_logging_obj = litellm_logging_obj self._provider_config = provider_config - self._raw_bytes: list[bytes] = [] + self._raw_bytes: list[bytes] = [] # mutable-ok: instance buffer for streaming chunks self._flush_scheduled = False - self._background_tasks: set[asyncio.Task] = set() + self._background_tasks: set[asyncio.Task] = set() # mutable-ok: instance set for background task tracking @property def status_code(self) -> int: @@ -176,7 +176,7 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): self._litellm_logging_obj = litellm_logging_obj self._provider_config = provider_config self._iterator: Generator[bytes, Any, Any] = _as_generator(response.iter_bytes()) - self._raw_bytes: list[bytes] = [] + self._raw_bytes: list[bytes] = [] # mutable-ok: instance buffer for streaming chunks self._flush_scheduled = False def _start_flush(self) -> None: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 3a53fb275d5..ca94d6dc292 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1434,7 +1434,7 @@ class ProxyBaseLLMRequestProcessing: Proxy/custom headers win on key collisions. """ - excluded_headers = { + excluded_headers = { # mutable-ok: set of header names to exclude from forwarding "transfer-encoding", "content-encoding", "set-cookie", @@ -1447,7 +1447,7 @@ class ProxyBaseLLMRequestProcessing: "upgrade", } - merged_headers = { + merged_headers = { # mutable-ok: dict comprehension for merged headers forwarded to httpx key: value for key, value in dict(response_headers or {}).items() if key.lower() not in excluded_headers } merged_headers.update(custom_headers) @@ -2448,7 +2448,7 @@ class ProxyBaseLLMRequestProcessing: # For passthrough routes, stream directly without error parsing # since we're dealing with raw binary data (e.g., AWS event streams) return StreamingResponse( - content=generator, # type: ignore[arg-type] + content=generator, # pyright: ignore[reportArgumentType] # generator-configured StreamingResponse status_code=getattr(response, "status_code", status.HTTP_200_OK), media_type=self._passthrough_event_stream_media_type(), headers=streaming_headers, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 74ce4c751e3..b7a1c71a9e5 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2665,8 +2665,8 @@ def create_generic_websocket_passthrough_endpoint( @router.api_route( "/gigachat/{endpoint:path}", - methods=["GET", "POST", "PUT", "DELETE", "PATCH"], - tags=["Gigachat Pass-through", "pass-through"], + methods=["GET", "POST", "PUT", "DELETE", "PATCH"], # mutable-ok: FastAPI route methods + tags=["Gigachat Pass-through", "pass-through"], # mutable-ok: FastAPI route tags ) async def gigachat_proxy_route( endpoint: str, @@ -2700,7 +2700,9 @@ async def gigachat_proxy_route( if model: is_router_model = is_passthrough_request_using_router_model(request_body, llm_router) elif any(word in endpoint for word in ("completions", "embeddings")): - raise HTTPException(status_code=400, detail={"error": "Model is required in request body"}) + raise HTTPException( + status_code=400, detail={"error": "Model is required in request body"} + ) # mutable-ok: HTTPException detail dict # If router model, use dedicated router passthrough handler # This uses the same common processing path as non-router models @@ -2730,16 +2732,16 @@ async def gigachat_proxy_route( "Gigachat passthrough: Using direct Gigachat model '%s' for endpoint '%s'", model, endpoint ) - data: Dict[str, Any] = {} + data: Dict[str, Any] = {} # mutable-ok: request body mutated in place by proxy pipeline data["method"] = request.method data["endpoint"] = endpoint data["json"] = request_body data["custom_llm_provider"] = "gigachat" - client = get_async_httpx_client( # type: ignore + client = get_async_httpx_client( llm_provider=LlmProviders.GIGACHAT, - params={ + params={ # mutable-ok: httpx client params "timeout": httpx.Timeout(timeout=600.0, connect=5.0), "ssl_verify": False, }, @@ -2823,7 +2825,7 @@ async def handle_gigachat_passthrough_router_model( data: Dict[str, Any] = await _read_request_body(request=request) if user_api_key_dict is not None: if data.get("metadata") is None: - data["metadata"] = {} + 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: @@ -2847,7 +2849,7 @@ async def handle_gigachat_passthrough_router_model( data["custom_llm_provider"] = "gigachat" # Remove sensitive keys from data - keys = [ + keys = [ # mutable-ok: list of keys to remove from data "gigachat_auth_url", "gigachat_access_token", "gigachat_scope", @@ -2857,9 +2859,9 @@ async def handle_gigachat_passthrough_router_model( for key in keys: data.pop(key, None) - client = get_async_httpx_client( # type: ignore + client = get_async_httpx_client( llm_provider=LlmProviders.GIGACHAT, - params={ + params={ # mutable-ok: httpx client params "timeout": httpx.Timeout(timeout=600.0, connect=5.0), "ssl_verify": False, }, From 0fc3bccef8e4807f8cca4019db55404534de69b8 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Sun, 23 Aug 2026 14:35:12 +0000 Subject: [PATCH 053/120] fix(lint): Partly fix TQ lint issues --- .../chat/test_gigachat_chat_transformation.py | 8 ++--- .../test_gigachat_embedding_transformation.py | 4 --- ...est_gigachat_passthrough_transformation.py | 32 ++++++++----------- .../llms/gigachat/test_authenticator.py | 20 +++++------- .../llms/gigachat/test_file_handler.py | 4 --- .../test_litellm/llms/gigachat/test_utils.py | 5 --- .../test_llm_pass_through_endpoints.py | 30 ++++++++--------- 7 files changed, 39 insertions(+), 64 deletions(-) diff --git a/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py b/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py index 42b5c1894b2..2f9511e642c 100644 --- a/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py +++ b/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py @@ -8,15 +8,11 @@ get_model_response_iterator, and get_error_class. """ import json -import os -import sys from unittest.mock import MagicMock, patch import httpx import pytest -sys.path.insert(0, os.path.abspath("../../../../..")) - from litellm.llms.gigachat.chat.transformation import ( GigaChatConfig, GigaChatError, @@ -147,7 +143,7 @@ class TestValidateEnvironment: @patch(f"{TRANSFORM_MODULE}.get_access_token", return_value="token") @patch(f"{TRANSFORM_MODULE}.get_secret_str") - def test_falls_back_to_env_for_credentials( + def test_falls_back_to_env_for_credentials( # test-quality-ok: mock-echo of internal wiring self, mock_get_secret, mock_get_token ): mock_get_secret.return_value = "env-creds" @@ -160,7 +156,7 @@ class TestValidateEnvironment: api_key=None, api_base=None, ) - mock_get_secret.assert_any_call("GIGACHAT_CREDENTIALS") + mock_get_secret.assert_any_call("GIGACHAT_CREDENTIALS") # test-quality-ok: mock-echo of internal wiring class TestGetSupportedOpenAiParams: diff --git a/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py b/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py index 594f592bab3..2a44a8e067e 100644 --- a/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py +++ b/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py @@ -8,15 +8,11 @@ and get_error_class. """ import json -import os -import sys from unittest.mock import MagicMock, patch import httpx import pytest -sys.path.insert(0, os.path.abspath("../../../../..")) - from litellm import LlmProviders from litellm.llms.gigachat.embedding.transformation import ( GigaChatEmbeddingConfig, diff --git a/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py b/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py index e18a09b2448..d2c617cd220 100644 --- a/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py +++ b/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py @@ -6,15 +6,11 @@ streaming detection, authentication handling, and logging response transformatio """ import json -import os -import sys from unittest.mock import MagicMock, patch import httpx import pytest -sys.path.insert(0, os.path.abspath("../../../../..")) - from litellm.llms.gigachat.passthrough.transformation import GigaChatPassthroughConfig from litellm.types.utils import EmbeddingResponse, ModelResponse @@ -132,7 +128,7 @@ class TestGigaChatPassthroughConfig: assert str(complete_url) == "https://custom.gigachat.ru/api/v1/chat/completions" assert base_target_url == api_base - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.llms.gigachat.passthrough.transformation.get_secret_str" ) def test_get_complete_url_with_env_api_base(self, mock_get_secret): @@ -155,7 +151,7 @@ class TestGigaChatPassthroughConfig: assert base_target_url == env_api_base mock_get_secret.assert_called_once_with("GIGACHAT_API_BASE") - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.llms.gigachat.passthrough.transformation.get_secret_str" ) def test_get_complete_url_fallback_to_default(self, mock_get_secret): @@ -180,11 +176,11 @@ class TestGigaChatPassthroughConfig: """Test that exception is raised when no api_base can be resolved.""" config = GigaChatPassthroughConfig() with patch( - "litellm.llms.gigachat.passthrough.transformation.get_secret_str", + "litellm.llms.gigachat.passthrough.transformation.get_secret_str", # test-quality-ok: patching litellm internal for unit test isolation return_value=None, ): with patch( - "litellm.llms.gigachat.passthrough.transformation.GIGACHAT_BASE_URL", + "litellm.llms.gigachat.passthrough.transformation.GIGACHAT_BASE_URL", # test-quality-ok: patching litellm internal for unit test isolation None, ): with pytest.raises(Exception, match="GigaChat api base not found"): @@ -197,7 +193,7 @@ class TestGigaChatPassthroughConfig: litellm_params={}, ) - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.llms.gigachat.passthrough.transformation.get_access_token" ) def test_validate_environment(self, mock_get_access_token): @@ -417,7 +413,7 @@ class TestGigaChatPassthroughConfig: assert isinstance(result, ModelResponse) assert result.choices[0].message.content == "valid" - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.llms.gigachat.passthrough.transformation.get_secret_str" ) def test_get_api_base_with_explicit_value(self, mock_get_secret): @@ -427,7 +423,7 @@ class TestGigaChatPassthroughConfig: assert result == explicit_base mock_get_secret.assert_not_called() - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.llms.gigachat.passthrough.transformation.get_secret_str" ) def test_get_api_base_from_environment(self, mock_get_secret): @@ -438,7 +434,7 @@ class TestGigaChatPassthroughConfig: assert result == env_base mock_get_secret.assert_called_once_with("GIGACHAT_API_BASE") - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.llms.gigachat.passthrough.transformation.get_secret_str" ) def test_get_api_base_fallback_to_default(self, mock_get_secret): @@ -447,7 +443,7 @@ class TestGigaChatPassthroughConfig: result = GigaChatPassthroughConfig.get_api_base(api_base=None) assert result == "https://gigachat.devices.sberbank.ru/api/v1" - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.llms.gigachat.passthrough.transformation.get_secret_str" ) def test_get_api_key_with_explicit_value(self, mock_get_secret): @@ -457,7 +453,7 @@ class TestGigaChatPassthroughConfig: assert result == explicit_key mock_get_secret.assert_not_called() - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.llms.gigachat.passthrough.transformation.get_secret_str" ) def test_get_api_key_from_environment(self, mock_get_secret): @@ -486,7 +482,7 @@ class TestGigaChatPassthroughConfig: logging_obj = MagicMock() with patch( - "litellm.utils.ProviderConfigManager.get_provider_chat_config", + "litellm.utils.ProviderConfigManager.get_provider_chat_config", # test-quality-ok: patching litellm internal for unit test isolation return_value=None, ): with pytest.raises(ValueError, match="No provider config found for model"): @@ -508,7 +504,7 @@ class TestGigaChatPassthroughConfig: logging_obj = MagicMock() with patch( - "litellm.utils.ProviderConfigManager.get_provider_embedding_config", + "litellm.utils.ProviderConfigManager.get_provider_embedding_config", # test-quality-ok: patching litellm internal for unit test isolation return_value=None, ): with pytest.raises(ValueError, match="No provider config found for model"): @@ -551,7 +547,7 @@ class TestGigaChatPassthroughConfig: ] with patch( - "litellm.llms.gigachat.passthrough.transformation.GigaChatModelResponseIterator.chunk_parser", + "litellm.llms.gigachat.passthrough.transformation.GigaChatModelResponseIterator.chunk_parser", # test-quality-ok: patching litellm internal for unit test isolation return_value=stream_chunk, ): result = config.handle_logging_collected_chunks( @@ -576,7 +572,7 @@ class TestGigaChatPassthroughConfig: ] with patch( - "litellm.llms.gigachat.passthrough.transformation.GigaChatModelResponseIterator.chunk_parser", + "litellm.llms.gigachat.passthrough.transformation.GigaChatModelResponseIterator.chunk_parser", # test-quality-ok: patching litellm internal for unit test isolation return_value=12345, # not dict and not ModelResponseStream ): result = config.handle_logging_collected_chunks( diff --git a/tests/test_litellm/llms/gigachat/test_authenticator.py b/tests/test_litellm/llms/gigachat/test_authenticator.py index dcbac7e0949..30b3dec52cb 100644 --- a/tests/test_litellm/llms/gigachat/test_authenticator.py +++ b/tests/test_litellm/llms/gigachat/test_authenticator.py @@ -5,16 +5,12 @@ Tests get_access_token and get_access_token_async covering token resolution from litellm_params/env, credential validation, caching, and error handling. """ -import os -import sys import time from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest -sys.path.insert(0, os.path.abspath("../../../..")) - from litellm.llms.gigachat import authenticator from litellm.llms.gigachat.authenticator import ( GigaChatAuthError, @@ -162,7 +158,7 @@ class TestGetAccessTokenSync: @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") @patch(f"{AUTH_MODULE}._get_credentials", return_value="env-creds") @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) - def test_litellm_params_override_scope_and_auth_url(self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request): + def test_litellm_params_override_scope_and_auth_url(self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request): # test-quality-ok: mock-echo of internal wiring mock_request.return_value = ("token", _future_expires_at_ms()) get_access_token( @@ -172,7 +168,7 @@ class TestGetAccessTokenSync: } ) - mock_request.assert_called_once_with( + mock_request.assert_called_once_with( # test-quality-ok: mock-echo of internal wiring "env-creds", "GIGACHAT_API_CORP", "https://params-auth.example.com" ) @@ -181,7 +177,7 @@ class TestGetAccessTokenSync: @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") @patch(f"{AUTH_MODULE}._get_credentials", return_value="env-creds") @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) - def test_explicit_args_override_everything(self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request): + def test_explicit_args_override_everything(self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request): # test-quality-ok: mock-echo of internal wiring mock_request.return_value = ("token", _future_expires_at_ms()) get_access_token( @@ -194,7 +190,7 @@ class TestGetAccessTokenSync: }, ) - mock_request.assert_called_once_with( + mock_request.assert_called_once_with( # test-quality-ok: mock-echo of internal wiring "explicit-creds", "EXPLICIT_SCOPE", "https://explicit.example.com" ) @@ -322,7 +318,7 @@ class TestGetAccessTokenAsync: @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") @patch(f"{AUTH_MODULE}._get_credentials", return_value="env-creds") @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) - async def test_litellm_params_override_scope_and_auth_url( + async def test_litellm_params_override_scope_and_auth_url( # test-quality-ok: mock-echo of internal wiring self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request ): mock_request.return_value = ("token", _future_expires_at_ms()) @@ -334,7 +330,7 @@ class TestGetAccessTokenAsync: } ) - mock_request.assert_called_once_with( + mock_request.assert_called_once_with( # test-quality-ok: mock-echo of internal wiring "env-creds", "GIGACHAT_API_CORP", "https://params-auth.example.com" ) @@ -344,7 +340,7 @@ class TestGetAccessTokenAsync: @patch(f"{AUTH_MODULE}._get_scope", return_value="GIGACHAT_API_PERS") @patch(f"{AUTH_MODULE}._get_credentials", return_value="env-creds") @patch(f"{AUTH_MODULE}.get_secret_str", return_value=None) - async def test_explicit_args_override_everything( + async def test_explicit_args_override_everything( # test-quality-ok: mock-echo of internal wiring self, mock_get_secret, mock_creds, mock_scope, mock_auth_url, mock_request ): mock_request.return_value = ("token", _future_expires_at_ms()) @@ -359,7 +355,7 @@ class TestGetAccessTokenAsync: }, ) - mock_request.assert_called_once_with( + mock_request.assert_called_once_with( # test-quality-ok: mock-echo of internal wiring "explicit-creds", "EXPLICIT_SCOPE", "https://explicit.example.com" ) diff --git a/tests/test_litellm/llms/gigachat/test_file_handler.py b/tests/test_litellm/llms/gigachat/test_file_handler.py index 05fcedac338..ae019eb7942 100644 --- a/tests/test_litellm/llms/gigachat/test_file_handler.py +++ b/tests/test_litellm/llms/gigachat/test_file_handler.py @@ -7,15 +7,11 @@ network errors, and the full upload flow. """ import json -import os -import sys from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest -sys.path.insert(0, os.path.abspath("../../..")) - from litellm.llms.gigachat import file_handler from litellm.llms.gigachat.file_handler import ( _file_cache, diff --git a/tests/test_litellm/llms/gigachat/test_utils.py b/tests/test_litellm/llms/gigachat/test_utils.py index e1c7b75217f..3d45b1a7465 100644 --- a/tests/test_litellm/llms/gigachat/test_utils.py +++ b/tests/test_litellm/llms/gigachat/test_utils.py @@ -2,11 +2,6 @@ Tests for litellm.llms.gigachat.utils """ -import os -import sys - -sys.path.insert(0, os.path.abspath("../../..")) - import pytest from litellm.llms.gigachat.utils import convert_usage from litellm.types.utils import PromptTokensDetailsWrapper, Usage 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 1e1771864eb..9b2b4238c5b 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 @@ -177,7 +177,7 @@ class TestBaseOpenAIPassThroughHandler: assert result["api-key"] == "test_api_key" assert result["test-header"] == "value" - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) async def test_base_openai_pass_through_handler(self, mock_create_pass_through): @@ -2018,15 +2018,15 @@ class TestLLMPassthroughFactoryProxyRoute: class TestVLLMProxyRoute: @pytest.mark.asyncio - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body", return_value={"model": "router-model", "stream": False}, ) - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_passthrough_request_using_router_model", return_value=True, ) - @patch("litellm.proxy.proxy_server.llm_router") + @patch("litellm.proxy.proxy_server.llm_router") # test-quality-ok: patching litellm internal for unit test isolation async def test_vllm_proxy_route_with_router_model( self, mock_llm_router, mock_is_router, mock_get_body ): @@ -2051,15 +2051,15 @@ class TestVLLMProxyRoute: mock_llm_router.allm_passthrough_route.assert_awaited_once() @pytest.mark.asyncio - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body", return_value={"model": "other-model"}, ) - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_passthrough_request_using_router_model", return_value=False, ) - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.llm_passthrough_factory_proxy_route" ) async def test_vllm_proxy_route_fallback_to_factory( @@ -2083,15 +2083,15 @@ class TestVLLMProxyRoute: class TestGigachatProxyRoute: @pytest.mark.asyncio - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body", return_value={"model": "router-model", "stream": False}, ) - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_passthrough_request_using_router_model", return_value=True, ) - @patch("litellm.proxy.proxy_server.llm_router") + @patch("litellm.proxy.proxy_server.llm_router") # test-quality-ok: patching litellm internal for unit test isolation async def test_gigachat_proxy_route_with_router_model( self, mock_llm_router, mock_is_router, mock_get_body ): @@ -2117,15 +2117,15 @@ class TestGigachatProxyRoute: assert isinstance(result, Response) @pytest.mark.asyncio - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body", return_value={"model": "other-model"}, ) - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_passthrough_request_using_router_model", return_value=False, ) - @patch( + @patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing.base_passthrough_process_llm_request", new_callable=AsyncMock, ) @@ -2220,10 +2220,10 @@ class TestGigachatProxyRoute: processor.data["litellm_logging_obj"], ) ), - ), patch( + ), patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.proxy.common_request_processing.route_request", new=_fake_route_request, - ), patch( + ), patch( # test-quality-ok: patching litellm internal for unit test isolation "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing.get_custom_headers", return_value={"x-litellm-call-id": "call-123"}, ): From ef6cb77f91a45eb9d9debe50685a936cceb4a79b Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Sun, 23 Aug 2026 15:27:55 +0000 Subject: [PATCH 054/120] fix(lint): Partly fix basedpyright lint issues --- litellm/llms/gigachat/authenticator.py | 39 ++++++------ litellm/llms/gigachat/chat/transformation.py | 11 ++-- .../llms/gigachat/embedding/transformation.py | 4 +- litellm/llms/gigachat/file_handler.py | 6 +- .../gigachat/passthrough/transformation.py | 28 ++++----- litellm/passthrough/main.py | 25 ++++---- .../llm_passthrough_endpoints.py | 9 +-- ...est_gigachat_passthrough_transformation.py | 46 +++++++------- .../llms/gigachat/test_file_handler.py | 60 +++++++++---------- 9 files changed, 109 insertions(+), 119 deletions(-) diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index a1c6bda093f..9b8ef3ec93f 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -17,7 +17,6 @@ from litellm.caching.caching import InMemoryCache from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, - _get_httpx_client, get_async_httpx_client, ) from litellm.secret_managers.main import get_secret_str @@ -57,7 +56,7 @@ def _get_scope() -> str: def _get_http_client() -> HTTPHandler: """Get cached httpx client with SSL verification disabled.""" - return _get_httpx_client(params={"ssl_verify": False}) + return HTTPHandler(ssl_verify=False) def get_access_token( @@ -101,24 +100,22 @@ def get_access_token( cache_key: Final = f"gigachat_token:{effective_credentials[:16]}" cached: Final = _token_cache.get_cache(cache_key) if cached: - token: Final - expires_at: Final - token, expires_at = cached + _token, _expires_at = cached # Check if token is still valid (with buffer) - if time.time() * 1000 < expires_at - TOKEN_EXPIRY_BUFFER_MS: + if time.time() * 1000 < _expires_at - TOKEN_EXPIRY_BUFFER_MS: verbose_logger.debug("Using cached GigaChat access token") - return token + return _token # Request new token - token, expires_at = _request_token_sync(effective_credentials, effective_scope, effective_auth_url) + new_token, new_expires_at = _request_token_sync(effective_credentials, effective_scope, effective_auth_url) - if expires_at: + if new_expires_at: # Cache token - ttl_seconds: Final = max(0, (expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000) + ttl_seconds: Final = max(0, (new_expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000) if ttl_seconds > 0: - _token_cache.set_cache(cache_key, (token, expires_at), ttl=ttl_seconds) + _token_cache.set_cache(cache_key, (new_token, new_expires_at), ttl=ttl_seconds) - return token + return new_token async def get_access_token_async( @@ -149,23 +146,21 @@ async def get_access_token_async( cache_key: Final = f"gigachat_token:{effective_credentials[:16]}" cached: Final = _token_cache.get_cache(cache_key) if cached: - token: Final - expires_at: Final - token, expires_at = cached - if time.time() * 1000 < expires_at - TOKEN_EXPIRY_BUFFER_MS: + _token, _expires_at = cached + if time.time() * 1000 < _expires_at - TOKEN_EXPIRY_BUFFER_MS: verbose_logger.debug("Using cached GigaChat access token") - return token + return _token # Request new token - token, expires_at = await _request_token_async(effective_credentials, effective_scope, effective_auth_url) + new_token, new_expires_at = await _request_token_async(effective_credentials, effective_scope, effective_auth_url) - if expires_at: + if new_expires_at: # Cache token - ttl_seconds: Final = max(0, (expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000) + ttl_seconds: Final = max(0, (new_expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000) if ttl_seconds > 0: - _token_cache.set_cache(cache_key, (token, expires_at), ttl=ttl_seconds) + _token_cache.set_cache(cache_key, (new_token, new_expires_at), ttl=ttl_seconds) - return token + return new_token def _request_token_sync( diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index b625e9722f8..20b1513ae15 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -289,13 +289,14 @@ class GigaChatConfig(BaseConfig): texts.append(part.get("text", "")) elif part.get("type") == "image_url": # Extract image URL and upload to GigaChat - image_url = part.get("image_url", {}) + image_url: object = part.get("image_url", {}) + upload_url: str if isinstance(image_url, str): - url: Final = image_url + upload_url = image_url else: - url: Final = image_url.get("url", "") - if url: - file_id = self._upload_image(url) # rebind-ok: inside for loop, no outer binding + upload_url = str(image_url.get("url", "")) if isinstance(image_url, dict) else "" + if upload_url: + file_id = self._upload_image(upload_url) if file_id: attachments.append(file_id) text: Final = "\n".join(texts) if texts else "" diff --git a/litellm/llms/gigachat/embedding/transformation.py b/litellm/llms/gigachat/embedding/transformation.py index 6c6ed3ce35f..6aaf63297cb 100644 --- a/litellm/llms/gigachat/embedding/transformation.py +++ b/litellm/llms/gigachat/embedding/transformation.py @@ -115,10 +115,8 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): # Normalize input to list if isinstance(input, str): input_list: list = [input] - elif isinstance(input, list): - input_list = input else: - input_list = [input] + input_list = input # Remove gigachat/ prefix from model if present model = model.removeprefix("gigachat/") diff --git a/litellm/llms/gigachat/file_handler.py b/litellm/llms/gigachat/file_handler.py index 163e944f124..900aa72f34c 100644 --- a/litellm/llms/gigachat/file_handler.py +++ b/litellm/llms/gigachat/file_handler.py @@ -14,7 +14,7 @@ from typing import Final from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( - _get_httpx_client, + HTTPHandler, get_async_httpx_client, ) from litellm.llms.gigachat.utils import get_api_base @@ -52,7 +52,7 @@ def _parse_data_url(data_url: str) -> tuple[bytes, str, str] | None: def _download_image_sync(url: str) -> tuple[bytes, str, str]: """Download image from URL synchronously.""" - client: Final = _get_httpx_client(params={"ssl_verify": False}) + client: Final = HTTPHandler(ssl_verify=False) response: Final = client.get(url) response.raise_for_status() @@ -120,7 +120,7 @@ def upload_file_sync( base_url: Final = get_api_base(api_base) upload_url: Final = f"{base_url}/files" - client: Final = _get_httpx_client(params={"ssl_verify": False}) + client: Final = HTTPHandler(ssl_verify=False) response: Final = client.post( upload_url, headers={"Authorization": f"Bearer {access_token}"}, diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index ec441c53cc4..531a53792b3 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -157,21 +157,13 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): all_translated_chunks: Final[list[object]] = [] # mutable-ok: accumulator for chunk in all_chunks: - if isinstance(chunk, bytes): - chunk = chunk.decode("utf-8", errors="ignore") - - if isinstance(chunk, str): - chunk = chunk.strip() - if not chunk or chunk == "[DONE]": - continue - chunk = chunk.removeprefix("data: ") - try: - message = json.loads(chunk) - except json.JSONDecodeError: - continue - elif isinstance(chunk, dict): - message = chunk - else: + chunk = chunk.strip() + if not chunk or chunk == "[DONE]": + continue + chunk = chunk.removeprefix("data: ") + try: + message = json.loads(chunk) + except json.JSONDecodeError: continue gigachat_iterator = GigaChatModelResponseIterator( @@ -180,7 +172,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): ) translated_chunk = gigachat_iterator.chunk_parser(chunk=message) - if isinstance(translated_chunk, dict) and generic_chunk_has_all_required_fields(translated_chunk): + if isinstance(translated_chunk, dict) and generic_chunk_has_all_required_fields(translated_chunk): # pyright: ignore[reportUnnecessaryIsInstance] # runtime guard for patched chunk_parser chunk_obj = convert_generic_chunk_to_model_response_stream( translated_chunk # pyright: ignore[reportArgumentType] # validated TypedDict ) @@ -212,5 +204,5 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): def get_base_model(model: str) -> str | None: return model - def get_models(self, api_key: str | None = None, api_base: str | None = None) -> Sequence[str]: - return super().get_models(api_key, api_base) + def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]: + return list(super().get_models(api_key, api_base)) diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 9d2ca66965b..5aef82531f8 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -8,6 +8,7 @@ import asyncio import contextvars from collections.abc import AsyncGenerator, AsyncIterator, Coroutine, Generator, Iterator from functools import partial +from types import TracebackType from typing import Any, Final, cast import httpx @@ -124,7 +125,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): async def __anext__(self) -> bytes: if not self._initialized: - await self + await self # pyright: ignore[reportGeneralTypeIssues] # structural type check misses __await__ try: chunk = await anext(self._iterator) self._raw_bytes.append(chunk) @@ -140,17 +141,17 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): async def asend(self, value: bytes) -> bytes: if not self._initialized: - await self + await self # pyright: ignore[reportGeneralTypeIssues] # structural type check misses __await__ return await self._iterator.asend(value) async def athrow( self, - typ: type[BaseException], - val: BaseException | None = None, - tb: type | None = None, + typ: BaseException | type[BaseException], + val: BaseException | object = None, + tb: TracebackType | None = None, ) -> bytes: if not self._initialized: - await self + await self # pyright: ignore[reportGeneralTypeIssues] # structural type check misses __await__ return await self._iterator.athrow(typ, val, tb) async def aclose(self) -> None: @@ -221,9 +222,9 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): def throw( self, - typ: type[BaseException], - val: BaseException | None = None, - tb: type | None = None, + typ: BaseException | type[BaseException], + val: BaseException | object = None, + tb: TracebackType | None = None, ) -> bytes: return self._iterator.throw(typ, val, tb) @@ -552,8 +553,8 @@ def llm_passthrough_route( else: return response except Exception as e: - if provider_config is None: - raise e + # provider_config is guaranteed non-None here due to the earlier guard + assert provider_config is not None raise base_llm_http_handler._handle_error( e=e, provider_config=provider_config, @@ -577,7 +578,7 @@ async def _async_passthrough_request( # Check if it's a coroutine and await it if asyncio.iscoroutine(response_result): if is_streaming_request: - return await AsyncPassthroughStreamingResponse( + return await AsyncPassthroughStreamingResponse( # pyright: ignore[reportGeneralTypeIssues] # structural type check misses __await__ response=response_result, litellm_logging_obj=litellm_logging_obj, provider_config=provider_config, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index b7a1c71a9e5..3dab313bd5c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -47,6 +47,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( create_websocket_passthrough_route, websocket_passthrough_request, ) +from litellm.proxy.utils import ProxyLogging as ProxyLoggingType from litellm.proxy.utils import is_known_model from litellm.proxy.vector_store_endpoints.utils import ( assert_proxy_admin_for_vector_store_index_management, @@ -1676,7 +1677,7 @@ def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict: def get_vertex_pass_through_handler( - call_type: Literal[discovery, aiplatform], + call_type: Literal["discovery", "aiplatform"], # noqa: UP037 ) -> BaseVertexAIPassThroughHandler: if call_type == "discovery": return VertexAIDiscoveryPassThroughHandler() @@ -2732,7 +2733,7 @@ async def gigachat_proxy_route( "Gigachat passthrough: Using direct Gigachat model '%s' for endpoint '%s'", model, endpoint ) - data: Dict[str, Any] = {} # mutable-ok: request body mutated in place by proxy pipeline + data: dict[str, Any] = {} # mutable-ok: request body mutated in place by proxy pipeline data["method"] = request.method data["endpoint"] = endpoint @@ -2784,7 +2785,7 @@ async def handle_gigachat_passthrough_router_model( fastapi_response: Response, llm_router: litellm.Router, user_api_key_dict: UserAPIKeyAuth, - proxy_logging_obj: ProxyLogging, + proxy_logging_obj: ProxyLoggingType, general_settings: dict, proxy_config: ProxyConfig, select_data_generator: Callable, @@ -2822,7 +2823,7 @@ async def handle_gigachat_passthrough_router_model( # Detect streaming based on request body is_streaming = request_body.get("stream", False) - data: Dict[str, Any] = await _read_request_body(request=request) + data: dict[str, Any] = await _read_request_body(request=request) if user_api_key_dict is not None: if data.get("metadata") is None: data["metadata"] = {} # mutable-ok: metadata dict mutated in place diff --git a/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py b/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py index d2c617cd220..0a6ef364954 100644 --- a/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py +++ b/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py @@ -299,13 +299,13 @@ class TestGigaChatPassthroughConfig: assert result.choices[0].message.content == "Hello world" def test_handle_logging_collected_chunks_with_bytes_chunks(self): - """Test converting bytes chunks to model response.""" + """Test converting string chunks to model response (bytes pre-decoded upstream).""" config = GigaChatPassthroughConfig() logging_obj = MagicMock() chunks = [ - b'{"choices": [{"delta": {"content": "Hi"}, "index": 0}]}', - b'{"choices": [{"delta": {}, "finish_reason": "stop", "index": 0}], "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}}', + '{"choices": [{"delta": {"content": "Hi"}, "index": 0}]}', + '{"choices": [{"delta": {}, "finish_reason": "stop", "index": 0}], "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}}', ] result = config.handle_logging_collected_chunks( @@ -343,26 +343,28 @@ class TestGigaChatPassthroughConfig: assert result.choices[0].message.content == "test" def test_handle_logging_collected_chunks_with_dict_chunks(self): - """Test converting dict chunks directly.""" + """Test converting string-serialized dict chunks (dicts pre-serialized upstream).""" config = GigaChatPassthroughConfig() logging_obj = MagicMock() chunks = [ - {"choices": [{"delta": {"content": "direct"}, "index": 0}]}, - { - "choices": [ - { - "delta": {}, - "finish_reason": "stop", - "index": 0, - } - ], - "usage": { - "prompt_tokens": 1, - "completion_tokens": 1, - "total_tokens": 2, - }, - }, + '{"choices": [{"delta": {"content": "direct"}, "index": 0}]}', + json.dumps( + { + "choices": [ + { + "delta": {}, + "finish_reason": "stop", + "index": 0, + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + } + ), ] result = config.handle_logging_collected_chunks( @@ -587,12 +589,12 @@ class TestGigaChatPassthroughConfig: assert result is None def test_handle_logging_collected_chunks_skips_unsupported_chunk_type(self): - """Test that unsupported chunk types (int, float, etc.) are skipped.""" + """Test that unsupported chunk types (non-JSON str) are skipped.""" config = GigaChatPassthroughConfig() logging_obj = MagicMock() - # The chunk is an int which doesn't match str/bytes/dict - chunks: list = [42, "not-a-real-chunk"] + # Both are valid str chunks; "not-a-valid-json" fails json.loads, int is not a str + chunks: list[str] = ["not-a-valid-json"] result = config.handle_logging_collected_chunks( all_chunks=chunks, diff --git a/tests/test_litellm/llms/gigachat/test_file_handler.py b/tests/test_litellm/llms/gigachat/test_file_handler.py index ae019eb7942..ced7c30ed9d 100644 --- a/tests/test_litellm/llms/gigachat/test_file_handler.py +++ b/tests/test_litellm/llms/gigachat/test_file_handler.py @@ -128,14 +128,14 @@ class TestParseDataUrl: class TestDownloadImageSync: - @patch(f"{FILE_MODULE}._get_httpx_client") - def test_downloads_image_successfully(self, mock_get_client): + @patch(f"{FILE_MODULE}.HTTPHandler") + def test_downloads_image_successfully(self, mock_http_handler_cls): mock_client = MagicMock() mock_response = MagicMock() mock_response.content = b"fake-image-bytes" mock_response.headers = {"content-type": "image/jpeg"} mock_client.get.return_value = mock_response - mock_get_client.return_value = mock_client + mock_http_handler_cls.return_value = mock_client content_bytes, content_type, ext = file_handler._download_image_sync("https://example.com/img.jpg") @@ -144,41 +144,41 @@ class TestDownloadImageSync: assert ext == "jpeg" mock_client.get.assert_called_once_with("https://example.com/img.jpg") - @patch(f"{FILE_MODULE}._get_httpx_client") - def test_raises_on_http_error(self, mock_get_client): + @patch(f"{FILE_MODULE}.HTTPHandler") + def test_raises_on_http_error(self, mock_http_handler_cls): mock_client = MagicMock() mock_client.get.side_effect = httpx.HTTPStatusError( "Not Found", request=httpx.Request("GET", "https://example.com/404"), response=httpx.Response(status_code=404, request=httpx.Request("GET", "https://example.com/404")), ) - mock_get_client.return_value = mock_client + mock_http_handler_cls.return_value = mock_client with pytest.raises(httpx.HTTPStatusError): file_handler._download_image_sync("https://example.com/404") - @patch(f"{FILE_MODULE}._get_httpx_client") - def test_parse_content_type_fallback(self, mock_get_client): + @patch(f"{FILE_MODULE}.HTTPHandler") + def test_parse_content_type_fallback(self, mock_http_handler_cls): mock_client = MagicMock() mock_response = MagicMock() mock_response.content = b"data" mock_response.headers = {} mock_client.get.return_value = mock_response - mock_get_client.return_value = mock_client + mock_http_handler_cls.return_value = mock_client _, content_type, ext = file_handler._download_image_sync("https://example.com/img") assert content_type == "image/jpeg" assert ext == "jpeg" - @patch(f"{FILE_MODULE}._get_httpx_client") - def test_extracts_extension_from_parametrized_type(self, mock_get_client): + @patch(f"{FILE_MODULE}.HTTPHandler") + def test_extracts_extension_from_parametrized_type(self, mock_http_handler_cls): mock_client = MagicMock() mock_response = MagicMock() mock_response.content = b"data" mock_response.headers = {"content-type": "image/png; charset=utf-8"} mock_client.get.return_value = mock_response - mock_get_client.return_value = mock_client + mock_http_handler_cls.return_value = mock_client _, _, ext = file_handler._download_image_sync("https://example.com/img.png") @@ -235,16 +235,16 @@ class TestDownloadImageAsync: class TestUploadFileSync: @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") - @patch(f"{FILE_MODULE}._get_httpx_client") + @patch(f"{FILE_MODULE}.HTTPHandler") def test_uploads_base64_image_and_caches( - self, mock_get_client, mock_get_token, mock_get_api_base + self, mock_http_handler_cls, mock_get_token, mock_get_api_base ): mock_client = MagicMock() mock_response = MagicMock() mock_response.json.return_value = {"id": "file-12345"} mock_response.raise_for_status = MagicMock() mock_client.post.return_value = mock_response - mock_get_client.return_value = mock_client + mock_http_handler_cls.return_value = mock_client result = upload_file_sync( image_url=_RED_PNG_DATA_URL, @@ -268,9 +268,9 @@ class TestUploadFileSync: @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") - @patch(f"{FILE_MODULE}._get_httpx_client") + @patch(f"{FILE_MODULE}.HTTPHandler") def test_returns_cached_file_id( - self, mock_get_client, mock_get_token, mock_get_api_base + self, mock_http_handler_cls, mock_get_token, mock_get_api_base ): # Pre-populate the cache url_hash = _get_url_hash(_RED_PNG_DATA_URL) @@ -280,14 +280,14 @@ class TestUploadFileSync: assert result == "cached-file-id" # No upload call was made - mock_get_client.return_value.post.assert_not_called() + mock_http_handler_cls.return_value.post.assert_not_called() - @patch(f"{FILE_MODULE}._get_httpx_client") + @patch(f"{FILE_MODULE}.HTTPHandler") @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") @patch(f"{FILE_MODULE}._download_image_sync") def test_downloads_and_uploads_url_image( - self, mock_download, mock_get_api_base, mock_get_token, mock_get_client + self, mock_download, mock_get_api_base, mock_get_token, mock_http_handler_cls ): mock_download.return_value = (b"remote-bytes", "image/png", "png") mock_client = MagicMock() @@ -295,7 +295,7 @@ class TestUploadFileSync: mock_response.json.return_value = {"id": "file-remote"} mock_response.raise_for_status = MagicMock() mock_client.post.return_value = mock_response - mock_get_client.return_value = mock_client + mock_http_handler_cls.return_value = mock_client result = upload_file_sync( image_url="https://example.com/remote.png", credentials="creds" @@ -304,11 +304,11 @@ class TestUploadFileSync: assert result == "file-remote" mock_download.assert_called_once_with("https://example.com/remote.png") - @patch(f"{FILE_MODULE}._get_httpx_client") + @patch(f"{FILE_MODULE}.HTTPHandler") @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") def test_returns_none_on_upload_failure( - self, mock_get_api_base, mock_get_token, mock_get_client + self, mock_get_api_base, mock_get_token, mock_http_handler_cls ): mock_client = MagicMock() mock_client.post.side_effect = httpx.HTTPStatusError( @@ -316,7 +316,7 @@ class TestUploadFileSync: request=httpx.Request("POST", "https://api.example.com/files"), response=httpx.Response(status_code=400, request=httpx.Request("POST", "https://api.example.com/files")), ) - mock_get_client.return_value = mock_client + mock_http_handler_cls.return_value = mock_client # upload_file_sync catches all exceptions and returns None result = upload_file_sync( @@ -325,18 +325,18 @@ class TestUploadFileSync: assert result is None - @patch(f"{FILE_MODULE}._get_httpx_client") + @patch(f"{FILE_MODULE}.HTTPHandler") @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") def test_returns_none_when_response_missing_id( - self, mock_get_api_base, mock_get_token, mock_get_client + self, mock_get_api_base, mock_get_token, mock_http_handler_cls ): mock_client = MagicMock() mock_response = MagicMock() mock_response.json.return_value = {"status": "ok"} # no "id" key mock_response.raise_for_status = MagicMock() mock_client.post.return_value = mock_response - mock_get_client.return_value = mock_client + mock_http_handler_cls.return_value = mock_client result = upload_file_sync( image_url=_RED_PNG_DATA_URL, credentials="creds" @@ -346,9 +346,9 @@ class TestUploadFileSync: @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") - @patch(f"{FILE_MODULE}._get_httpx_client") + @patch(f"{FILE_MODULE}.HTTPHandler") def test_uploads_without_optional_args( - self, mock_get_client, mock_get_token, mock_get_api_base + self, mock_http_handler_cls, mock_get_token, mock_get_api_base ): """Verify that credentials, api_base, and litellm_params are optional.""" mock_client = MagicMock() @@ -356,7 +356,7 @@ class TestUploadFileSync: mock_response.json.return_value = {"id": "file-no-args"} mock_response.raise_for_status = MagicMock() mock_client.post.return_value = mock_response - mock_get_client.return_value = mock_client + mock_http_handler_cls.return_value = mock_client result = upload_file_sync(image_url=_RED_PNG_DATA_URL) From 5af54d26b81a607c2a1e528d4cb6e24f052a029e Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Wed, 26 Aug 2026 19:28:42 +0000 Subject: [PATCH 055/120] fix(lint): Partly fix LIT issues --- litellm/llms/gigachat/authenticator.py | 6 ++-- litellm/llms/gigachat/chat/streaming.py | 8 ++--- .../llms/gigachat/embedding/transformation.py | 4 +-- .../gigachat/passthrough/transformation.py | 2 +- litellm/llms/gigachat/utils.py | 8 +++-- litellm/passthrough/main.py | 6 ++-- litellm/proxy/common_request_processing.py | 8 ++--- .../llm_passthrough_endpoints.py | 36 ++++++++++--------- 8 files changed, 42 insertions(+), 36 deletions(-) diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index 9b8ef3ec93f..a2e8030a2d1 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -249,12 +249,12 @@ def _parse_token_response(response: httpx.Response) -> tuple[str, int]: message=f"Invalid token response: {data}", ) - expires_at: int # expires_at is in milliseconds + expires_at: int # rebind-ok: conditionally assigned from str or int if isinstance(expires_at_raw, str): - expires_at = int(expires_at_raw) + expires_at = int(expires_at_raw) # rebind-ok: conditionally assigned from str or int else: - expires_at = expires_at_raw # pyright: ignore[reportAssignmentType] # raw value is int or str; converted above + expires_at = expires_at_raw # pyright: ignore[reportAssignmentType] # raw value is int or str; converted above; rebind-ok: conditionally assigned from str or int verbose_logger.debug("GigaChat access token obtained successfully") return access_token, expires_at diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index 338d37d810a..64f3057a66f 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -48,15 +48,15 @@ class GigaChatModelResponseIterator: # Extract text content text: Final = delta.get("content", "") or "" - usage_block: ChatCompletionUsageBlock | None = None - tool_use: ChatCompletionToolCallChunk | None = None + usage_block: ChatCompletionUsageBlock | None = None # rebind-ok: conditionally assigned after stop detection + tool_use: ChatCompletionToolCallChunk | None = None # rebind-ok: conditionally assigned on function_call finish_reason: str | None = chunk_finish_reason # Handle function_call in stream if chunk_finish_reason == "function_call" and delta.get("function_call"): func_call: Final = delta["function_call"] args_raw: Final = func_call.get("arguments") or {} - args_str: str + args_str: str # rebind-ok: conditionally assigned from dict or str if isinstance(args_raw, dict): args_str = json.dumps(args_raw, ensure_ascii=False) # rebind-ok: build from dict else: @@ -76,7 +76,7 @@ class GigaChatModelResponseIterator: if chunk_finish_reason == "stop": usage_data: Final = chunk.get("usage") or {} # mutable-ok: empty dict default if usage_data: - usage: Final = convert_usage(usage_data) + usage = convert_usage(usage_data) # rebind-ok: conditional usage assignment usage_block = ChatCompletionUsageBlock( prompt_tokens=usage.prompt_tokens, completion_tokens=usage.completion_tokens, diff --git a/litellm/llms/gigachat/embedding/transformation.py b/litellm/llms/gigachat/embedding/transformation.py index 6aaf63297cb..2ec8324e33c 100644 --- a/litellm/llms/gigachat/embedding/transformation.py +++ b/litellm/llms/gigachat/embedding/transformation.py @@ -114,12 +114,12 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): """ # Normalize input to list if isinstance(input, str): - input_list: list = [input] + input_list: list = [input] # rebind-ok: locally scoped conversion else: input_list = input # Remove gigachat/ prefix from model if present - model = model.removeprefix("gigachat/") + model = model.removeprefix("gigachat/") # rebind-ok: parameter reassignment for normalization return { "model": model, diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index 531a53792b3..fe65cf5561f 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -92,7 +92,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): if provider_chat_config is None: raise ValueError(f"No provider config found for model: {model}") - litellm_model_response: ModelResponse = provider_chat_config.transform_response( + litellm_model_response: Final = provider_chat_config.transform_response( model=model, messages=request_data.get("messages", []), # mutable-ok: empty list default for transform_response raw_response=httpx_response, diff --git a/litellm/llms/gigachat/utils.py b/litellm/llms/gigachat/utils.py index aa30552a7ec..b1db8f685e5 100644 --- a/litellm/llms/gigachat/utils.py +++ b/litellm/llms/gigachat/utils.py @@ -17,9 +17,13 @@ def convert_usage(usage_data: Mapping[str, int]) -> Usage: prompt_tokens_total: Final = prompt_tokens + precached_prompt_tokens total_tokens_total: Final = total_tokens + precached_prompt_tokens - prompt_tokens_details: PromptTokensDetailsWrapper | None = None + prompt_tokens_details: PromptTokensDetailsWrapper | None = ( + None # rebind-ok: conditionally assigned when cached tokens exist + ) if precached_prompt_tokens > 0: - prompt_tokens_details = PromptTokensDetailsWrapper(cached_tokens=precached_prompt_tokens) + prompt_tokens_details = PromptTokensDetailsWrapper( + cached_tokens=precached_prompt_tokens + ) # rebind-ok: conditionally assigned when cached tokens exist return Usage( prompt_tokens=prompt_tokens_total, diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 956badfcfe1..a64aa894221 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -101,7 +101,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): self._flush_scheduled = True try: - task = asyncio.create_task( + task: Final = asyncio.create_task( self._litellm_logging_obj.async_flush_passthrough_collected_chunks( raw_bytes=self._raw_bytes, provider_config=self._provider_config, @@ -127,7 +127,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): if not self._initialized: await self # pyright: ignore[reportGeneralTypeIssues] # structural type check misses __await__ try: - chunk = await anext(self._iterator) + chunk: Final = await anext(self._iterator) self._raw_bytes.append(chunk) except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic self._start_flush() @@ -205,7 +205,7 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): def __next__(self) -> bytes: try: - chunk = next(self._iterator) + chunk: Final = next(self._iterator) self._raw_bytes.append(chunk) except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic self._start_flush() diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 90eae550fe0..69a2ff53672 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1448,7 +1448,7 @@ class ProxyBaseLLMRequestProcessing: Proxy/custom headers win on key collisions. """ - excluded_headers = { # mutable-ok: set of header names to exclude from forwarding + excluded_headers: Final = { # mutable-ok: set of header names to exclude from forwarding "transfer-encoding", "content-encoding", "set-cookie", @@ -1461,7 +1461,7 @@ class ProxyBaseLLMRequestProcessing: "upgrade", } - merged_headers = { # mutable-ok: dict comprehension for merged headers forwarded to httpx + merged_headers: Final = { # mutable-ok: dict comprehension for merged headers forwarded to httpx key: value for key, value in dict(response_headers or {}).items() if key.lower() not in excluded_headers } merged_headers.update(custom_headers) @@ -2425,9 +2425,9 @@ class ProxyBaseLLMRequestProcessing: logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete if route_type == "allm_passthrough_route": - streaming_headers = custom_headers + streaming_headers = custom_headers # rebind-ok: initial assignment before header merge if hasattr(response, "headers"): - streaming_headers = ProxyBaseLLMRequestProcessing._merge_passthrough_streaming_headers( + streaming_headers = ProxyBaseLLMRequestProcessing._merge_passthrough_streaming_headers( # rebind-ok: merge result replaces initial assignment response_headers=getattr(response, "headers", None), custom_headers=custom_headers, ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index a3fe2911038..3b554f5b56a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -76,9 +76,9 @@ if TYPE_CHECKING: from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig from litellm.router import Router - ProxyConfig = _ProxyConfig + ProxyConfig = _ProxyConfig # rebind-ok: conditional type alias else: - ProxyConfig = Any + ProxyConfig = Any # rebind-ok: runtime fallback vertex_llm_base: Final = VertexBase() router: Final = APIRouter() @@ -508,8 +508,8 @@ async def milvus_proxy_route( status_code=400, detail=f"collectionName must be a string. Got {type(_raw_collection_name).__name__}", ) - collection_name: str | None = _raw_collection_name - extra_headers = {} + collection_name: str | None = _raw_collection_name # rebind-ok: locally scoped conversion + extra_headers: Final = {} # mutable-ok: dict for extra headers base_target_url: str | None = None if not collection_name: raise HTTPException( @@ -2839,12 +2839,14 @@ async def gigachat_proxy_route( ) ## check for streaming - request_body = await get_request_body(request) - is_router_model = False + request_body: Final = await get_request_body(request) + is_router_model = False # rebind-ok: conditionally set to True when model uses router - model = request_body.get("model") + model: Final = request_body.get("model") if model: - is_router_model = is_passthrough_request_using_router_model(request_body, llm_router) + is_router_model = is_passthrough_request_using_router_model( + request_body, llm_router + ) # rebind-ok: conditionally set to True elif any(word in endpoint for word in ("completions", "embeddings")): raise HTTPException( status_code=400, detail={"error": "Model is required in request body"} @@ -2878,14 +2880,14 @@ async def gigachat_proxy_route( "Gigachat passthrough: Using direct Gigachat model '%s' for endpoint '%s'", model, endpoint ) - data: dict[str, Any] = {} # mutable-ok: request body mutated in place by proxy pipeline + data: Final[dict[str, Any]] = {} # mutable-ok: request body mutated in place by proxy pipeline data["method"] = request.method data["endpoint"] = endpoint data["json"] = request_body data["custom_llm_provider"] = "gigachat" - client = get_async_httpx_client( + client: Final = get_async_httpx_client( llm_provider=LlmProviders.GIGACHAT, params={ # mutable-ok: httpx client params "timeout": httpx.Timeout(timeout=600.0, connect=5.0), @@ -2894,7 +2896,7 @@ async def gigachat_proxy_route( ) data["client"] = client - base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) + base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: return await base_llm_response_processor.base_passthrough_process_llm_request( @@ -2966,9 +2968,9 @@ async def handle_gigachat_passthrough_router_model( from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing # Detect streaming based on request body - is_streaming = request_body.get("stream", False) + is_streaming: Final = request_body.get("stream", False) - data: dict[str, Any] = await _read_request_body(request=request) + data: Final[dict[str, Any]] = await _read_request_body(request=request) if user_api_key_dict is not None: if data.get("metadata") is None: data["metadata"] = {} # mutable-ok: metadata dict mutated in place @@ -2995,7 +2997,7 @@ async def handle_gigachat_passthrough_router_model( data["custom_llm_provider"] = "gigachat" # Remove sensitive keys from data - keys = [ # mutable-ok: list of keys to remove from data + keys: Final = [ # mutable-ok: list of keys to remove from data "gigachat_auth_url", "gigachat_access_token", "gigachat_scope", @@ -3005,7 +3007,7 @@ async def handle_gigachat_passthrough_router_model( for key in keys: data.pop(key, None) - client = get_async_httpx_client( + client: Final = get_async_httpx_client( llm_provider=LlmProviders.GIGACHAT, params={ # mutable-ok: httpx client params "timeout": httpx.Timeout(timeout=600.0, connect=5.0), @@ -3014,12 +3016,12 @@ async def handle_gigachat_passthrough_router_model( ) data["client"] = client - base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) + base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) # Use the common passthrough processing to handle metadata and hooks # This also handles all response formatting (streaming/non-streaming) and exceptions try: - result = await base_llm_response_processor.base_passthrough_process_llm_request( + result: Final = await base_llm_response_processor.base_passthrough_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, From 278fdbecb76f64a683089422cdf113ac9f45652c Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Wed, 26 Aug 2026 20:21:20 +0000 Subject: [PATCH 056/120] fix(lint): Partly fix basedpyright lint issues --- litellm/llms/gigachat/authenticator.py | 8 ++++---- litellm/llms/gigachat/chat/streaming.py | 6 +++--- .../llm_passthrough_endpoints.py | 18 +++++++++++------- 3 files changed, 18 insertions(+), 14 deletions(-) diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index a2e8030a2d1..293176d8931 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -107,7 +107,7 @@ def get_access_token( return _token # Request new token - new_token, new_expires_at = _request_token_sync(effective_credentials, effective_scope, effective_auth_url) + new_token, new_expires_at = _request_token_sync(effective_credentials, effective_scope, effective_auth_url) # pyright: ignore[reportArgumentType] # credential keys may be broader than str if new_expires_at: # Cache token @@ -152,7 +152,7 @@ async def get_access_token_async( return _token # Request new token - new_token, new_expires_at = await _request_token_async(effective_credentials, effective_scope, effective_auth_url) + new_token, new_expires_at = await _request_token_async(effective_credentials, effective_scope, effective_auth_url) # pyright: ignore[reportArgumentType] # credential keys may be broader than str if new_expires_at: # Cache token @@ -187,7 +187,7 @@ def _request_token_sync( client: Final = _get_http_client() response: Final = client.post(auth_url, headers=headers, data=data, timeout=30) response.raise_for_status() - return _parse_token_response(response) + return _parse_token_response(response) # pyright: ignore[reportArgumentType] # httpx Response may be None at type level except httpx.HTTPStatusError as e: raise GigaChatAuthError( status_code=e.response.status_code, @@ -222,7 +222,7 @@ async def _request_token_async( ) response: Final = await client.post(auth_url, headers=headers, data=data, timeout=30) response.raise_for_status() - return _parse_token_response(response) + return _parse_token_response(response) # pyright: ignore[reportArgumentType] # httpx Response may be None at type level except httpx.HTTPStatusError as e: raise GigaChatAuthError( status_code=e.response.status_code, diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index 64f3057a66f..3ed31c62d7d 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -5,7 +5,7 @@ GigaChat Streaming Response Handler import json import uuid from collections.abc import Mapping, Sequence -from typing import Any, Final +from typing import Any, Final, cast from litellm.llms.gigachat.utils import convert_usage from litellm.types.llms.openai import ( @@ -76,7 +76,7 @@ class GigaChatModelResponseIterator: if chunk_finish_reason == "stop": usage_data: Final = chunk.get("usage") or {} # mutable-ok: empty dict default if usage_data: - usage = convert_usage(usage_data) # rebind-ok: conditional usage assignment + usage = convert_usage(cast("Mapping[str, int]", usage_data)) # cast-ok: dict from chunk has int values usage_block = ChatCompletionUsageBlock( prompt_tokens=usage.prompt_tokens, completion_tokens=usage.completion_tokens, @@ -90,7 +90,7 @@ class GigaChatModelResponseIterator: ) return GenericStreamingChunk( - text=text, + text=cast(str, text), tool_use=tool_use, is_finished=chunk_finish_reason is not None, finish_reason=finish_reason or "", diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 3b554f5b56a..c7c290dcd2d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -509,7 +509,7 @@ async def milvus_proxy_route( detail=f"collectionName must be a string. Got {type(_raw_collection_name).__name__}", ) collection_name: str | None = _raw_collection_name # rebind-ok: locally scoped conversion - extra_headers: Final = {} # mutable-ok: dict for extra headers + extra_headers = {} # mutable-ok: dict for extra headers; rebind-ok: reassigned later from credentials base_target_url: str | None = None if not collection_name: raise HTTPException( @@ -2839,10 +2839,10 @@ async def gigachat_proxy_route( ) ## check for streaming - request_body: Final = await get_request_body(request) + request_body: Final = await get_request_body(request) # pyright: ignore[reportUnknownVariableType] # get_request_body returns Unknown is_router_model = False # rebind-ok: conditionally set to True when model uses router - model: Final = request_body.get("model") + model: Final = request_body.get("model") # pyright: ignore[reportUnknownVariableType] # get_request_body returns Unknown if model: is_router_model = is_passthrough_request_using_router_model( request_body, llm_router @@ -2880,7 +2880,9 @@ async def gigachat_proxy_route( "Gigachat passthrough: Using direct Gigachat model '%s' for endpoint '%s'", model, endpoint ) - data: Final[dict[str, Any]] = {} # mutable-ok: request body mutated in place by proxy pipeline + data: dict[ + str, Any + ] = {} # mutable-ok: request body mutated in place by proxy pipeline; pyright: ignore[reportExplicitAny] # Any needed for proxy pipeline flexibility data["method"] = request.method data["endpoint"] = endpoint @@ -2968,9 +2970,11 @@ async def handle_gigachat_passthrough_router_model( from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing # Detect streaming based on request body - is_streaming: Final = request_body.get("stream", False) + is_streaming: Final = request_body.get("stream", False) # pyright: ignore[reportUnknownVariableType] # request_body is dict[Unknown, Unknown] - data: Final[dict[str, Any]] = await _read_request_body(request=request) + data: dict[str, Any] = await _read_request_body( + 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 @@ -3021,7 +3025,7 @@ async def handle_gigachat_passthrough_router_model( # Use the common passthrough processing to handle metadata and hooks # This also handles all response formatting (streaming/non-streaming) and exceptions try: - result: Final = await base_llm_response_processor.base_passthrough_process_llm_request( + result = await base_llm_response_processor.base_passthrough_process_llm_request( # rebind-ok: assigned once in try block request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, From 6c36030e60c3181f2605b9ca69cce45817f83077 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Wed, 26 Aug 2026 20:33:27 +0000 Subject: [PATCH 057/120] fix(lint): Partly fix TID251 lint issues --- litellm/llms/gigachat/chat/streaming.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index 3ed31c62d7d..e7ca0bec83d 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -5,7 +5,7 @@ GigaChat Streaming Response Handler import json import uuid from collections.abc import Mapping, Sequence -from typing import Any, Final, cast +from typing import Any, Final from litellm.llms.gigachat.utils import convert_usage from litellm.types.llms.openai import ( @@ -76,7 +76,10 @@ class GigaChatModelResponseIterator: if chunk_finish_reason == "stop": usage_data: Final = chunk.get("usage") or {} # mutable-ok: empty dict default if usage_data: - usage = convert_usage(cast("Mapping[str, int]", usage_data)) # cast-ok: dict from chunk has int values + validated_usage: Final = { + k: int(v) for k, v in dict(usage_data).items() + } # rebind-ok: dict comprehension + usage = convert_usage(validated_usage) usage_block = ChatCompletionUsageBlock( prompt_tokens=usage.prompt_tokens, completion_tokens=usage.completion_tokens, @@ -90,7 +93,7 @@ class GigaChatModelResponseIterator: ) return GenericStreamingChunk( - text=cast(str, text), + text=str(text), tool_use=tool_use, is_finished=chunk_finish_reason is not None, finish_reason=finish_reason or "", From 014d7b87b5cabd889a41e942efa7b213ebeeb4c8 Mon Sep 17 00:00:00 2001 From: KnyazSh Date: Wed, 26 Aug 2026 20:57:38 +0000 Subject: [PATCH 058/120] fix(lint): Partly fix basedpyright lint issues --- litellm/llms/gigachat/chat/streaming.py | 22 +++++++++++----------- litellm/passthrough/main.py | 14 ++++++++++---- 2 files changed, 21 insertions(+), 15 deletions(-) diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index e7ca0bec83d..7bdca61fd0f 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -75,21 +75,21 @@ class GigaChatModelResponseIterator: if chunk_finish_reason == "stop": usage_data: Final = chunk.get("usage") or {} # mutable-ok: empty dict default - if usage_data: - validated_usage: Final = { - k: int(v) for k, v in dict(usage_data).items() - } # rebind-ok: dict comprehension + 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) - usage_block = ChatCompletionUsageBlock( + _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=( - usage.prompt_tokens_details.model_dump() if usage.prompt_tokens_details else None - ), - completion_tokens_details=( - usage.completion_tokens_details.model_dump() if usage.completion_tokens_details else None - ), + prompt_tokens_details=_prompt_details, + completion_tokens_details=_completion_details, ) return GenericStreamingChunk( diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index a64aa894221..1715ec10f30 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -152,7 +152,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): ) -> bytes: if not self._initialized: await self # pyright: ignore[reportGeneralTypeIssues] # structural type check misses __await__ - return await self._iterator.athrow(typ, val, tb) + return await self._iterator.athrow(typ, val, tb) # pyright: ignore[reportCallIssue, reportArgumentType] # matches one of the athrow overloads async def aclose(self) -> None: self._start_flush() @@ -226,7 +226,7 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): val: BaseException | object = None, tb: TracebackType | None = None, ) -> bytes: - return self._iterator.throw(typ, val, tb) + return self._iterator.throw(typ, val, tb) # pyright: ignore[reportCallIssue, reportArgumentType] # matches one of the throw overloads def close(self) -> None: self._start_flush() @@ -488,10 +488,13 @@ def llm_passthrough_route( forward_headers=False, ) + _request_data: dict | None = ( + data if isinstance(data, dict) else (json if isinstance(json, dict) else None) + ) # rebind-ok: conditional headers, signed_json_body = provider_config.sign_request( headers=headers, litellm_params=litellm_params_dict, - request_data=data if data else json, + request_data=_request_data, api_base=str(updated_url), model=model, ) @@ -513,9 +516,12 @@ def llm_passthrough_route( ) ## IS STREAMING REQUEST + _streaming_request_data: dict = ( + data if isinstance(data, dict) else (json if isinstance(json, dict) else {}) + ) # rebind-ok: conditional is_streaming_request: Final = provider_config.is_streaming_request( endpoint=endpoint, - request_data=data or json or {}, + request_data=_streaming_request_data, ) # Update logging object with streaming status From f0c336cccb6e136de44c045709c60e60a5f62dcc Mon Sep 17 00:00:00 2001 From: yatishgoel Date: Thu, 27 Aug 2026 12:42:41 +0530 Subject: [PATCH 059/120] fix(router): apply model renames to the in-memory deployment list --- litellm/router.py | 3 +- tests/test_litellm/test_router.py | 65 +++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index f0ebb539bb7..e19038da914 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9075,7 +9075,8 @@ class Router: if _deployment_on_router is not None: # deployment with this model_id exists on the router if ( - deployment.litellm_params == _deployment_on_router.litellm_params + deployment.model_name == _deployment_on_router.model_name + and deployment.litellm_params == _deployment_on_router.litellm_params and deployment.model_info == _deployment_on_router.model_info ): # No need to update diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 8716e6d6b25..c4ccfb55b4f 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8096,6 +8096,71 @@ class TestUpsertDeploymentRollback: assert len(router.model_list) == 1 +class TestUpsertDeploymentRename: + """ + Issue #38360: renaming a model wrote the new `model_name` to the db, but the reload's + `upsert_deployment` compared only `litellm_params` and `model_info`. A rename with no + other edit therefore compared equal and the router kept the old name until a restart, + so `/model/info` and `/v1/models` served the stale name and the new one was unroutable. + """ + + @staticmethod + def _router() -> "litellm.Router": + return litellm.Router( + model_list=[ + { + "model_name": "old-name", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test"}, + "model_info": {"id": "rename-1", "db_model": True}, + } + ] + ) + + @staticmethod + def _deployment(model_name: str, tpm: int | None = None): + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + return Deployment( + model_name=model_name, + litellm_params=LiteLLM_Params(model="openai/gpt-4o", api_key="sk-test", tpm=tpm), + model_info=ModelInfo(id="rename-1", db_model=True), + ) + + def test_rename_only_updates_the_router(self): + router = self._router() + + assert router.upsert_deployment(deployment=self._deployment("new-name")) is not None + + assert [model["model_name"] for model in router.model_list] == ["new-name"] + renamed = router.get_deployment(model_id="rename-1") + assert renamed is not None + assert renamed.model_name == "new-name" + + def test_rename_only_makes_the_new_name_routable(self): + router = self._router() + + router.upsert_deployment(deployment=self._deployment("new-name")) + + assert router.get_model_ids(model_name="new-name") == ["rename-1"] + assert router.get_model_ids(model_name="old-name") == [] + + def test_rename_alongside_another_edit_still_updates(self): + router = self._router() + + router.upsert_deployment(deployment=self._deployment("new-name", tpm=1234)) + + assert router.get_model_ids(model_name="new-name") == ["rename-1"] + renamed = router.get_deployment(model_id="rename-1") + assert renamed is not None + assert renamed.litellm_params.tpm == 1234 + + def test_unchanged_deployment_is_still_a_no_op(self): + router = self._router() + + assert router.upsert_deployment(deployment=self._deployment("old-name")) is None + assert [model["model_name"] for model in router.model_list] == ["old-name"] + + class TestConsumedRequestTagsStamp: """Issue #36621: when a request's tags select a tagged pre-routing strategy, those tags are consumed by the selection; the hook must stamp the rewritten model group so From b7a7754b0529721dc000cd1cadf68272da7e287d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:07:40 +0000 Subject: [PATCH 060/120] fix(bedrock): surface Nova Sonic user transcripts, speech events, and usage in realtime API Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/bedrock/realtime/transformation.py | 154 ++++++++++++++-- ...odel_prices_and_context_window_backup.json | 20 +++ litellm/types/llms/openai.py | 39 ++++ model_prices_and_context_window.json | 20 +++ .../test_bedrock_realtime_transformation.py | 168 ++++++++++++++++++ 5 files changed, 383 insertions(+), 18 deletions(-) diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 951bf636b2f..20996512295 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -7,7 +7,7 @@ Transforms between OpenAI Realtime API format and Bedrock Nova Sonic format. import base64 import json import uuid as uuid_lib -from typing import Any, Final +from typing import Any, Final, cast from pydantic import BaseModel @@ -20,16 +20,21 @@ from litellm.types.llms.openai import ( OpenAIRealtimeContentPartDone, OpenAIRealtimeDoneEvent, OpenAIRealtimeEvents, + OpenAIRealtimeInputAudioBufferSpeechEvent, + OpenAIRealtimeInputAudioTranscriptionCompleted, + OpenAIRealtimeInputAudioTranscriptionDelta, OpenAIRealtimeOutputItemDone, OpenAIRealtimeResponseAudioDone, OpenAIRealtimeResponseContentPartAdded, OpenAIRealtimeResponseDelta, OpenAIRealtimeResponseDoneObject, OpenAIRealtimeResponseTextDone, + OpenAIRealtimeResponseUsage, OpenAIRealtimeStreamResponseBaseObject, OpenAIRealtimeStreamResponseOutputItemAdded, OpenAIRealtimeStreamSession, OpenAIRealtimeStreamSessionEvents, + OpenAIRealtimeUsageTokenDetails, ) from litellm.types.realtime import ( ALL_DELTA_TYPES, @@ -43,6 +48,27 @@ class BedrockContentEnd(BaseModel): stopReason: str | None = None +class BedrockUsageTokenDetails(BaseModel): + speechTokens: int = 0 + textTokens: int = 0 + + +class BedrockUsageDetailsTotal(BaseModel): + input: BedrockUsageTokenDetails = BedrockUsageTokenDetails() + output: BedrockUsageTokenDetails = BedrockUsageTokenDetails() + + +class BedrockUsageDetails(BaseModel): + total: BedrockUsageDetailsTotal = BedrockUsageDetailsTotal() + + +class BedrockUsageEvent(BaseModel): + totalInputTokens: int = 0 + totalOutputTokens: int = 0 + totalTokens: int = 0 + details: BedrockUsageDetails = BedrockUsageDetails() + + TRIGGER_AUDIO_SAMPLE_RATE_HERTZ: Final = 16000 TRIGGER_AUDIO_BYTES_PER_SECOND: Final = TRIGGER_AUDIO_SAMPLE_RATE_HERTZ * 2 TRIGGER_LEADING_SILENCE: Final = bytes(TRIGGER_AUDIO_BYTES_PER_SECOND // 2) @@ -87,6 +113,12 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Text configuration self.text_media_type = "text/plain" + # Response-stream state (Bedrock events carry no role on textOutput, + # so the USER/ASSISTANT split from contentStart is tracked here) + self._user_transcript_active = False + self._user_transcript_generation_stage: str | None = None + self._latest_usage: OpenAIRealtimeResponseUsage | None = None + def validate_environment(self, headers: dict, model: str, api_key: str | None = None) -> dict: """Validate environment - no special validation needed for Bedrock.""" return headers @@ -691,6 +723,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): role: Final = content_start.get("role") if role != "ASSISTANT": + if role == "USER" and content_start.get("type") == "TEXT": + self._user_transcript_active = True + self._user_transcript_generation_stage = self._parse_generation_stage( + content_start.get("additionalModelFields") + ) return ( [], current_response_id, @@ -700,6 +737,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): ) verbose_logger.debug("Handling ASSISTANT contentStart") + is_new_response: Final = current_response_id is None # Initialize IDs if needed if not current_response_id: @@ -715,7 +753,8 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): returned_messages: Final[list[OpenAIRealtimeEvents]] = [] - # Send response.created + # Send response.created only once per response (a response can contain + # multiple content blocks, e.g. TEXT then AUDIO) response_created: Final = OpenAIRealtimeStreamResponseBaseObject( type="response.created", event_id=f"event_{uuid.uuid4()}", @@ -727,7 +766,8 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "conversation_id": current_conversation_id, }, ) - returned_messages.append(response_created) + if is_new_response: + returned_messages.append(response_created) # Send response.output_item.added output_item_added: Final = OpenAIRealtimeStreamResponseOutputItemAdded( @@ -767,6 +807,70 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_delta_type, ) + @staticmethod + def _parse_generation_stage(additional_model_fields: object) -> str | None: + if not isinstance(additional_model_fields, str): + return None + try: + parsed: Final = json.loads(additional_model_fields) + except json.JSONDecodeError: + return None + stage: Final = parsed.get("generationStage") if isinstance(parsed, dict) else None + return stage if isinstance(stage, str) else None + + def transform_user_speech_event(self, is_speech_start: bool) -> tuple[OpenAIRealtimeEvents, ...]: + """Transform Bedrock userSpeechStart/userSpeechEnd to OpenAI speech boundary events.""" + verbose_logger.debug("Handling userSpeech%s", "Start" if is_speech_start else "End") + speech_event: Final[OpenAIRealtimeInputAudioBufferSpeechEvent] = { + "type": "input_audio_buffer.speech_started" if is_speech_start else "input_audio_buffer.speech_stopped", + "event_id": f"event_{uuid.uuid4()}", + "item_id": f"item_{uuid.uuid4()}", + } + return (speech_event,) + + def transform_usage_event(self, usage_event: BedrockUsageEvent) -> None: + """Record Bedrock usageEvent token totals for the next response.done.""" + verbose_logger.debug("Handling usageEvent") + input_details: Final[OpenAIRealtimeUsageTokenDetails] = { + "audio_tokens": usage_event.details.total.input.speechTokens, + "text_tokens": usage_event.details.total.input.textTokens, + "cached_tokens": 0, + } + output_details: Final[OpenAIRealtimeUsageTokenDetails] = { + "audio_tokens": usage_event.details.total.output.speechTokens, + "text_tokens": usage_event.details.total.output.textTokens, + } + latest_usage: Final[OpenAIRealtimeResponseUsage] = { + "input_tokens": usage_event.totalInputTokens, + "output_tokens": usage_event.totalOutputTokens, + "total_tokens": usage_event.totalTokens, + "input_token_details": input_details, + "output_token_details": output_details, + } + self._latest_usage = latest_usage + + def transform_user_transcript_event(self, transcript: str) -> tuple[OpenAIRealtimeEvents, ...]: + """Transform a USER-role Bedrock textOutput (ASR transcript) to OpenAI transcription events.""" + verbose_logger.debug("Handling USER textOutput (ASR transcript)") + item_id: Final = f"item_{uuid.uuid4()}" + delta_event: Final[OpenAIRealtimeInputAudioTranscriptionDelta] = { + "type": "conversation.item.input_audio_transcription.delta", + "event_id": f"event_{uuid.uuid4()}", + "item_id": item_id, + "content_index": 0, + "delta": transcript, + } + if self._user_transcript_generation_stage == "SPECULATIVE": + return (delta_event,) + completed_event: Final[OpenAIRealtimeInputAudioTranscriptionCompleted] = { + "type": "conversation.item.input_audio_transcription.completed", + "event_id": f"event_{uuid.uuid4()}", + "item_id": item_id, + "content_index": 0, + "transcript": transcript, + } + return (delta_event, completed_event) + def transform_text_output_event( self, event: dict, @@ -985,7 +1089,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): if not current_response_id or not current_conversation_id: return [], None, None, None - usage_obj: Final = get_empty_usage() + empty_usage: Final = get_empty_usage() + zero_usage: Final[OpenAIRealtimeResponseUsage] = { + "input_tokens": empty_usage.prompt_tokens, + "output_tokens": empty_usage.completion_tokens, + "total_tokens": empty_usage.total_tokens, + } + usage: Final = self._latest_usage or zero_usage + self._latest_usage = None response_done: Final = OpenAIRealtimeDoneEvent( type="response.done", event_id=f"event_{uuid.uuid4()}", @@ -995,11 +1106,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): status="completed", output=[], conversation_id=current_conversation_id, - usage={ - "prompt_tokens": usage_obj.prompt_tokens, - "completion_tokens": usage_obj.completion_tokens, - "total_tokens": usage_obj.total_tokens, - }, + usage=dict(usage), ), ) @@ -1042,8 +1149,6 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Create a function call arguments done event # This is a custom event format that matches what clients expect - from typing import cast - function_call_event: Final[dict[str, Any]] = { "type": "response.function_call_arguments.done", "event_id": f"event_{uuid.uuid4()}", @@ -1194,18 +1299,25 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): returned_messages.extend(events) elif "textOutput" in event: - events, current_delta_chunks = self.transform_text_output_event( - event, - current_output_item_id, - current_response_id, - current_delta_chunks, - ) - returned_messages.extend(events) + if self._user_transcript_active: + returned_messages.extend(self.transform_user_transcript_event(event["textOutput"].get("content", ""))) + else: + events, current_delta_chunks = self.transform_text_output_event( + event, + current_output_item_id, + current_response_id, + current_delta_chunks, + ) + returned_messages.extend(events) elif "audioOutput" in event: events = self.transform_audio_output_event(event, current_output_item_id, current_response_id) returned_messages.extend(events) + elif "contentEnd" in event and self._user_transcript_active: + self._user_transcript_active = False + self._user_transcript_generation_stage = None + elif "contentEnd" in event: events, current_delta_chunks = self.transform_content_end_event( event, @@ -1224,6 +1336,12 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): ) = self._response_done_events(current_response_id, current_conversation_id) returned_messages.extend(done_events) + elif "userSpeechStart" in event or "userSpeechEnd" in event: + returned_messages.extend(self.transform_user_speech_event("userSpeechStart" in event)) + + elif "usageEvent" in event: + self.transform_usage_event(BedrockUsageEvent.model_validate(event["usageEvent"])) + elif "toolUse" in event: events, tool_call_id, tool_name = self.transform_tool_use_event( event, current_output_item_id, current_response_id diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 77572f69b8b..8fb6108bdda 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -553,6 +553,26 @@ "supports_response_schema": true, "supports_vision": true }, + "amazon.nova-sonic-v1:0": { + "input_cost_per_audio_token": 3.4e-06, + "input_cost_per_token": 6e-08, + "litellm_provider": "bedrock", + "mode": "realtime", + "output_cost_per_audio_token": 1.36e-05, + "output_cost_per_token": 2.4e-07, + "supports_audio_input": true, + "supports_audio_output": true + }, + "amazon.nova-2-sonic-v1:0": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 3.3e-07, + "litellm_provider": "bedrock", + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2.75e-06, + "supports_audio_input": true, + "supports_audio_output": true + }, "amazon.rerank-v1:0": { "input_cost_per_query": 0.001, "input_cost_per_token": 0.0, diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index a6115640d78..fcade835cce 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -2162,6 +2162,42 @@ class OpenAIRealtimeDoneEvent(TypedDict): type: Literal["response.done"] +class OpenAIRealtimeInputAudioBufferSpeechEvent(TypedDict): + type: ReadOnly[Literal["input_audio_buffer.speech_started", "input_audio_buffer.speech_stopped"]] + event_id: ReadOnly[str] + item_id: ReadOnly[str] + + +class OpenAIRealtimeInputAudioTranscriptionDelta(TypedDict): + type: ReadOnly[Literal["conversation.item.input_audio_transcription.delta"]] + event_id: ReadOnly[str] + item_id: ReadOnly[str] + content_index: ReadOnly[int] + delta: ReadOnly[str] + + +class OpenAIRealtimeInputAudioTranscriptionCompleted(TypedDict): + type: ReadOnly[Literal["conversation.item.input_audio_transcription.completed"]] + event_id: ReadOnly[str] + item_id: ReadOnly[str] + content_index: ReadOnly[int] + transcript: ReadOnly[str] + + +class OpenAIRealtimeUsageTokenDetails(TypedDict): + audio_tokens: ReadOnly[int] + text_tokens: ReadOnly[int] + cached_tokens: NotRequired[ReadOnly[int]] + + +class OpenAIRealtimeResponseUsage(TypedDict): + input_tokens: ReadOnly[int] + output_tokens: ReadOnly[int] + total_tokens: ReadOnly[int] + input_token_details: NotRequired[ReadOnly[OpenAIRealtimeUsageTokenDetails]] + output_token_details: NotRequired[ReadOnly[OpenAIRealtimeUsageTokenDetails]] + + class OpenAIRealtimeEventTypes(Enum): SESSION_CREATED = "session.created" # Beta delta event names @@ -2199,6 +2235,9 @@ OpenAIRealtimeEvents = ( | OpenAIRealtimeOutputItemDone | OpenAIRealtimeFunctionCallArgumentsDone | OpenAIRealtimeDoneEvent + | OpenAIRealtimeInputAudioBufferSpeechEvent + | OpenAIRealtimeInputAudioTranscriptionDelta + | OpenAIRealtimeInputAudioTranscriptionCompleted ) OpenAIRealtimeStreamList = list[OpenAIRealtimeEvents] diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 77572f69b8b..8fb6108bdda 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -553,6 +553,26 @@ "supports_response_schema": true, "supports_vision": true }, + "amazon.nova-sonic-v1:0": { + "input_cost_per_audio_token": 3.4e-06, + "input_cost_per_token": 6e-08, + "litellm_provider": "bedrock", + "mode": "realtime", + "output_cost_per_audio_token": 1.36e-05, + "output_cost_per_token": 2.4e-07, + "supports_audio_input": true, + "supports_audio_output": true + }, + "amazon.nova-2-sonic-v1:0": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 3.3e-07, + "litellm_provider": "bedrock", + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2.75e-06, + "supports_audio_input": true, + "supports_audio_output": true + }, "amazon.rerank-v1:0": { "input_cost_per_query": 0.001, "input_cost_per_token": 0.0, diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py index ae6b1febd6b..b910018e6c4 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py @@ -827,5 +827,173 @@ class TestBedrockRealtimeSessionEvents: assert event["session"]["modalities"] == ["text", "audio"] +class TestBedrockRealtimeUserEventsAndUsage: + """Regression tests for #38346: USER ASR transcripts, speech boundary events, + usage propagation, and duplicate response.created""" + + @staticmethod + def _run(config, messages): + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_123" + state = { + "session_configuration_request": json.dumps({"configured": True}), + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + } + all_events = [] + for msg in messages: + result = config.transform_realtime_response( + json.dumps(msg), + "amazon.nova-2-sonic-v1:0", + logging_obj, + realtime_response_transform_input=dict(state), + ) + all_events.extend(result["response"]) + state.update( + { + "current_output_item_id": result["current_output_item_id"], + "current_response_id": result["current_response_id"], + "current_conversation_id": result["current_conversation_id"], + "current_delta_chunks": result["current_delta_chunks"], + "current_delta_type": result["current_delta_type"], + } + ) + return all_events + + def test_user_speech_start_and_stop_events(self): + events = self._run( + BedrockRealtimeConfig(), + [{"event": {"userSpeechStart": {}}}, {"event": {"userSpeechEnd": {}}}], + ) + assert [e["type"] for e in events] == [ + "input_audio_buffer.speech_started", + "input_audio_buffer.speech_stopped", + ] + assert all(e["event_id"] and e["item_id"] for e in events) + + def test_user_transcript_emits_input_audio_transcription_events(self): + events = self._run( + BedrockRealtimeConfig(), + [ + { + "event": { + "contentStart": { + "role": "USER", + "type": "TEXT", + "additionalModelFields": json.dumps({"generationStage": "FINAL"}), + } + } + }, + {"event": {"textOutput": {"content": "ready"}}}, + {"event": {"contentEnd": {"stopReason": "PARTIAL_TURN"}}}, + ], + ) + deltas = [e for e in events if e["type"] == "conversation.item.input_audio_transcription.delta"] + completed = [e for e in events if e["type"] == "conversation.item.input_audio_transcription.completed"] + assert len(deltas) == 1 and deltas[0]["delta"] == "ready" + assert len(completed) == 1 and completed[0]["transcript"] == "ready" + assert deltas[0]["item_id"] == completed[0]["item_id"] + assert not any(e["type"] == "response.text.delta" for e in events) + + def test_speculative_user_transcript_emits_delta_only(self): + events = self._run( + BedrockRealtimeConfig(), + [ + { + "event": { + "contentStart": { + "role": "USER", + "type": "TEXT", + "additionalModelFields": json.dumps({"generationStage": "SPECULATIVE"}), + } + } + }, + {"event": {"textOutput": {"content": "rea"}}}, + ], + ) + assert [e["type"] for e in events] == ["conversation.item.input_audio_transcription.delta"] + + def test_user_transcript_state_resets_on_content_end(self): + events = self._run( + BedrockRealtimeConfig(), + [ + {"event": {"contentStart": {"role": "USER", "type": "TEXT"}}}, + {"event": {"contentEnd": {"stopReason": "PARTIAL_TURN"}}}, + {"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}, + {"event": {"textOutput": {"content": "Hi there"}}}, + ], + ) + text_deltas = [e for e in events if e["type"] == "response.text.delta"] + assert len(text_deltas) == 1 and text_deltas[0]["delta"] == "Hi there" + assert not any(e["type"].startswith("conversation.item.input_audio_transcription") for e in events) + + def test_response_created_emitted_once_per_response(self): + events = self._run( + BedrockRealtimeConfig(), + [ + {"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}, + {"event": {"textOutput": {"content": "Hi"}}}, + {"event": {"contentEnd": {"stopReason": "PARTIAL_TURN"}}}, + {"event": {"contentStart": {"role": "ASSISTANT", "type": "AUDIO"}}}, + ], + ) + assert sum(1 for e in events if e["type"] == "response.created") == 1 + + def test_usage_event_propagates_to_response_done(self): + events = self._run( + BedrockRealtimeConfig(), + [ + { + "event": { + "usageEvent": { + "totalInputTokens": 25, + "totalOutputTokens": 40, + "totalTokens": 65, + "details": { + "total": { + "input": {"speechTokens": 20, "textTokens": 5}, + "output": {"speechTokens": 30, "textTokens": 10}, + } + }, + } + } + }, + {"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}, + {"event": {"textOutput": {"content": "Hi"}}}, + {"event": {"contentEnd": {"stopReason": "END_TURN"}}}, + ], + ) + done_events = [e for e in events if e["type"] == "response.done"] + assert len(done_events) == 1 + usage = done_events[0]["response"]["usage"] + assert usage["input_tokens"] == 25 + assert usage["output_tokens"] == 40 + assert usage["total_tokens"] == 65 + assert usage["input_token_details"]["audio_tokens"] == 20 + assert usage["input_token_details"]["text_tokens"] == 5 + assert usage["output_token_details"]["audio_tokens"] == 30 + assert usage["output_token_details"]["text_tokens"] == 10 + + def test_response_done_without_usage_event_reports_zero_usage(self): + events = self._run( + BedrockRealtimeConfig(), + [ + {"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}, + {"event": {"textOutput": {"content": "Hi"}}}, + {"event": {"contentEnd": {"stopReason": "END_TURN"}}}, + ], + ) + done_events = [e for e in events if e["type"] == "response.done"] + assert len(done_events) == 1 + usage = done_events[0]["response"]["usage"] + assert usage["input_tokens"] == 0 + assert usage["output_tokens"] == 0 + assert usage["total_tokens"] == 0 + + if __name__ == "__main__": pytest.main([__file__, "-v"]) From 5eee3bd9f952caa50edfd42927945da86848ed8d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:28:22 +0000 Subject: [PATCH 061/120] fix(bedrock): dispatch success handlers for realtime sessions so spend is logged Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/realtime/handler.py | 12 +++++ .../realtime/test_bedrock_realtime_handler.py | 48 +++++++++++++++++++ 2 files changed, 60 insertions(+) diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 3eeb3cb9fc6..b1b598039f8 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -13,6 +13,8 @@ from pydantic import TypeAdapter from litellm._logging import _redact_string, verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.types.llms.openai import OpenAIRealtimeEvents from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError @@ -154,6 +156,7 @@ class BedrockRealtime(BaseAWSLLM): ) ) + logged_events: Final[list[OpenAIRealtimeEvents]] = [] # mutable-ok: events accumulate across stream loop iterations for spend logging bedrock_to_client_task: Final = asyncio.create_task( self._forward_bedrock_to_client( bedrock_stream, @@ -162,6 +165,7 @@ class BedrockRealtime(BaseAWSLLM): model, logging_obj, session_state, + logged_events, ) ) @@ -172,6 +176,11 @@ class BedrockRealtime(BaseAWSLLM): return_exceptions=True, ) + if logged_events: + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( + logging_obj.dispatch_success_handlers(logged_events, prefer_async_handlers=True) + ) + except Exception as e: verbose_proxy_logger.exception("Error in BedrockRealtime.async_realtime: %s", e) try: @@ -252,6 +261,7 @@ class BedrockRealtime(BaseAWSLLM): model: str, logging_obj: LiteLLMLogging, session_state: dict, + logged_events: "list[OpenAIRealtimeEvents] | None" = None, ): """Forward messages from Bedrock stream to client WebSocket.""" try: @@ -304,6 +314,8 @@ class BedrockRealtime(BaseAWSLLM): # Send transformed messages to client openai_messages = transformed_response.get("response", []) for openai_message in openai_messages: + if logged_events is not None and isinstance(openai_message, dict): + logged_events.append(openai_message) message_json = json.dumps(openai_message) await client_ws.send_text(message_json) verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200]) diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 9efcee192b1..76a564845dd 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -104,6 +104,26 @@ class RealtimeClientWS: self.closed = True +class ScriptedBedrockReceiver: + def __init__(self, payloads): + self._payloads = list(payloads) + + async def receive(self): + if not self._payloads: + return None + payload = self._payloads.pop(0) + return SimpleNamespace(value=SimpleNamespace(bytes_=payload.encode("utf-8"))) + + +class ScriptedBedrockStream: + def __init__(self, payloads): + self.input_stream = FakeInputStream() + self._receiver = ScriptedBedrockReceiver(payloads) + + async def await_output(self): + return (None, self._receiver) + + class ImmediatelyEndingBedrockStream: def __init__(self): self.input_stream = FakeInputStream() @@ -271,6 +291,34 @@ class TestBedrockRealtimeHandler: assert "sessionEnd" in event_names assert stream.input_stream.closed + @pytest.mark.asyncio + async def test_forwarded_events_are_collected_for_spend_logging(self): + handler = BedrockRealtime() + stream = ScriptedBedrockStream( + [ + json.dumps({"event": {"userSpeechStart": {}}}), + json.dumps({"event": {"userSpeechEnd": {}}}), + ] + ) + client_ws = RealtimeClientWS() + logged_events = [] + + await handler._forward_bedrock_to_client( + stream, + client_ws, + BedrockRealtimeConfig(), + "amazon.nova-sonic-v1:0", + FakeLogging(), + {}, + logged_events, + ) + + assert [event["type"] for event in logged_events] == [ + "input_audio_buffer.speech_started", + "input_audio_buffer.speech_stopped", + ] + assert client_ws.closed + @pytest.mark.asyncio async def test_bedrock_stream_end_closes_client_websocket(self): handler = BedrockRealtime() From eee47dcdaa9a3531f5f78e3bb2a70315626ea175 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:39:10 +0000 Subject: [PATCH 062/120] fix(bedrock): share one item_id across a user utterance's realtime events Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/realtime/handler.py | 4 ++- .../llms/bedrock/realtime/transformation.py | 11 +++++-- .../realtime/test_bedrock_realtime_handler.py | 8 ++--- .../test_bedrock_realtime_transformation.py | 31 +++++++++++++++++++ 4 files changed, 45 insertions(+), 9 deletions(-) diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index b1b598039f8..0891379ef8c 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -156,7 +156,9 @@ class BedrockRealtime(BaseAWSLLM): ) ) - logged_events: Final[list[OpenAIRealtimeEvents]] = [] # mutable-ok: events accumulate across stream loop iterations for spend logging + logged_events: Final[ + list[OpenAIRealtimeEvents] + ] = [] # mutable-ok: events accumulate across stream loop iterations for spend logging bedrock_to_client_task: Final = asyncio.create_task( self._forward_bedrock_to_client( bedrock_stream, diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 20996512295..ec5eff0a6fc 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -117,6 +117,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # so the USER/ASSISTANT split from contentStart is tracked here) self._user_transcript_active = False self._user_transcript_generation_stage: str | None = None + self._user_item_id: str | None = None self._latest_usage: OpenAIRealtimeResponseUsage | None = None def validate_environment(self, headers: dict, model: str, api_key: str | None = None) -> dict: @@ -818,13 +819,19 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): stage: Final = parsed.get("generationStage") if isinstance(parsed, dict) else None return stage if isinstance(stage, str) else None + def _current_user_item_id(self, new_utterance: bool = False) -> str: + """Item id shared by all events of one user utterance (speech boundaries and transcript).""" + if new_utterance or self._user_item_id is None: + self._user_item_id = f"item_{uuid.uuid4()}" + return self._user_item_id + def transform_user_speech_event(self, is_speech_start: bool) -> tuple[OpenAIRealtimeEvents, ...]: """Transform Bedrock userSpeechStart/userSpeechEnd to OpenAI speech boundary events.""" verbose_logger.debug("Handling userSpeech%s", "Start" if is_speech_start else "End") speech_event: Final[OpenAIRealtimeInputAudioBufferSpeechEvent] = { "type": "input_audio_buffer.speech_started" if is_speech_start else "input_audio_buffer.speech_stopped", "event_id": f"event_{uuid.uuid4()}", - "item_id": f"item_{uuid.uuid4()}", + "item_id": self._current_user_item_id(new_utterance=is_speech_start), } return (speech_event,) @@ -852,7 +859,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): def transform_user_transcript_event(self, transcript: str) -> tuple[OpenAIRealtimeEvents, ...]: """Transform a USER-role Bedrock textOutput (ASR transcript) to OpenAI transcription events.""" verbose_logger.debug("Handling USER textOutput (ASR transcript)") - item_id: Final = f"item_{uuid.uuid4()}" + item_id: Final = self._current_user_item_id() delta_event: Final[OpenAIRealtimeInputAudioTranscriptionDelta] = { "type": "conversation.item.input_audio_transcription.delta", "event_id": f"event_{uuid.uuid4()}", diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 76a564845dd..aa6573884e2 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -368,9 +368,7 @@ class TestBedrockRealtimeSessionLifecycle: [json.dumps({"type": "session.update", "session": {"instructions": "hi", "modalities": ["text"]}})] ) - await handler._forward_client_to_bedrock( - client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, FakeLogging() - ) + await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, FakeLogging()) acked = [json.loads(message) for message in client_ws.sent_to_client] updated = [event for event in acked if event["type"] == "session.updated"] @@ -382,9 +380,7 @@ class TestBedrockRealtimeSessionLifecycle: handler = BedrockRealtime() config = BedrockRealtimeConfig() stream = FakeBedrockStream() - client_ws = DisconnectingClientWS( - [json.dumps({"type": "session.update", "session": {"instructions": "hi"}})] - ) + client_ws = DisconnectingClientWS([json.dumps({"type": "session.update", "session": {"instructions": "hi"}})]) await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}) diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py index b910018e6c4..dbe84cacb29 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py @@ -874,6 +874,37 @@ class TestBedrockRealtimeUserEventsAndUsage: "input_audio_buffer.speech_stopped", ] assert all(e["event_id"] and e["item_id"] for e in events) + assert events[0]["item_id"] == events[1]["item_id"] + + def test_utterance_lifecycle_shares_one_item_id(self): + events = self._run( + BedrockRealtimeConfig(), + [ + {"event": {"userSpeechStart": {}}}, + {"event": {"userSpeechEnd": {}}}, + { + "event": { + "contentStart": { + "role": "USER", + "type": "TEXT", + "additionalModelFields": json.dumps({"generationStage": "FINAL"}), + } + } + }, + {"event": {"textOutput": {"content": "ready"}}}, + {"event": {"contentEnd": {"stopReason": "PARTIAL_TURN"}}}, + ], + ) + item_ids = {e["item_id"] for e in events if "item_id" in e} + assert len(item_ids) == 1 + + def test_new_utterance_gets_new_item_id(self): + config = BedrockRealtimeConfig() + first = self._run(config, [{"event": {"userSpeechStart": {}}}, {"event": {"userSpeechEnd": {}}}]) + second = self._run(config, [{"event": {"userSpeechStart": {}}}, {"event": {"userSpeechEnd": {}}}]) + assert first[0]["item_id"] == first[1]["item_id"] + assert second[0]["item_id"] == second[1]["item_id"] + assert first[0]["item_id"] != second[0]["item_id"] def test_user_transcript_emits_input_audio_transcription_events(self): events = self._run( From 0a31ab38b9cda4b8b9830a8e1bf1a945033cb1ad Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 15:17:22 -0700 Subject: [PATCH 063/120] fix(speech): stop forwarding response_format as a chat param for Gemini TTS --- .../transformation.py | 70 ++++++++++++------- tests/test_litellm/endpoints/__init__.py | 0 .../test_litellm/endpoints/speech/__init__.py | 0 .../speech_to_completion_bridge/__init__.py | 0 .../test_transformation.py | 60 ++++++++++++++++ 5 files changed, 106 insertions(+), 24 deletions(-) create mode 100644 tests/test_litellm/endpoints/__init__.py create mode 100644 tests/test_litellm/endpoints/speech/__init__.py create mode 100644 tests/test_litellm/endpoints/speech/speech_to_completion_bridge/__init__.py create mode 100644 tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py index fb66edbf272..5e38b644300 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py @@ -1,10 +1,14 @@ +from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast +from typing_extensions import NotRequired, ReadOnly, TypedDict + from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS if TYPE_CHECKING: from litellm import Logging as LiteLLMLoggingObj - from litellm.types.llms.openai import HttpxBinaryResponseContent + from litellm.types.llms.openai import ChatCompletionUserMessage, HttpxBinaryResponseContent from litellm.types.utils import ModelResponse @@ -16,7 +20,42 @@ def _completion_response_cost(model_response: "ModelResponse") -> float | None: return response_cost if isinstance(response_cost, float) else None +GEMINI_TTS_CHAT_AUDIO_FORMAT: Final = "pcm16" + + +class ChatAudioParam(TypedDict): + voice: ReadOnly[str] + format: ReadOnly[NotRequired[str]] + + class SpeechToCompletionBridgeTransformationHandler: + def _chat_completion_params(self, optional_params: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType( + { + param: value + for param, value in optional_params.items() + if param in OPENAI_CHAT_COMPLETION_PARAMS and param != "response_format" + } + ) + + def _chat_audio_format(self, model: str, optional_params: Mapping[str, object]) -> str | None: + if self._is_gemini_tts_model(model): + return GEMINI_TTS_CHAT_AUDIO_FORMAT + response_format: Final = optional_params.get("response_format") + return response_format if isinstance(response_format, str) else None + + def _chat_audio_param( + self, model: str, voice: str | dict | None, optional_params: Mapping[str, object] + ) -> ChatAudioParam | None: + if not isinstance(voice, str): + return None + audio_format: Final = self._chat_audio_format(model, optional_params) + if audio_format is None: + voice_only: Final[ChatAudioParam] = {"voice": voice} + return voice_only + audio: Final[ChatAudioParam] = {"voice": voice, "format": audio_format} + return audio + def transform_request( self, model: str, @@ -28,36 +67,19 @@ class SpeechToCompletionBridgeTransformationHandler: litellm_logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str, ) -> dict: - passed_optional_params: Final = {} - for op in optional_params: - if op in OPENAI_CHAT_COMPLETION_PARAMS: - passed_optional_params[op] = optional_params[op] - - if voice is not None: - if isinstance(voice, str): - passed_optional_params["audio"] = {"voice": voice} - if "response_format" in optional_params: - passed_optional_params["audio"]["format"] = optional_params["response_format"] - - return_kwargs = { + user_message: Final[ChatCompletionUserMessage] = {"role": "user", "content": input} + return_kwargs: Final = { "model": model, - "messages": [ - { - "role": "user", - "content": input, - } - ], + "messages": [user_message], "modalities": ["audio"], - **passed_optional_params, + **self._chat_completion_params(optional_params), + "audio": self._chat_audio_param(model, voice, optional_params), **litellm_params, "headers": headers, "litellm_logging_obj": litellm_logging_obj, "custom_llm_provider": custom_llm_provider, } - - # filter out None values - return_kwargs = {k: v for k, v in return_kwargs.items() if v is not None} - return return_kwargs + return {k: v for k, v in return_kwargs.items() if v is not None} def _convert_pcm16_to_wav(self, pcm_data: bytes, sample_rate: int = 24000, channels: int = 1) -> bytes: """ diff --git a/tests/test_litellm/endpoints/__init__.py b/tests/test_litellm/endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/endpoints/speech/__init__.py b/tests/test_litellm/endpoints/speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/__init__.py b/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py b/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py new file mode 100644 index 00000000000..9367d9d6d16 --- /dev/null +++ b/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py @@ -0,0 +1,60 @@ +from typing import Final +from unittest.mock import MagicMock + +import pytest + +import litellm +from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS +from litellm.endpoints.speech.speech_to_completion_bridge.transformation import ( + SpeechToCompletionBridgeTransformationHandler, +) + +GEMINI_TTS_MODEL: Final = "gemini-3.1-flash-tts-preview" + + +def _bridge_request(response_format: str | None) -> dict: + optional_params: Final = {"temperature": 0.4} if response_format is None else {"temperature": 0.4, "response_format": response_format} + return SpeechToCompletionBridgeTransformationHandler().transform_request( + model=GEMINI_TTS_MODEL, + input="Hello from LiteLLM", + voice="Kore", + optional_params=optional_params, + litellm_params={}, + headers={}, + litellm_logging_obj=MagicMock(), + custom_llm_provider="gemini", + ) + + +@pytest.mark.parametrize("response_format", ["wav", "mp3", "pcm", None]) +def test_gemini_tts_request_keeps_speech_response_format_out_of_chat_params(response_format: str | None) -> None: + request: Final = _bridge_request(response_format) + + assert "response_format" not in request + assert request["audio"] == {"voice": "Kore", "format": "pcm16"} + assert request["temperature"] == 0.4 + assert request["modalities"] == ["audio"] + + gemini_params: Final = litellm.get_optional_params( + model=GEMINI_TTS_MODEL, + custom_llm_provider="gemini", + **{param: value for param, value in request.items() if param in OPENAI_CHAT_COMPLETION_PARAMS}, + ) + assert gemini_params["speechConfig"] == {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}} + assert "responseMimeType" not in gemini_params + + +def test_non_gemini_request_forwards_speech_response_format_as_audio_format() -> None: + request: Final = SpeechToCompletionBridgeTransformationHandler().transform_request( + model="gpt-4o-audio-preview", + input="Hello from LiteLLM", + voice="alloy", + optional_params={"response_format": "wav"}, + litellm_params={}, + headers={}, + litellm_logging_obj=MagicMock(), + custom_llm_provider="openai", + ) + + assert "response_format" not in request + assert request["audio"] == {"voice": "alloy", "format": "wav"} From b34064fe30430079200aea1d244942669d0f590b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 15:21:39 -0700 Subject: [PATCH 064/120] ci: add tests/test_litellm/endpoints to the misc unit-test shard --- .github/workflows/test-unit.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index ed9d8800202..c2dff805772 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -103,6 +103,7 @@ jobs: tests/test_litellm/completion_extras tests/test_litellm/compression tests/test_litellm/containers + tests/test_litellm/endpoints tests/test_litellm/experimental_mcp_client tests/test_litellm/models tests/test_litellm/repositories From f269b52a1591d7c9495bde478acb584928fad2f3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 15:26:05 -0700 Subject: [PATCH 065/120] refactor(speech): type the bridge voice param as a mapping and wrap a long test line --- .../speech/speech_to_completion_bridge/transformation.py | 2 +- .../speech/speech_to_completion_bridge/test_transformation.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py index 5e38b644300..9b757ce86be 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py @@ -45,7 +45,7 @@ class SpeechToCompletionBridgeTransformationHandler: return response_format if isinstance(response_format, str) else None def _chat_audio_param( - self, model: str, voice: str | dict | None, optional_params: Mapping[str, object] + self, model: str, voice: str | Mapping[str, object] | None, optional_params: Mapping[str, object] ) -> ChatAudioParam | None: if not isinstance(voice, str): return None diff --git a/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py b/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py index 9367d9d6d16..c0c720bbaf6 100644 --- a/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py +++ b/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py @@ -13,7 +13,9 @@ GEMINI_TTS_MODEL: Final = "gemini-3.1-flash-tts-preview" def _bridge_request(response_format: str | None) -> dict: - optional_params: Final = {"temperature": 0.4} if response_format is None else {"temperature": 0.4, "response_format": response_format} + optional_params: Final = ( + {"temperature": 0.4} if response_format is None else {"temperature": 0.4, "response_format": response_format} + ) return SpeechToCompletionBridgeTransformationHandler().transform_request( model=GEMINI_TTS_MODEL, input="Hello from LiteLLM", From abfb6adc2b2e153be51a8f829bbe605cce60c12e Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 17:13:26 -0700 Subject: [PATCH 066/120] refactor(proxy): bound the budget window seed by time instead of request ids The one-time seed for a budget window row subtracted the batch's own LiteLLM_SpendLogs rows by request_id, and request_id is the client's x-litellm-call-id whenever the response carries no id of its own. Carrying that set through the queue meant an unbounded, client-controlled aggregate that the commit-failure requeue kept alive across retries. Every log row at or after a batch's earliest start is owed by an increment that still reaches the row, so summing only rows before it needs nothing from the request. That drops request_ids end to end and closes the cross-pod double count the id list could not see. --- .../proxy/db/budget_window_spend_writer.py | 67 +++++------ litellm/proxy/db/db_spend_update_writer.py | 7 +- .../window_spend_update_queue.py | 20 +--- .../proxy/hooks/proxy_track_cost_callback.py | 3 +- litellm/proxy/proxy_server.py | 11 +- .../test_redis_update_buffer.py | 9 +- .../test_window_spend_update_queue.py | 73 ++---------- .../db/test_budget_window_spend_writer.py | 108 ++++++++---------- .../proxy/db/test_db_spend_update_writer.py | 70 ------------ .../hooks/test_proxy_track_cost_callback.py | 60 +--------- tests/test_litellm/proxy/test_proxy_server.py | 55 +-------- 11 files changed, 104 insertions(+), 379 deletions(-) diff --git a/litellm/proxy/db/budget_window_spend_writer.py b/litellm/proxy/db/budget_window_spend_writer.py index f9188f95cfd..8b1ad0e24c7 100644 --- a/litellm/proxy/db/budget_window_spend_writer.py +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -7,17 +7,14 @@ instead of aggregating LiteLLM_SpendLogs every time a window counter goes cold (issue #35766). Raw SQL rather than the Prisma upsert helper because the conditional roll cannot be expressed through the query builder. -Seeding a row that does not exist yet reads LiteLLM_SpendLogs once, excluding -the requests whose increments are in the same batch so neither source counts -them twice. One gap survives that exclusion: without the Redis transaction -buffer every pod flushes its own increments, so a row seeded by one pod can -include spend logs whose increments are still queued on another pod, and those -increments are added again when that pod flushes. That is bounded by a single -flush interval, happens at most once per window row, and only ever over-counts: -the seed never omits spend, because every increment not yet in the row still -reaches it on its own pod's next flush. A row therefore lags real spend by at -most one flush interval of queued increments, the same lag the SpendLogs -aggregate it replaces (and every other spend column) already has. +Seeding a row that does not exist yet reads LiteLLM_SpendLogs once, summing +only rows that started before the batch being flushed so neither source counts +the same request twice. Anything at or after that cutoff is owed by an +increment that still reaches the row, on this pod's next flush or another +pod's, so a row lags real spend by at most one flush interval of queued +increments: the same lag the SpendLogs aggregate it replaces (and every other +spend column) already has. A request whose increment is lost before it flushes, +which today means the pod dying, is missed by both sources and stays missing. """ from collections.abc import Sequence @@ -69,13 +66,13 @@ _ROLL_WINDOW_SPEND_SQL: Final = ( _SEED_FROM_SPEND_LOGS_KEY_SQL: Final = ( 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') " - "AND NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" + "AND \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')" ) _SEED_FROM_SPEND_LOGS_TEAM_SQL: Final = ( 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') " - "AND NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" + "AND \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')" ) _SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL: Final = ( @@ -92,8 +89,8 @@ _UPSERT_TRANSACTION_TIMEOUT: Final = timedelta(seconds=60) class WindowSpendLogsAggregate(Protocol): - """Sums LiteLLM_SpendLogs for one entity since window_start, ignoring the - requests whose ids are handed in. + """Sums LiteLLM_SpendLogs for one entity between window_start and the + batch's earliest request. Injected so the flush can be exercised without a database and so the expensive aggregate stays swappable. @@ -105,21 +102,19 @@ class WindowSpendLogsAggregate(Protocol): entity_type: str, entity_id: str, window_start: datetime, - exclude_request_ids: Sequence[str], - exclude_started_at: datetime | None, + batch_started_at: datetime | None, ) -> float | None: ... -async def spend_logs_total_excluding( +async def spend_logs_total_before_batch( prisma_client: "PrismaClient", entity_type: str, entity_id: str, window_start: datetime, - exclude_request_ids: Sequence[str], - exclude_started_at: datetime | None, + batch_started_at: datetime | None, ) -> float | None: - """LiteLLM_SpendLogs spend for one entity since window_start, minus the - requests already accounted for by the increments being flushed. + """LiteLLM_SpendLogs spend for one entity since window_start, stopping + before the requests the increments being flushed already cover. The spend log writer drains its own queue on a ~2s poll whenever anything is queued, while window increments flush on the much slower batch tick, so @@ -127,13 +122,11 @@ async def spend_logs_total_excluding( already in the table. Counting them in the seed and again in the increment is what made a fresh row land at twice the true spend. - The exclusion is bounded to rows that started at or after the batch's - earliest request. request_id can be chosen by the client - (x-litellm-call-id), so an unbounded exclusion would let a replayed old id - erase a historical row from the seed while its increment still lands. - Without a known start the batch's ids are not excluded at all: that can - only over-count once, which enforcement tolerates, whereas under-counting - is a budget bypass. + Every log row at or after the cutoff belongs to a request whose own + increment still reaches this row, on this pod's next flush or another pod's, + so bounding the sum by time needs nothing from the request itself. Without a + known start the whole window is summed: that can only over-count once, which + enforcement tolerates, whereas under-counting is a budget bypass. """ if entity_type == Litellm_EntityType.KEY.value: bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_KEY_SQL, _SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL @@ -143,13 +136,12 @@ async def spend_logs_total_excluding( return None rows: Final = ( await prisma_client.db.query_raw(unbounded_sql, entity_id, window_start) - if exclude_started_at is None or not exclude_request_ids + if batch_started_at is None else await prisma_client.db.query_raw( bounded_sql, entity_id, window_start, - tuple(exclude_request_ids), - _exclusion_lower_bound(exclude_started_at), + _exclusion_upper_bound(batch_started_at), ) ) if not rows: @@ -157,7 +149,7 @@ async def spend_logs_total_excluding( return float(rows[0].get("total") or 0.0) -def _exclusion_lower_bound(started_at: datetime) -> datetime: +def _exclusion_upper_bound(started_at: datetime) -> datetime: """LiteLLM_SpendLogs.startTime is TIMESTAMP(3); floor to the second so a millisecond rounding of the batch's own earliest row cannot slip under it.""" return to_naive_utc(started_at).replace(microsecond=0) @@ -194,8 +186,8 @@ async def _seed_base_for_missing_row( This is the LiteLLM_SpendLogs aggregate the window counter reseed runs on every cold counter today, but here it runs once per window lifetime and off - the request path, and it excludes this batch's own requests so they are - counted by their increments alone. + the request path, and it stops before the queued increments so they are + counted once. """ if _primary_key(transaction) in existing_primary_keys: return 0.0 @@ -204,8 +196,7 @@ async def _seed_base_for_missing_row( entity_type=transaction["entity_type"], entity_id=transaction["entity_id"], window_start=datetime.fromisoformat(transaction["window_start"]).replace(tzinfo=timezone.utc), - exclude_request_ids=transaction["request_ids"], - exclude_started_at=_transaction_started_at(transaction), + batch_started_at=_transaction_started_at(transaction), ) return float(base or 0.0) @@ -241,7 +232,7 @@ def _upsert_params( async def commit_window_spend_updates( prisma_client: "PrismaClient", transactions: Sequence[WindowSpendTransaction], - spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_total_excluding, + spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_total_before_batch, ) -> None: """Apply aggregated window increments to LiteLLM_BudgetWindowSpend. diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 641c07914d9..b3fd2c3f22c 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -215,11 +215,7 @@ class DBSpendUpdateWriter: start_time: datetime | None, end_time: datetime | None, response_cost: float | None, - ) -> str | None: - """Returns the LiteLLM_SpendLogs request_id this call was recorded - under, so the caller can tell the budget-window writer which log rows - its increments already cover. None when the payload could not be built. - """ + ) -> None: from litellm.proxy.proxy_server import ( disable_spend_logs, litellm_proxy_budget_name, @@ -310,7 +306,6 @@ class DBSpendUpdateWriter: ) verbose_proxy_logger.debug("Runs spend update on all tables") - return payload.get("request_id") except Exception: spend_log_error( "Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue " diff --git a/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py index 04dea66165e..43b069fd8c2 100644 --- a/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py @@ -26,17 +26,11 @@ class WindowSpendTransaction(TypedDict): window_start is an ISO-8601 string rather than a datetime so the transaction survives the JSON round trip through the Redis buffer. - request_ids carries the LiteLLM_SpendLogs ids this spend came from. The - one-time seed for a window that has no row yet subtracts them from its - LiteLLM_SpendLogs aggregate, because the spend log writer flushes on its - own ~2s poll and will usually have persisted these rows before the window - queue flushes; without the exclusion the seed and the increment would each - count them. - - started_at is the earliest request start in the batch. The seed only - subtracts a request_id whose LiteLLM_SpendLogs.startTime is at or after it, - so a client that replays an old id through x-litellm-call-id cannot make the - seed drop the historical row that id already paid for. + started_at is the earliest request start in the batch. The one-time seed for + a window that has no row yet sums only LiteLLM_SpendLogs rows that started + before it, because the spend log writer flushes on its own ~2s poll and will + usually have persisted this batch's rows before the window queue flushes; + without the bound the seed and the increment would each count them. """ entity_type: ReadOnly[str] @@ -44,7 +38,6 @@ class WindowSpendTransaction(TypedDict): window_duration: ReadOnly[str] window_start: ReadOnly[str] spend: ReadOnly[float] - request_ids: ReadOnly[Sequence[str]] started_at: ReadOnly[str | None] @@ -72,7 +65,6 @@ def build_window_spend_transaction( window_duration: str, window_start: datetime, spend: float, - request_id: str | None = None, started_at: datetime | None = None, ) -> WindowSpendTransaction: return WindowSpendTransaction( @@ -81,7 +73,6 @@ def build_window_spend_transaction( window_duration=window_duration, window_start=to_naive_utc(window_start).isoformat(timespec="microseconds"), spend=spend, - request_ids=() if request_id is None else (request_id,), started_at=None if started_at is None else to_naive_utc(started_at.astimezone(timezone.utc)).isoformat(timespec="microseconds"), @@ -101,7 +92,6 @@ def _merge_window_spend_transactions( window_duration=first["window_duration"], window_start=first["window_start"], spend=math.fsum(payload["spend"] for payload in payloads), - request_ids=tuple(sorted(frozenset(chain.from_iterable(payload["request_ids"] for payload in payloads)))), started_at=min(started_ats) if started_ats else None, ) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index f02901f0e97..47aafda2337 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -587,7 +587,7 @@ async def _update_database_and_spend_counters( model_access_groups: Sequence[str] | None = None, ) -> None: try: - spend_log_request_id = await proxy_logging_obj.db_spend_update_writer.update_database( + await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key, response_cost=response_cost, user_id=user_id, @@ -623,7 +623,6 @@ async def _update_database_and_spend_counters( budget_reservation=budget_reservation, end_user_id=end_user_id, tags=request_tags, - request_id=spend_log_request_id, request_started_at=start_time, model_access_groups=model_access_groups, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 654fe4a3f2e..7c47eb1bc42 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2658,7 +2658,6 @@ async def increment_spend_counters( budget_reservation: dict | None = None, end_user_id: str | None = None, tags: list[str] | None = None, - request_id: str | None = None, request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, ): @@ -2733,7 +2732,6 @@ async def increment_spend_counters( window_duration=duration, window_start=key_window_start, increment=cost, - request_id=request_id, request_started_at=request_started_at, ) @@ -2777,7 +2775,6 @@ async def increment_spend_counters( window_duration=duration, window_start=team_window_start, increment=cost, - request_id=request_id, request_started_at=request_started_at, ) @@ -3005,16 +3002,15 @@ async def _enqueue_window_spend_row_update( window_duration: str, window_start: datetime | None, increment: float, - request_id: str | None, request_started_at: datetime | None, ) -> None: """Queue this request's cost against the LiteLLM_BudgetWindowSpend row for the window, so enforcement can read a maintained total instead of aggregating LiteLLM_SpendLogs. - request_id is the LiteLLM_SpendLogs id this cost was recorded under and - request_started_at its startTime; the flush uses them to keep the one-time - seed from counting a request that its increment already covers. + request_started_at is this request's LiteLLM_SpendLogs startTime; the flush + stops the one-time seed there so a request its increment already covers is + not counted twice. Enqueued even when the cache increment was skipped for a reserved counter: the reservation only pre-charged the counter, and the row still owes the @@ -3035,7 +3031,6 @@ async def _enqueue_window_spend_row_update( window_duration=window_duration, window_start=window_start, spend=increment, - request_id=request_id, started_at=request_started_at, ) ) diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 504654e103a..99ac1fe8b50 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -203,7 +203,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buff "window_duration": "30d", "window_start": "2026-08-01T00:00:00.000000", "spend": 3.0, - "request_ids": ["req-1"], + "started_at": None, } ] ) @@ -233,13 +233,11 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buff window_spend, ) = result - # Budget window spend from two pods is summed per window, not overwritten, - # and both pods' request ids reach the seed exclusion. + # Budget window spend from two pods is summed per window, not overwritten. assert window_spend is not None assert len(window_spend) == 1 assert window_spend[0]["spend"] == 6.0 assert window_spend[0]["entity_id"] == "hashed-token" - assert window_spend[0]["request_ids"] == ("req-1",) # Verify db spend was parsed correctly assert db_spend is not None @@ -326,7 +324,6 @@ async def test_restored_window_spend_transactions_drain_back_unchanged(redis_upd window_duration="30d", window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), spend=3.0, - request_id="req-1", started_at=datetime(2026, 8, 10, 12, 0, tzinfo=timezone.utc), ), ) @@ -500,7 +497,6 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_up window_duration="30d", window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), spend=1.25, - request_id="req-1", started_at=datetime(2026, 8, 10, 12, 0, 0, tzinfo=timezone.utc), ) ) @@ -526,7 +522,6 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_up "window_duration": "30d", "window_start": "2026-08-01T00:00:00.000000", "spend": 1.25, - "request_ids": ["req-1"], "started_at": "2026-08-10T12:00:00.000000", } ] diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py index b1ecda57afa..6632b1c8e35 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py @@ -19,7 +19,6 @@ def _txn( spend: float, duration: str = "30d", entity_type: str = "key", - request_id: str | None = None, started_at: datetime | None = None, ): return build_window_spend_transaction( @@ -28,7 +27,6 @@ def _txn( window_duration=duration, window_start=window_start, spend=spend, - request_id=request_id, started_at=started_at, ) @@ -38,13 +36,12 @@ def test_build_window_spend_transaction_stores_naive_utc_iso(): TIMESTAMP(3) column, so a non-UTC input must be converted, not truncated.""" non_utc = datetime(2026, 8, 1, 20, 0, tzinfo=timezone(timedelta(hours=-4))) - assert _txn("k1", non_utc, 1.0, request_id="req-1") == { + assert _txn("k1", non_utc, 1.0) == { "entity_type": "key", "entity_id": "k1", "window_duration": "30d", "window_start": "2026-08-02T00:00:00.000000", "spend": 1.0, - "request_ids": ("req-1",), "started_at": None, } @@ -59,19 +56,19 @@ def test_build_window_spend_transaction_stores_started_at_as_naive_utc_iso(): @pytest.mark.asyncio async def test_aggregation_keeps_the_earliest_started_at_of_the_batch(): - """The seed bounds its request-id exclusion at the batch's earliest start, - so a later start must never win the merge.""" + """The seed stops at the batch's earliest start, so a later start must never + win the merge: it would push the cutoff forward and count a request the + increments already cover.""" queue = WindowSpendUpdateQueue() earliest = datetime(2026, 8, 10, 12, 0, 0, tzinfo=timezone.utc) - await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-2", started_at=earliest + timedelta(seconds=5))) - await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-1", started_at=earliest)) - await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-3")) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, started_at=earliest + timedelta(seconds=5))) + await queue.add_update(_txn("k1", WINDOW_A, 1.0, started_at=earliest)) + await queue.add_update(_txn("k1", WINDOW_A, 1.0)) aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() assert len(aggregated) == 1 assert aggregated[0]["started_at"] == "2026-08-10T12:00:00.000000" - assert aggregated[0]["request_ids"] == ("req-1", "req-2", "req-3") def test_to_naive_utc_leaves_naive_values_alone(): @@ -208,62 +205,12 @@ def test_aggregation_survives_the_redis_json_round_trip(): assert reloaded == aggregated -@pytest.mark.asyncio -async def test_aggregation_unions_the_request_ids_of_merged_increments(): - """The seed excludes exactly the requests its batch already covers, so every - merged increment's id has to survive aggregation.""" - queue = WindowSpendUpdateQueue() - await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-1")) - await queue.add_update(_txn("k1", WINDOW_A, 2.0, request_id="req-2")) - - aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() - - assert len(aggregated) == 1 - assert aggregated[0]["request_ids"] == ("req-1", "req-2") - - -@pytest.mark.asyncio -async def test_request_ids_stay_with_their_own_window(): - queue = WindowSpendUpdateQueue() - await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a")) - await queue.add_update(_txn("k1", WINDOW_B, 2.0, request_id="req-b")) - - aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() - - assert {payload["window_start"]: payload["request_ids"] for payload in aggregated} == { - "2026-08-01T00:00:00.000000": ("req-a",), - "2026-08-31T00:00:00.000000": ("req-b",), - } - - -@pytest.mark.asyncio -async def test_request_ids_are_deduplicated_and_ordered(): - queue = WindowSpendUpdateQueue() - await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-b")) - await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a")) - await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a")) - - aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() - - assert aggregated[0]["request_ids"] == ("req-a", "req-b") - - -@pytest.mark.asyncio -async def test_increment_without_a_request_id_carries_no_exclusion(): - queue = WindowSpendUpdateQueue() - await queue.add_update(_txn("k1", WINDOW_A, 1.0)) - - aggregated = await queue.flush_and_get_aggregated_window_spend_transactions() - - assert aggregated[0]["request_ids"] == () - - -def test_request_ids_survive_the_redis_json_round_trip(): +def test_started_at_survives_the_redis_json_round_trip(): aggregated = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( - [(_txn("k1", WINDOW_A, 1.0, request_id="req-1"),)] + [(_txn("k1", WINDOW_A, 1.0, started_at=datetime(2026, 8, 10, 12, 0, tzinfo=timezone.utc)),)] ) reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions([json.loads(json.dumps(aggregated))]) - assert reloaded[0]["request_ids"] == ("req-1",) + assert reloaded[0]["started_at"] == "2026-08-10T12:00:00.000000" assert reloaded[0]["spend"] == 1.0 diff --git a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py index a849317c930..74b71f63401 100644 --- a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py +++ b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py @@ -8,7 +8,7 @@ import pytest from litellm.proxy.db.budget_window_spend_writer import ( commit_window_spend_updates, roll_window_spend_row, - spend_logs_total_excluding, + spend_logs_total_before_batch, ) from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( build_window_spend_transaction, @@ -82,16 +82,14 @@ class _RecordingAggregate: entity_type: str, entity_id: str, window_start: datetime, - exclude_request_ids: Any, - exclude_started_at: datetime | None, + batch_started_at: datetime | None, ) -> float | None: self.calls.append( { "entity_type": entity_type, "entity_id": entity_id, "window_start": window_start, - "exclude_request_ids": tuple(exclude_request_ids), - "exclude_started_at": exclude_started_at, + "batch_started_at": batch_started_at, } ) return self.value @@ -99,8 +97,8 @@ class _RecordingAggregate: class _SpendLogsFake: """Sums the LiteLLM_SpendLogs rows (request_id, spend, startTime) it holds, - honouring the exclusion exactly as the real aggregate's - NOT (request_id = ANY(...) AND startTime >= bound) does.""" + honouring the cutoff exactly as the real aggregate's + startTime < bound does.""" def __init__(self, rows: tuple[tuple[str, float, datetime], ...]) -> None: self.rows = rows @@ -111,25 +109,22 @@ class _SpendLogsFake: entity_type: str, entity_id: str, window_start: datetime, - exclude_request_ids: Any, - exclude_started_at: datetime | None, + batch_started_at: datetime | None, ) -> float | None: - excluded = frozenset(exclude_request_ids) if exclude_started_at is not None else frozenset() return math.fsum( spend - for request_id, spend, started_at in self.rows - if not (request_id in excluded and started_at >= exclude_started_at) + for _request_id, spend, started_at in self.rows + if batch_started_at is None or started_at < batch_started_at ) -def _batch(request_ids: tuple[str, ...], spend: float, started_at: datetime | None = BATCH_STARTED_AT) -> dict: +def _batch(spend: float, started_at: datetime | None = BATCH_STARTED_AT) -> dict: return { "entity_type": "key", "entity_id": "k1", "window_duration": "30d", "window_start": "2026-08-01T00:00:00.000000", "spend": spend, - "request_ids": request_ids, "started_at": None if started_at is None else started_at.replace(tzinfo=None).isoformat(timespec="microseconds"), @@ -345,7 +340,7 @@ async def test_unknown_entity_type_contributes_no_seed(): db = _FakeDB(existing_rows=[]) async def no_such_column( - prisma_client, entity_type, entity_id, window_start, exclude_request_ids, exclude_started_at + prisma_client, entity_type, entity_id, window_start, batch_started_at ): return None @@ -363,7 +358,7 @@ async def test_unknown_entity_type_contributes_no_seed(): async def test_unavailable_spend_logs_aggregate_seeds_zero_rather_than_failing(): db = _FakeDB(existing_rows=[]) - async def unavailable(prisma_client, entity_type, entity_id, window_start, exclude_request_ids, exclude_started_at): + async def unavailable(prisma_client, entity_type, entity_id, window_start, batch_started_at): return None await commit_window_spend_updates( @@ -399,18 +394,17 @@ async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_o @pytest.mark.asyncio -async def test_seed_receives_the_batch_request_ids_and_earliest_start_to_exclude(): +async def test_seed_receives_the_batch_earliest_start_as_its_cutoff(): db = _FakeDB(existing_rows=[]) aggregate = _RecordingAggregate(value=0.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), - transactions=(_batch(("req-1", "req-2", "req-3"), 3.0),), + transactions=(_batch(3.0),), spend_logs_aggregate=aggregate, ) - assert aggregate.calls[0]["exclude_request_ids"] == ("req-1", "req-2", "req-3") - assert aggregate.calls[0]["exclude_started_at"] == BATCH_STARTED_AT + assert aggregate.calls[0]["batch_started_at"] == BATCH_STARTED_AT @pytest.mark.asyncio @@ -420,11 +414,11 @@ async def test_seed_passes_no_start_bound_when_the_batch_has_none(): await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), - transactions=(_batch(("req-1",), 1.0, started_at=None),), + transactions=(_batch(1.0, started_at=None),), spend_logs_aggregate=aggregate, ) - assert aggregate.calls[0]["exclude_started_at"] is None + assert aggregate.calls[0]["batch_started_at"] is None @pytest.mark.asyncio @@ -444,7 +438,7 @@ async def test_new_row_is_not_double_counted_when_the_batch_logs_already_flushed await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), - transactions=(_batch(("req-1", "req-2", "req-3"), 0.000141),), + transactions=(_batch(0.000141),), spend_logs_aggregate=already_flushed, ) @@ -460,7 +454,7 @@ async def test_new_row_still_covers_spend_that_predates_the_batch(): await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), - transactions=(_batch(("req-1",), 0.000047),), + transactions=(_batch(0.000047),), spend_logs_aggregate=spend_logs, ) @@ -469,17 +463,20 @@ async def test_new_row_still_covers_spend_that_predates_the_batch(): @pytest.mark.asyncio -async def test_replayed_request_id_cannot_erase_historical_spend_from_the_seed(): - """request_id can be chosen by the client via x-litellm-call-id. A request - that replays an id from before this batch writes no new LiteLLM_SpendLogs - row (the insert skips duplicates), so the seed must keep counting the - historical row that id belongs to; only its increment is new.""" +async def test_seed_skips_logs_from_requests_this_batch_never_saw(): + """A concurrent request on another pod can land its spend log before this + pod seeds the row. Its increment is still queued over there, so the cutoff + has to drop it from the seed even though this batch has no way to know its + id; counting it here and again on that pod's flush is the double count the + old id list could not catch.""" db = _FakeDB(existing_rows=[]) - spend_logs = _SpendLogsFake(rows=(("replayed", 0.5, BEFORE_BATCH),)) + spend_logs = _SpendLogsFake( + rows=(("older", 0.5, BEFORE_BATCH), ("other-pod", 0.25, BATCH_STARTED_AT + timedelta(seconds=1))), + ) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), - transactions=(_batch(("replayed",), 0.000047),), + transactions=(_batch(0.000047),), spend_logs_aggregate=spend_logs, ) @@ -496,7 +493,7 @@ async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet(): await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), - transactions=(_batch(("req-1", "req-2", "req-3"), 0.000141),), + transactions=(_batch(0.000141),), spend_logs_aggregate=nothing_flushed, ) @@ -509,57 +506,48 @@ async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet(): "entity_type, expected_column", [("key", "api_key = $1"), ("team", "team_id = $1")], ) -async def test_seed_aggregate_sql_excludes_the_request_ids_only_within_the_batch_start_bound( - entity_type, expected_column -): +async def test_seed_aggregate_sql_stops_at_the_batch_start(entity_type, expected_column): db = _FakeDB(existing_rows=[{"total": 1.25}]) - total = await spend_logs_total_excluding( + total = await spend_logs_total_before_batch( prisma_client=_FakePrismaClient(db), entity_type=entity_type, entity_id="e1", window_start=WINDOW_A, - exclude_request_ids=("req-1", "req-2"), - exclude_started_at=BATCH_STARTED_AT, + batch_started_at=BATCH_STARTED_AT, ) assert total == pytest.approx(1.25) ((query, params),) = db.query_raw_calls normalized = " ".join(query.split()) assert expected_column in normalized - assert "NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" in normalized + assert "AND \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')" in normalized assert 'FROM "LiteLLM_SpendLogs"' in normalized # startTime is TIMESTAMP(3): the bound is floored to the second so the # batch's own earliest row cannot round under it. - assert params == ("e1", WINDOW_A, ("req-1", "req-2"), datetime(2026, 8, 10, 12, 0, 0)) - # The ids are bound, never spliced into the statement. - assert "req-1" not in query + assert params == ("e1", WINDOW_A, datetime(2026, 8, 10, 12, 0, 0)) + # Nothing the caller supplied reaches the statement text. + assert "e1" not in query @pytest.mark.asyncio -@pytest.mark.parametrize( - "exclude_request_ids, exclude_started_at", - [(("req-1",), None), ((), BATCH_STARTED_AT)], -) -async def test_seed_aggregate_excludes_nothing_without_both_ids_and_a_start_bound( - exclude_request_ids, exclude_started_at -): - """Ids without a start bound would reopen the replayed-id hole, so the - seed counts everything instead; at worst that over-counts one batch.""" +async def test_seed_aggregate_sums_the_whole_window_without_a_start_bound(): + """A batch with no known start cannot place the cutoff, so the seed counts + everything; at worst that over-counts one batch, which enforcement + tolerates, where under-counting is a budget bypass.""" db = _FakeDB(existing_rows=[{"total": 1.25}]) - total = await spend_logs_total_excluding( + total = await spend_logs_total_before_batch( prisma_client=_FakePrismaClient(db), entity_type="key", entity_id="e1", window_start=WINDOW_A, - exclude_request_ids=exclude_request_ids, - exclude_started_at=exclude_started_at, + batch_started_at=None, ) assert total == pytest.approx(1.25) ((query, params),) = db.query_raw_calls - assert "request_id" not in query + assert '"startTime" <' not in query assert params == ("e1", WINDOW_A) @@ -567,13 +555,12 @@ async def test_seed_aggregate_excludes_nothing_without_both_ids_and_a_start_boun async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs_column(): db = _FakeDB(existing_rows=[]) - total = await spend_logs_total_excluding( + total = await spend_logs_total_before_batch( prisma_client=_FakePrismaClient(db), entity_type="user", entity_id="u1", window_start=WINDOW_A, - exclude_request_ids=(), - exclude_started_at=None, + batch_started_at=None, ) assert total is None @@ -584,13 +571,12 @@ async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs async def test_seed_aggregate_treats_an_entity_with_no_rows_as_zero(): db = _FakeDB(existing_rows=[]) - total = await spend_logs_total_excluding( + total = await spend_logs_total_before_batch( prisma_client=_FakePrismaClient(db), entity_type="key", entity_id="k-unknown", window_start=WINDOW_A, - exclude_request_ids=(), - exclude_started_at=None, + batch_started_at=None, ) assert total == 0.0 diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index d28cf8c9c6a..11ef911de3e 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2581,7 +2581,6 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_ window_duration="30d", window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), spend=0.5, - request_id="req-1", ) await db_writer.window_spend_update_queue.add_update(transaction) db = _WindowSpendFakeDB() @@ -2611,7 +2610,6 @@ async def test_failed_window_spend_commit_from_redis_is_restored_to_redis(): window_duration="7d", window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), spend=2.0, - request_id="req-1", ), ) mock_redis_update_buffer = AsyncMock() @@ -2638,74 +2636,6 @@ async def test_failed_window_spend_commit_from_redis_is_restored_to_redis(): db_writer.pod_lock_manager.release_lock.assert_awaited_once() -@pytest.mark.asyncio -async def test_update_database_returns_the_spend_log_request_id(): - """The budget-window seed excludes the log rows its increments already - cover, so the caller needs the id this call was recorded under. It cannot - be re-derived: cache hits append time.time() to the id.""" - db_writer = DBSpendUpdateWriter() - db_writer._insert_spend_log_to_db = AsyncMock() - db_writer._enqueue_tool_usage_transaction = AsyncMock() - - with ( - patch.multiple( # test-quality-ok: update_database lazily imports these proxy_server globals; no injection seam - "litellm.proxy.proxy_server", - disable_spend_logs=False, - prisma_client=MagicMock(), - litellm_proxy_budget_name="test-budget", - ) - ): - request_id = await db_writer.update_database( - token="test-token", - user_id="test-user", - end_user_id=None, - team_id="test-team", - org_id=None, - kwargs={"model": "gpt-4", "custom_llm_provider": "openai", "litellm_call_id": "call-xyz"}, - completion_response=MagicMock(), - start_time=datetime.now(), - end_time=datetime.now(), - response_cost=0.1, - ) - await asyncio.sleep(0) - - assert request_id is not None - # Same id the spend log row was queued under. - assert request_id == db_writer._insert_spend_log_to_db.call_args[1]["payload"]["request_id"] - - -@pytest.mark.asyncio -async def test_update_database_returns_none_when_the_payload_cannot_be_built(): - db_writer = DBSpendUpdateWriter() - - with ( - patch.multiple( # test-quality-ok: update_database lazily imports these proxy_server globals; no injection seam - "litellm.proxy.proxy_server", - disable_spend_logs=False, - prisma_client=MagicMock(), - litellm_proxy_budget_name="test-budget", - ), - patch( # test-quality-ok: the payload builder is called by name inside update_database; no injection seam - "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", - side_effect=Exception("payload boom"), - ), - ): - request_id = await db_writer.update_database( - token="test-token", - user_id="test-user", - end_user_id=None, - team_id="test-team", - org_id=None, - kwargs={"model": "gpt-4"}, - completion_response=MagicMock(), - start_time=datetime.now(), - end_time=datetime.now(), - response_cost=0.1, - ) - - assert request_id is None - - @pytest.mark.asyncio async def test_commit_spend_updates_to_db_does_not_stamp_key_settings_updated_at(): """Spend flushes must leave settings_updated_at alone, or it decays into diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 2a540f4f522..8043a1aca3f 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -567,9 +567,7 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re @pytest.mark.asyncio async def test_update_database_and_spend_counters_updates_counters_after_db_update(): proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( - return_value="chatcmpl-abc123" - ) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} start_time = datetime.now() @@ -602,7 +600,6 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda budget_reservation=budget_reservation, end_user_id="test_end_user_id", tags=["tag-a"], - request_id="chatcmpl-abc123", request_started_at=start_time, model_access_groups=("premium",), ) @@ -1884,61 +1881,6 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( ) -@pytest.mark.asyncio -async def test_update_database_and_spend_counters_forwards_the_spend_log_request_id(): - """The budget-window flush excludes the log rows its increments already - cover. That only works if the id update_database recorded the row under is - handed to the counter update, so this seam is load-bearing.""" - proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( - return_value="chatcmpl-abc123" - ) - increment_spend_counters = AsyncMock() - - await _update_database_and_spend_counters( - proxy_logging_obj=proxy_logging_obj, - increment_spend_counters=increment_spend_counters, - user_api_key="test_api_key", - user_id="test_user_id", - end_user_id=None, - team_id="test_team_id", - org_id="test_org_id", - kwargs={}, - completion_response=None, - start_time=datetime.now(), - end_time=datetime.now(), - response_cost=0.2, - budget_reservation=None, - ) - - assert increment_spend_counters.await_args.kwargs["request_id"] == "chatcmpl-abc123" - - -@pytest.mark.asyncio -async def test_update_database_and_spend_counters_forwards_a_missing_request_id_as_none(): - proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(return_value=None) - increment_spend_counters = AsyncMock() - - await _update_database_and_spend_counters( - proxy_logging_obj=proxy_logging_obj, - increment_spend_counters=increment_spend_counters, - user_api_key="test_api_key", - user_id="test_user_id", - end_user_id=None, - team_id="test_team_id", - org_id="test_org_id", - kwargs={}, - completion_response=None, - start_time=datetime.now(), - end_time=datetime.now(), - response_cost=0.2, - budget_reservation=None, - ) - - assert increment_spend_counters.await_args.kwargs["request_id"] is None - - class _FakeDeploymentLookup: """Deployment lookup returning the access groups each deployment declares.""" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 1de3ed6e56d..43280258153 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11399,35 +11399,10 @@ async def test_no_window_spend_row_enqueued_without_budget_limits(): assert enqueued == [] -@pytest.mark.asyncio -async def test_window_spend_row_carries_the_spend_log_request_id(): - """The flush excludes these ids from its one-time seed, so the id threaded - here has to be the same one the LiteLLM_SpendLogs row was written under.""" - from litellm.proxy.proxy_server import increment_spend_counters - - reset_at = datetime.now(timezone.utc) + timedelta(days=10) - key_obj = MagicMock() - key_obj.budget_limits = [ - {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} - ] - - with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: - await increment_spend_counters( - token="hashed-token", - team_id=None, - user_id=None, - response_cost=0.25, - request_id="chatcmpl-abc123", - ) - enqueued = await _drain(queue) - - assert enqueued[0]["request_ids"] == ("chatcmpl-abc123",) - - @pytest.mark.asyncio async def test_window_spend_row_carries_the_request_start_time(): - """The seed only excludes a batch id whose LiteLLM_SpendLogs.startTime is at - or after this, so it must be the same start the spend log was written with.""" + """The seed sums LiteLLM_SpendLogs only up to this point, so it must be the + same start the spend log row was written with.""" from litellm.proxy.proxy_server import increment_spend_counters reset_at = datetime.now(timezone.utc) + timedelta(days=10) @@ -11442,7 +11417,6 @@ async def test_window_spend_row_carries_the_request_start_time(): team_id=None, user_id=None, response_cost=0.25, - request_id="chatcmpl-abc123", request_started_at=datetime(2026, 8, 10, 12, 0, 0, 500_000, tzinfo=timezone.utc), ) enqueued = await _drain(queue) @@ -11451,26 +11425,7 @@ async def test_window_spend_row_carries_the_request_start_time(): @pytest.mark.asyncio -async def test_window_spend_row_without_a_request_id_excludes_nothing(): - from litellm.proxy.proxy_server import increment_spend_counters - - reset_at = datetime.now(timezone.utc) + timedelta(days=10) - key_obj = MagicMock() - key_obj.budget_limits = [ - {"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()} - ] - - with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue: - await increment_spend_counters( - token="hashed-token", team_id=None, user_id=None, response_cost=0.25 - ) - enqueued = await _drain(queue) - - assert enqueued[0]["request_ids"] == () - - -@pytest.mark.asyncio -async def test_team_window_spend_row_carries_the_request_id(): +async def test_team_window_spend_row_carries_the_request_start_time(): from litellm.proxy.proxy_server import increment_spend_counters reset_at = datetime.now(timezone.utc) + timedelta(days=3) @@ -11485,11 +11440,11 @@ async def test_team_window_spend_row_carries_the_request_id(): team_id="team-1", user_id=None, response_cost=1.5, - request_id="chatcmpl-team", + request_started_at=datetime(2026, 8, 10, 12, 0, 0, 500_000, tzinfo=timezone.utc), ) enqueued = await _drain(queue) - assert enqueued[0]["request_ids"] == ("chatcmpl-team",) + assert enqueued[0]["started_at"] == "2026-08-10T12:00:00.500000" def _mock_startup_prisma_client(health_check_error=None, connect_error=None): From 26e71ddc5452484903175bcf2c27211f02c7e714 Mon Sep 17 00:00:00 2001 From: samzong Date: Sun, 2 Aug 2026 10:06:15 -0400 Subject: [PATCH 067/120] fix(proxy): serialize model block responses Signed-off-by: samzong --- litellm/models/model.py | 2 + .../test_model_tag_accessgroup_e2e.py | 51 +++++++++---------- tests/test_litellm/models/test_models.py | 28 ++++++++++ 3 files changed, 55 insertions(+), 26 deletions(-) diff --git a/litellm/models/model.py b/litellm/models/model.py index 209f26d4837..a0c840341ab 100644 --- a/litellm/models/model.py +++ b/litellm/models/model.py @@ -29,6 +29,8 @@ class LiteLLM_ProxyModelTable(LiteLLMPydanticObjectBase): @model_validator(mode="before") @classmethod def check_potential_json_str(cls, values): + if not isinstance(values, dict): + return values if isinstance(values.get("litellm_params"), str): try: values["litellm_params"] = json.loads(values["litellm_params"]) diff --git a/tests/e2e/management/test_model_tag_accessgroup_e2e.py b/tests/e2e/management/test_model_tag_accessgroup_e2e.py index e6a187ae105..63e51e4abcb 100644 --- a/tests/e2e/management/test_model_tag_accessgroup_e2e.py +++ b/tests/e2e/management/test_model_tag_accessgroup_e2e.py @@ -180,6 +180,12 @@ class ModelBlockBody(BaseModel): model_id: str +class ModelBlockResponse(BaseModel): + model_config = ConfigDict(protected_namespaces=()) + model_id: str + blocked: bool + + class ModelInfoBlockDetail(BaseModel): id: str | None = None blocked: bool | None = None @@ -245,11 +251,6 @@ class TestModelRoutes: def test_block_then_unblock_persists_to_model_info( self, client: ManagementClient, resources: ResourceManager ) -> None: - """The blocked flag's persistence is read back from /model/info, not from the - /model/block response: that route currently returns a non-2xx serialization - envelope even though the DB write lands, so the /model/info read-back is the - authoritative persistence contract and keeps this test valid once the - response shape is fixed.""" model_name = f"e2e-mgmt-model-block-{unique_marker()}" model_id = _create_db_model(client, resources, model_name) @@ -257,27 +258,25 @@ class TestModelRoutes: f"{model_name!r} already reports blocked in /model/info before /model/block ran" ) - _ = client.proxy.transport.send( - "/model/block", - headers=client.proxy.transport.master, - json=ModelBlockBody(model_id=model_id), - ) - _ = _poll( - client.proxy, - lambda: True if _model_blocked_flag(client, model_id) is True else None, - f"/model/info never reported {model_name!r} blocked after /model/block", - ) - - _ = client.proxy.transport.send( - "/model/unblock", - headers=client.proxy.transport.master, - json=ModelBlockBody(model_id=model_id), - ) - _ = _poll( - client.proxy, - lambda: True if _model_blocked_flag(client, model_id) is not True else None, - f"/model/info never cleared blocked for {model_name!r} after /model/unblock", - ) + for action, expected in (("block", True), ("unblock", False)): + response = unwrap( + client.proxy.transport.post( + f"/model/{action}", + headers=client.proxy.transport.master, + json=ModelBlockBody(model_id=model_id), + response_type=ModelBlockResponse, + ) + ) + assert response.model_id == model_id + assert response.blocked is expected + _ = _poll( + client.proxy, + lambda: True + if _model_blocked_flag(client, model_id) is expected + else None, + f"/model/info never reported blocked={expected} for {model_name!r} " + f"after /model/{action}", + ) class TestTagRoutes: diff --git a/tests/test_litellm/models/test_models.py b/tests/test_litellm/models/test_models.py index 669dba8e466..9ae9b732066 100644 --- a/tests/test_litellm/models/test_models.py +++ b/tests/test_litellm/models/test_models.py @@ -5,6 +5,7 @@ Tests for backend domain models. from datetime import datetime import pytest +from pydantic import BaseModel, TypeAdapter from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.models.budget import ( @@ -130,6 +131,33 @@ class TestModel: assert model.litellm_params == {"model": "gpt-4"} assert model.model_info == {"team_id": "t1"} + def test_response_type_adapter_accepts_pydantic_row(self): + class PrismaModelRow(BaseModel): + model_id: str + model_name: str + litellm_params: dict[str, str] + model_info: dict[str, str] | None = None + blocked: bool = False + + row = PrismaModelRow( + model_id="m1", + model_name="gpt-4", + litellm_params={"model": "gpt-4"}, + model_info={"team_id": "t1"}, + blocked=True, + ) + + model = TypeAdapter(LiteLLM_ProxyModelTable | None).validate_python( + row, + from_attributes=True, + ) + + assert model is not None + assert model.model_id == "m1" + assert model.litellm_params == {"model": "gpt-4"} + assert model.model_info == {"team_id": "t1"} + assert model.blocked is True + def test_team_helpers_none_when_no_model_info(self): model = LiteLLM_ProxyModelTable( model_id="m1", model_name="gpt-4", litellm_params={}, model_info=None From fd72ae830c3535ec206869dd7b682eded13c1f25 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 21:39:07 -0700 Subject: [PATCH 068/120] test(model_management): drive /model/block and /model/unblock through response serialization The route-level regression test returns a real prisma row from a mocked update and asserts both routes serialize it to a 200 with the toggled blocked flag, which is exactly the path that raised AttributeError before the validator guard. Also binds the loop variable in the e2e poll lambda (ruff B023). --- .../test_model_tag_accessgroup_e2e.py | 4 +- .../test_model_management_endpoints.py | 68 +++++++++++++++++++ 2 files changed, 70 insertions(+), 2 deletions(-) diff --git a/tests/e2e/management/test_model_tag_accessgroup_e2e.py b/tests/e2e/management/test_model_tag_accessgroup_e2e.py index 63e51e4abcb..eb3a6093c69 100644 --- a/tests/e2e/management/test_model_tag_accessgroup_e2e.py +++ b/tests/e2e/management/test_model_tag_accessgroup_e2e.py @@ -271,8 +271,8 @@ class TestModelRoutes: assert response.blocked is expected _ = _poll( client.proxy, - lambda: True - if _model_blocked_flag(client, model_id) is expected + lambda want=expected: True + if _model_blocked_flag(client, model_id) is want else None, f"/model/info never reported blocked={expected} for {model_name!r} " f"after /model/{action}", diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index dc9fede1f65..13c8a26377c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -4466,3 +4466,71 @@ class TestEnforceRpmTpmOnModelAdd: _raise_if_rate_limits_required_but_missing(litellm_params=params, enforced=True) assert expected_missing in str(exc_info.value.message) assert exc_info.value.code == "400" + + +class TestBlockModelResponseSerialization: + """POST /model/block and /model/unblock return the raw prisma row through this + route's `LiteLLM_ProxyModelTable | None` response validation. The row is not a + dict, so the dict-assuming before-validator used to raise AttributeError inside + FastAPI's serialization layer: a 500 for the caller after the DB write already + landed. The routes must serialize the row to a 200 with the updated blocked flag.""" + + @pytest.mark.parametrize( + ("route", "blocked"), [("/model/block", True), ("/model/unblock", False)] + ) + def test_block_routes_serialize_prisma_row_to_200(self, route, blocked): + from datetime import datetime, timezone + + from prisma import models as prisma_models + + import litellm.proxy.proxy_server as ps + from litellm.proxy.proxy_server import app + + written_at = datetime(2026, 8, 29, tzinfo=timezone.utc) + row_fields = { + "model_id": "m-block-1", + "model_name": "gpt-4o-mini", + "litellm_params": json.dumps({"model": "openai/gpt-4o-mini", "api_key": "encrypted-value"}), + "model_info": json.dumps({"id": "m-block-1"}), + "created_at": written_at, + "created_by": "admin", + "updated_at": written_at, + "updated_by": "admin", + } + existing_row = prisma_models.LiteLLM_ProxyModelTable(blocked=not blocked, **row_fields) + updated_row = prisma_models.LiteLLM_ProxyModelTable(blocked=blocked, **row_fields) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row) + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + app.dependency_overrides[ps.user_api_key_auth] = lambda: admin + try: + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point + "litellm.proxy.proxy_server.llm_router", + MagicMock(**{"get_model_ids.return_value": ["m-block-1"]}), + ), + patch("litellm.proxy.proxy_server.redis_usage_cache", None), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch( # test-quality-ok: stubs the cache write so the test observes only response serialization + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + patch( # test-quality-ok: audit logging is a background side effect outside this test's contract + "litellm.proxy.management_endpoints.model_management_endpoints.create_object_audit_log", + new=AsyncMock(return_value=None), + ), + ): + client = TestClient(app) + response = client.post(route, json={"model_id": "m-block-1"}) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + body = response.json() + assert body["model_id"] == "m-block-1" + assert body["blocked"] is blocked + assert body["litellm_params"] == {"model": "openai/gpt-4o-mini", "api_key": "encrypted-value"} From b0ce17c755a7023fb85e91ec0cefc5e24bd4c70d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 22:08:54 -0700 Subject: [PATCH 069/120] fix(gigachat): generic env-credential passthrough fallback plus type hardening - forward unrouted /gigachat/* requests with env credentials like other passthrough providers (the old fallback returned 400 on any request without a routed model, /gigachat/models included) - fix basedpyright budget breaches across the gigachat provider, common_request_processing, and llm_passthrough_endpoints with real narrowing, no new suppressions - add regression tests for the fallback target, auth header, and model-less endpoints --- litellm/litellm_core_utils/litellm_logging.py | 43 ++++---- litellm/llms/gigachat/chat/streaming.py | 10 +- litellm/llms/gigachat/chat/transformation.py | 33 ++++--- .../gigachat/passthrough/transformation.py | 13 ++- litellm/proxy/common_request_processing.py | 11 ++- .../llm_passthrough_endpoints.py | 61 ++++-------- .../test_gigachat_embedding_transformation.py | 3 +- .../test_llm_pass_through_endpoints.py | 98 +++++++++++++++---- 8 files changed, 168 insertions(+), 104 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 53ec3983242..e34c647efc2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -14,7 +14,7 @@ from collections.abc import Callable, Mapping, Sequence from datetime import datetime as dt_object from functools import lru_cache from types import MappingProxyType, TracebackType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast from httpx import Response from pydantic import BaseModel @@ -86,7 +86,6 @@ from litellm.litellm_core_utils.redact_messages import ( redact_streaming_responses_for_custom_logger, ) from litellm.llms.base_llm.ocr.transformation import OCRResponse -from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.agents import LiteLLMSendMessageResponse @@ -1603,24 +1602,26 @@ class Logging(LiteLLMLoggingBaseClass): def _response_cost_calculator( self, - result: ModelResponse - | ModelResponseStream - | EmbeddingResponse - | ImageResponse - | TranscriptionResponse - | TextCompletionResponse - | HttpxBinaryResponseContent - | RerankResponse - | Batch - | FineTuningJob - | ResponsesAPIResponse - | ResponseCompletedEvent - | OpenAIFileObject - | LiteLLMRealtimeStreamLoggingObject - | OpenAIModerationResponse - | SearchResponse - | dict - | list, + result: Union[ + ModelResponse, + ModelResponseStream, + EmbeddingResponse, + ImageResponse, + TranscriptionResponse, + TextCompletionResponse, + HttpxBinaryResponseContent, + RerankResponse, + Batch, + FineTuningJob, + ResponsesAPIResponse, + ResponseCompletedEvent, + OpenAIFileObject, + LiteLLMRealtimeStreamLoggingObject, + OpenAIModerationResponse, + "SearchResponse", + dict, + list, + ], cache_hit: bool | None = None, litellm_model_name: str | None = None, router_model_id: str | None = None, @@ -6262,7 +6263,7 @@ def _get_traceback_str_for_error(error_str: str) -> str: from decimal import Decimal # used for unit testing -from typing import Any, Optional +from typing import Any, Optional, Union def create_dummy_standard_logging_payload() -> StandardLoggingPayload: diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index 7bdca61fd0f..c471582dc9e 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -53,20 +53,22 @@ class GigaChatModelResponseIterator: finish_reason: str | None = chunk_finish_reason # Handle function_call in stream - if chunk_finish_reason == "function_call" and delta.get("function_call"): - func_call: Final = delta["function_call"] - args_raw: Final = func_call.get("arguments") or {} + raw_function_call: Final = delta.get("function_call") + if chunk_finish_reason == "function_call" and isinstance(raw_function_call, Mapping) and raw_function_call: + func_call: Final[Mapping[str, object]] = raw_function_call + args_raw: Final[object] = func_call.get("arguments") or {} args_str: str # rebind-ok: conditionally assigned from dict or str if isinstance(args_raw, dict): args_str = json.dumps(args_raw, ensure_ascii=False) # rebind-ok: build from dict else: args_str = str(args_raw) + name_raw: Final = func_call.get("name") tool_use = ChatCompletionToolCallChunk( id=f"call_{uuid.uuid4().hex[:24]}", type="function", function=ChatCompletionToolCallFunctionChunk( - name=func_call.get("name", ""), + name=name_raw if isinstance(name_raw, str) else "", arguments=args_str, ), index=0, diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index 643adf43ef8..8f23c5175ec 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -167,19 +167,21 @@ class GigaChatConfig(BaseConfig): pass elif param == "tools": # Convert tools to functions format - optional_params["functions"] = self._convert_tools_to_functions(value) + if isinstance(value, Sequence): + optional_params["functions"] = self._convert_tools_to_functions(value) elif param == "tool_choice": # Map OpenAI tool_choice to GigaChat function_call - mapped_choice = self._map_tool_choice(value) - if mapped_choice is not None: - optional_params["function_call"] = mapped_choice + if isinstance(value, (str, Mapping)): + mapped_choice = self._map_tool_choice(value) + if mapped_choice is not None: + optional_params["function_call"] = mapped_choice elif param == "functions": optional_params["functions"] = value elif param == "function_call": optional_params["function_call"] = value elif param == "response_format": # Handle structured output via function calling - if value.get("type") == "json_schema": + if isinstance(value, Mapping) and value.get("type") == "json_schema": json_schema = value.get("json_schema", {}) schema_name = json_schema.get("name", "structured_output") schema = json_schema.get("schema", {}) @@ -190,9 +192,15 @@ class GigaChatConfig(BaseConfig): "parameters": schema, } - if "functions" not in optional_params: - optional_params["functions"] = [] # mutable-ok: list for httpx - optional_params["functions"].append(function_def) + existing_functions = optional_params.get("functions") + optional_params["functions"] = [ + *( + existing_functions + if isinstance(existing_functions, Sequence) and not isinstance(existing_functions, str) + else () + ), + function_def, + ] optional_params["function_call"] = {"name": schema_name} # mutable-ok: request payload optional_params["_structured_output"] = True @@ -246,8 +254,9 @@ class GigaChatConfig(BaseConfig): # OpenAI format: {"type": "function", "function": {"name": "func_name"}} # GigaChat format: {"name": "func_name"} if tool_choice.get("type") == "function": - func_name: Final = tool_choice.get("function", {}).get("name") - if func_name: + function_spec: Final = tool_choice.get("function") + func_name: Final = function_spec.get("name") if isinstance(function_spec, Mapping) else None + if isinstance(func_name, str) and func_name: return {"name": func_name} # Default to None (don't set function_call) @@ -317,7 +326,7 @@ class GigaChatConfig(BaseConfig): giga_messages: Final = self._transform_messages(messages) # Build request - request_data: Final = { + request_data: Final[dict[str, object]] = { "model": model.replace("gigachat/", ""), "messages": giga_messages, } @@ -407,7 +416,7 @@ class GigaChatConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: tiktoken.Encoding | None, api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index fe65cf5561f..a0edc6f5682 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -92,16 +92,19 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): if provider_chat_config is None: raise ValueError(f"No provider config found for model: {model}") + raw_messages: Final = request_data.get("messages") litellm_model_response: Final = provider_chat_config.transform_response( model=model, - messages=request_data.get("messages", []), # mutable-ok: empty list default for transform_response + messages=list(raw_messages) + if isinstance(raw_messages, list) + else [], # mutable-ok: transform_response wants a list raw_response=httpx_response, model_response=ModelResponse(), logging_obj=logging_obj, optional_params={}, # mutable-ok: empty dict kwarg for transform_response litellm_params={}, # mutable-ok: empty dict kwarg for transform_response api_key="", - request_data=request_data, + request_data=dict(request_data), # mutable-ok: transform_response wants a dict encoding=encoding, ) @@ -124,7 +127,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): logging_obj=logging_obj, optional_params={}, # mutable-ok: empty dict kwarg for transform_embedding_response api_key="", - request_data=request_data, + request_data=dict(request_data), # mutable-ok: transform_embedding_response wants a dict litellm_params={}, # mutable-ok: empty dict kwarg for transform_embedding_response ) ) @@ -172,7 +175,9 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): ) translated_chunk = gigachat_iterator.chunk_parser(chunk=message) - if isinstance(translated_chunk, dict) and generic_chunk_has_all_required_fields(translated_chunk): # pyright: ignore[reportUnnecessaryIsInstance] # runtime guard for patched chunk_parser + if isinstance(translated_chunk, dict) and generic_chunk_has_all_required_fields( # pyright: ignore[reportUnnecessaryIsInstance] # runtime guard for patched chunk_parser + dict(translated_chunk) + ): chunk_obj = convert_generic_chunk_to_model_response_stream( translated_chunk # pyright: ignore[reportArgumentType] # validated TypedDict ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 38867519d80..8ffb259473c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2457,12 +2457,15 @@ class ProxyBaseLLMRequestProcessing: logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete if route_type == "allm_passthrough_route": - streaming_headers = custom_headers # rebind-ok: initial assignment before header merge - if hasattr(response, "headers"): - streaming_headers = ProxyBaseLLMRequestProcessing._merge_passthrough_streaming_headers( # rebind-ok: merge result replaces initial assignment - response_headers=getattr(response, "headers", None), + upstream_response_headers: Final = getattr(response, "headers", None) + streaming_headers: Final = ( + ProxyBaseLLMRequestProcessing._merge_passthrough_streaming_headers( + response_headers=upstream_response_headers, custom_headers=custom_headers, ) + if upstream_response_headers is not None + else custom_headers + ) # Check if response is an async generator if self._is_streaming_response(response): diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 42d45bce3fa..0ddcf99a938 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2824,7 +2824,6 @@ async def gigachat_proxy_route( """ [Docs](https://docs.litellm.ai/docs/pass_through/gigachat) """ - from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.proxy_server import ( general_settings, llm_router, @@ -2876,55 +2875,35 @@ async def gigachat_proxy_route( version=version, ) - # Fall back to existing implementation for direct GigaChat models verbose_proxy_logger.debug( "Gigachat passthrough: Using direct Gigachat model '%s' for endpoint '%s'", model, endpoint ) - data: dict[ - str, Any - ] = {} # mutable-ok: request body mutated in place by proxy pipeline; pyright: ignore[reportExplicitAny] # Any needed for proxy pipeline flexibility + from litellm.llms.gigachat.authenticator import get_access_token + from litellm.llms.gigachat.utils import GIGACHAT_BASE_URL - data["method"] = request.method - data["endpoint"] = endpoint - data["json"] = request_body - data["custom_llm_provider"] = "gigachat" + base_target_url: Final = get_secret_str("GIGACHAT_API_BASE") or GIGACHAT_BASE_URL + request_path: Final = httpx.URL(endpoint).path + encoded_endpoint: Final = request_path if request_path.startswith("/") else f"/{request_path}" - client: Final = get_async_httpx_client( - llm_provider=LlmProviders.GIGACHAT, - params={ # mutable-ok: httpx client params - "timeout": httpx.Timeout(timeout=600.0, connect=5.0), - "ssl_verify": False, - }, + base_url: Final = httpx.URL(base_target_url) + updated_url: Final = base_url.copy_with( + path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint) ) - data["client"] = client - base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) + is_streaming_request: Final = await is_streaming_request_fn(request) - try: - return await base_llm_response_processor.base_passthrough_process_llm_request( - request=request, - fastapi_response=fastapi_response, - user_api_key_dict=user_api_key_dict, - proxy_logging_obj=proxy_logging_obj, - llm_router=llm_router, - general_settings=general_settings, - proxy_config=proxy_config, - select_data_generator=select_data_generator, - model=model, - user_model=user_model, - user_temperature=user_temperature, - user_request_timeout=user_request_timeout, - user_max_tokens=user_max_tokens, - user_api_base=user_api_base, - version=version, - ) - except Exception as e: # noqa: BLE001 # Safe catch-all for handle exception - raise await base_llm_response_processor._handle_llm_api_exception( - e=e, - user_api_key_dict=user_api_key_dict, - proxy_logging_obj=proxy_logging_obj, - ) + endpoint_func: Final = create_pass_through_route( + endpoint=endpoint, + target=str(updated_url), + custom_headers={"Authorization": f"Bearer {get_access_token()}"}, + is_streaming_request=is_streaming_request, + ) + return await endpoint_func( + request, + fastapi_response, + user_api_key_dict, + ) async def handle_gigachat_passthrough_router_model( diff --git a/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py b/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py index 2a44a8e067e..8537793ea72 100644 --- a/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py +++ b/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py @@ -92,7 +92,8 @@ class TestGetOpenaiCompatibleProviderInfo: assert api_base == "https://api.example.com" assert api_key == "test-key" - def test_resolves_api_base_when_none(self): + def test_resolves_api_base_when_none(self, monkeypatch): + monkeypatch.delenv("GIGACHAT_API_BASE", raising=False) provider, api_base, api_key = self.config._get_openai_compatible_provider_info( api_base=None, api_key="key" ) 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 47fafa4c5e0..b969917c8ab 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 @@ -2133,37 +2133,101 @@ class TestGigachatProxyRoute: return_value=False, ) @patch( # test-quality-ok: patching litellm internal for unit test isolation - "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing.base_passthrough_process_llm_request", + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn", new_callable=AsyncMock, + return_value=False, ) - async def test_gigachat_proxy_route_fallback_to_http_pass_through( + @patch( # test-quality-ok: patching litellm internal for unit test isolation + "litellm.llms.gigachat.authenticator.get_access_token", + return_value="gigachat-test-token", + ) + async def test_gigachat_proxy_route_fallback_forwards_to_gigachat_api( self, - mock_base_passthrough, + mock_get_token, + mock_is_streaming, mock_is_router, mock_get_body, + monkeypatch, ): + monkeypatch.delenv("GIGACHAT_API_BASE", raising=False) mock_request = MagicMock(spec=Request) mock_fastapi_response = MagicMock(spec=Response) mock_user_api_key_dict = MagicMock() - expected_response = Response( - content=b'{"response": "success"}', - status_code=200, - media_type="application/json", - ) - mock_base_passthrough.return_value = expected_response + captured_kwargs = {} - result = await gigachat_proxy_route( - endpoint="/chat/completions", - request=mock_request, - fastapi_response=mock_fastapi_response, - user_api_key_dict=mock_user_api_key_dict, - ) + async def fake_endpoint(request, fastapi_response, user_api_key_dict): + return Response(content=b'{"response": "success"}', status_code=200) + + def fake_create_pass_through_route(**kwargs): + captured_kwargs.update(kwargs) + return fake_endpoint + + with patch( # test-quality-ok: patching litellm internal for unit test isolation + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route", + side_effect=fake_create_pass_through_route, + ): + result = await gigachat_proxy_route( + endpoint="/chat/completions", + request=mock_request, + fastapi_response=mock_fastapi_response, + user_api_key_dict=mock_user_api_key_dict, + ) assert isinstance(result, Response) assert result.status_code == 200 - assert result.body == b'{"response": "success"}' - mock_base_passthrough.assert_awaited_once() + assert captured_kwargs["target"] == "https://gigachat.devices.sberbank.ru/api/v1/chat/completions" + assert captured_kwargs["custom_headers"] == {"Authorization": "Bearer gigachat-test-token"} + + @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", + return_value={}, + ) + @patch( # test-quality-ok: patching litellm internal for unit test isolation + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn", + new_callable=AsyncMock, + return_value=False, + ) + @patch( # test-quality-ok: patching litellm internal for unit test isolation + "litellm.llms.gigachat.authenticator.get_access_token", + return_value="gigachat-test-token", + ) + async def test_gigachat_proxy_route_models_endpoint_without_model( + self, + mock_get_token, + mock_is_streaming, + mock_get_body, + monkeypatch, + ): + monkeypatch.delenv("GIGACHAT_API_BASE", raising=False) + mock_request = MagicMock(spec=Request) + mock_fastapi_response = MagicMock(spec=Response) + mock_user_api_key_dict = MagicMock() + + captured_kwargs = {} + + async def fake_endpoint(request, fastapi_response, user_api_key_dict): + return Response(content=b'{"data": []}', status_code=200) + + def fake_create_pass_through_route(**kwargs): + captured_kwargs.update(kwargs) + return fake_endpoint + + with patch( # test-quality-ok: patching litellm internal for unit test isolation + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route", + side_effect=fake_create_pass_through_route, + ): + result = await gigachat_proxy_route( + endpoint="models", + request=mock_request, + fastapi_response=mock_fastapi_response, + user_api_key_dict=mock_user_api_key_dict, + ) + + assert isinstance(result, Response) + assert result.status_code == 200 + assert captured_kwargs["target"] == "https://gigachat.devices.sberbank.ru/api/v1/models" @pytest.mark.asyncio async def test_allm_passthrough_streaming_preserves_upstream_headers(self): From a928c1429e5f91b2aa1d95e3e2e55ebe9f223305 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 22:12:35 -0700 Subject: [PATCH 070/120] fix(proxy): preserve model table columns on master key rotation --- .../key_management_endpoints.py | 53 ++++++++++------- .../test_key_management_endpoints.py | 58 +++++++++---------- 2 files changed, 61 insertions(+), 50 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 0d12b012c18..f1c31c3ff35 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -25,6 +25,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeV import fastapi import yaml from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -154,6 +155,7 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: + import prisma from prisma import Prisma from prisma import models as prisma_models @@ -181,6 +183,14 @@ class _TxTables(Protocol): litellm_proxymodeltable: TableActions[object] +class _ModelParamsUpdate(TypedDict): + litellm_params: ReadOnly["prisma.Json"] + + +class _ModelRowWhere(TypedDict): + model_id: ReadOnly[str] + + class _ConfigTableActions(Protocol): """Config table surface this module needs; the shared repository seam exposes no ``update``.""" @@ -4427,28 +4437,29 @@ async def _rotate_master_key( if models: decrypted_models: Final = proxy_config.decrypt_model_list_from_db(new_models=models) verbose_proxy_logger.debug("ABLE TO DECRYPT MODELS - len(decrypted_models): %s", len(decrypted_models)) - new_models: Final[list[dict[str, object]]] = [] - for model in decrypted_models: - new_model = await _add_model_to_db( - model_params=Deployment(**model), - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - new_encryption_key=new_master_key, - should_create_model_in_db=False, - ) - if new_model: - _dumped = new_model.model_dump(exclude_none=True) - _dumped["litellm_params"] = prisma.Json(_dumped["litellm_params"]) - _dumped["model_info"] = prisma.Json(_dumped["model_info"]) - new_models.append(_dumped) - verbose_proxy_logger.debug("Resetting proxy model table") - async with prisma_client.db.tx() as tx_ctx: + reencrypted_models: Final = tuple( + [ + reencrypted + for model in decrypted_models + if ( + reencrypted := await _add_model_to_db( + model_params=Deployment(**model), + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + new_encryption_key=new_master_key, + should_create_model_in_db=False, + ) + ) + ] + ) + verbose_proxy_logger.debug("Re-encrypting litellm_params on %s model rows", len(reencrypted_models)) + async with prisma_client.db.tx(timeout=timedelta(minutes=2)) as tx_ctx: tx: Final[_TxTables] = tx_ctx - await tx.litellm_proxymodeltable.delete_many() - verbose_proxy_logger.debug("Creating %s models", len(new_models)) - await tx.litellm_proxymodeltable.create_many( - data=new_models, - ) + for reencrypted_model in reencrypted_models: + await tx.litellm_proxymodeltable.update_many( + data=_ModelParamsUpdate(litellm_params=prisma.Json(reencrypted_model.litellm_params)), + where=_ModelRowWhere(model_id=reencrypted_model.model_id), + ) await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") # 3. process config table try: diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 42e56ceabd3..1c80aa5683f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -8312,15 +8312,17 @@ async def test_key_does_not_override_explicit_budget_duration(): @patch( "litellm.proxy.management_endpoints.key_management_endpoints.rotate_mcp_server_credentials_master_key" ) -async def test_rotate_master_key_model_data_valid_for_prisma( +async def test_rotate_master_key_reencrypts_model_params_in_place( mock_rotate_mcp, ): """ - Test that _rotate_master_key produces valid data for Prisma create_many(). - - Regression test for: master key rotation fails with Prisma validation error - because created_at/updated_at are None (non-nullable DateTime) and - litellm_params/model_info are JSON strings (create_many expects dicts). + Regression test for: master key rotation wipes every non-credential column + on LiteLLM_ProxyModelTable. Rotation used to rebuild the table via + delete_many + create_many from Deployment objects, which carry no + blocked/created_at/created_by/updated_at/updated_by, so every rotation + reset blocked to False (silently unblocking blocked models) and rewrote the + audit columns. Rotation must instead update only litellm_params (the sole + encrypted column) on each existing row, keyed by model_id. """ from unittest.mock import AsyncMock, MagicMock @@ -8352,6 +8354,7 @@ async def test_rotate_master_key_model_data_valid_for_prisma( mock_tx.litellm_proxymodeltable = MagicMock() mock_tx.litellm_proxymodeltable.delete_many = AsyncMock() mock_tx.litellm_proxymodeltable.create_many = AsyncMock() + mock_tx.litellm_proxymodeltable.update_many = AsyncMock() mock_prisma_client.db.tx = MagicMock( return_value=AsyncMock( __aenter__=AsyncMock(return_value=mock_tx), @@ -8400,36 +8403,33 @@ async def test_rotate_master_key_model_data_valid_for_prisma( new_master_key="sk-new-master-key", ) - # Verify create_many was called - mock_tx.litellm_proxymodeltable.create_many.assert_called_once() + # Rotation must never rewrite whole rows: no delete + recreate + mock_tx.litellm_proxymodeltable.delete_many.assert_not_called() + mock_tx.litellm_proxymodeltable.create_many.assert_not_called() - # Get the data passed to create_many - call_args = mock_tx.litellm_proxymodeltable.create_many.call_args - created_models = call_args.kwargs.get("data") or call_args[1].get("data") + mock_tx.litellm_proxymodeltable.update_many.assert_called_once() + call_args = mock_tx.litellm_proxymodeltable.update_many.call_args - assert len(created_models) == 1 - model_data = created_models[0] + assert call_args.kwargs["where"] == { + "model_id": "model-1" + }, "the re-encrypted params must land on the same row, keyed by model_id" - # Verify timestamps are NOT present (Prisma @default(now()) should apply) - assert ( - "created_at" not in model_data - ), "created_at should be excluded so Prisma @default(now()) applies" - assert ( - "updated_at" not in model_data - ), "updated_at should be excluded so Prisma @default(now()) applies" + update_data = call_args.kwargs["data"] + assert set(update_data.keys()) == {"litellm_params"}, ( + "rotation must touch only the encrypted litellm_params column; writing any " + f"other column wipes it (blocked, audit columns), got {sorted(update_data.keys())}" + ) - # Verify litellm_params and model_info are prisma.Json wrappers, NOT JSON strings import prisma assert isinstance( - model_data["litellm_params"], prisma.Json - ), f"litellm_params should be prisma.Json for create_many(), got {type(model_data['litellm_params'])}" - assert isinstance( - model_data["model_info"], prisma.Json - ), f"model_info should be prisma.Json for create_many(), got {type(model_data['model_info'])}" - - # Verify delete_many was called inside the transaction (before create_many) - mock_tx.litellm_proxymodeltable.delete_many.assert_called_once() + update_data["litellm_params"], prisma.Json + ), f"litellm_params should be prisma.Json for update_many(), got {type(update_data['litellm_params'])}" + reencrypted_params = update_data["litellm_params"].data + assert set(reencrypted_params.keys()) >= {"model", "api_key"} + assert ( + reencrypted_params["api_key"] != "sk-decrypted-key" + ), "api_key must be stored re-encrypted under the new master key, not in plaintext" async def test_default_key_generate_params_duration(monkeypatch): From 3ea4b715ba8d6c1666205443a12475d74a1af040 Mon Sep 17 00:00:00 2001 From: siyoon Date: Sun, 30 Aug 2026 14:22:41 +0900 Subject: [PATCH 071/120] feat(friendli): add zai-org/GLM-5.3-Flash model pricing Per https://api.friendli.ai/serverless/v1/models: - $0.15 input / $0.50 output / $0.03 cached input per MTok - 1M context, 1M max output, reasoning with effort low/high/max (per HF chat_template.jinja: low/high honored, anything else -> max) - tool calling, parallel tool calls, structured output, prompt caching - image + video input (native multimodal) --- model_prices_and_context_window.json | 25 +++++++++++++ ...t_friendli_glm_5_3_flash_model_metadata.py | 37 +++++++++++++++++++ 2 files changed, 62 insertions(+) create mode 100644 tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 05c1cfd3179..fd6b9e11adf 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19564,6 +19564,31 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "friendliai/zai-org/GLM-5.3-Flash": { + "litellm_provider": "friendliai", + "supports_reasoning": true, + "supports_function_calling": true, + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "max_output_tokens": 1048576, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 3e-08, + "supports_prompt_caching": true, + "supports_max_reasoning_effort": true, + "supports_low_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "mode": "chat", + "comment": "Native multimodal GLM model for efficient coding and long-horizon agent tasks", + "source": "https://api.friendli.ai/serverless/v1/models", + "supports_vision": true, + "supports_image_input": true, + "supports_video_input": true + }, "ft:babbage-002": { "deprecation_date": "2026-10-23", "input_cost_per_token": 1.6e-06, diff --git a/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py b/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py new file mode 100644 index 00000000000..5110307e5bb --- /dev/null +++ b/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py @@ -0,0 +1,37 @@ +import json +from pathlib import Path + +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + +def test_friendli_glm_5_3_flash_model_info(): + model = "friendliai/zai-org/GLM-5.3-Flash" + json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" + with open(json_path) as f: + model_cost = json.load(f) + + info = model_cost.get(model) + assert ( + info is not None + ), f"{model} not found in model_prices_and_context_window.json" + assert info["litellm_provider"] == "friendliai" + assert info["mode"] == "chat" + assert info["input_cost_per_token"] == 1.5e-07 + assert info["output_cost_per_token"] == 5e-07 + assert info["cache_read_input_token_cost"] == 3e-08 + assert info["max_input_tokens"] == 1048576 + assert info["max_output_tokens"] == 1048576 + assert info["supports_function_calling"] is True + assert info["supports_reasoning"] is True + assert info["supports_low_reasoning_effort"] is True + assert info["supports_max_reasoning_effort"] is True + assert info["supports_tool_choice"] is True + assert info["supports_prompt_caching"] is True + # Friendli serves image and video input + assert info["supports_vision"] is True + assert info["supports_image_input"] is True + assert info["supports_video_input"] is True + + routed_model, provider, _, _ = get_llm_provider(model=model) + assert routed_model == "zai-org/GLM-5.3-Flash" + assert provider == "friendliai" From e7bfe99cd3a0a471071e402aec733351e6e76039 Mon Sep 17 00:00:00 2001 From: siyoon Date: Sun, 30 Aug 2026 14:23:36 +0900 Subject: [PATCH 072/120] feat(friendli): add zai-org/GLM-5.3 model pricing Per https://api.friendli.ai/serverless/v1/models: - $1.40 input / $4.40 output / $0.26 cached input per MTok - 1M context, 1M max output, reasoning with effort low/high/max (per HF chat_template.jinja: low/high honored, anything else -> max) - tool calling, parallel tool calls, structured output, prompt caching - text-only (no vision), flagship GLM model --- model_prices_and_context_window.json | 24 +++++++++++++ .../test_friendli_glm_5_3_model_metadata.py | 36 +++++++++++++++++++ 2 files changed, 60 insertions(+) create mode 100644 tests/test_litellm/test_friendli_glm_5_3_model_metadata.py diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 05c1cfd3179..24aad7c90c8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19564,6 +19564,30 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "friendliai/zai-org/GLM-5.3": { + "litellm_provider": "friendliai", + "supports_reasoning": true, + "supports_function_calling": true, + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "max_output_tokens": 1048576, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 2.6e-07, + "supports_prompt_caching": true, + "supports_max_reasoning_effort": true, + "supports_low_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "mode": "chat", + "comment": "Flagship GLM model for long-horizon coding, agents, and complex project delivery", + "source": "https://api.friendli.ai/serverless/v1/models", + "supports_vision": false, + "supports_image_input": false + }, "ft:babbage-002": { "deprecation_date": "2026-10-23", "input_cost_per_token": 1.6e-06, diff --git a/tests/test_litellm/test_friendli_glm_5_3_model_metadata.py b/tests/test_litellm/test_friendli_glm_5_3_model_metadata.py new file mode 100644 index 00000000000..c0b07358d23 --- /dev/null +++ b/tests/test_litellm/test_friendli_glm_5_3_model_metadata.py @@ -0,0 +1,36 @@ +import json +from pathlib import Path + +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + +def test_friendli_glm_5_3_model_info(): + model = "friendliai/zai-org/GLM-5.3" + json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" + with open(json_path) as f: + model_cost = json.load(f) + + info = model_cost.get(model) + assert ( + info is not None + ), f"{model} not found in model_prices_and_context_window.json" + assert info["litellm_provider"] == "friendliai" + assert info["mode"] == "chat" + assert info["input_cost_per_token"] == 1.4e-06 + assert info["output_cost_per_token"] == 4.4e-06 + assert info["cache_read_input_token_cost"] == 2.6e-07 + assert info["max_input_tokens"] == 1048576 + assert info["max_output_tokens"] == 1048576 + assert info["supports_function_calling"] is True + assert info["supports_reasoning"] is True + assert info["supports_low_reasoning_effort"] is True + assert info["supports_max_reasoning_effort"] is True + assert info["supports_tool_choice"] is True + assert info["supports_prompt_caching"] is True + # GLM-5.3 (non-flash) is text-only on Friendli's catalog + assert info["supports_vision"] is False + assert info["supports_image_input"] is False + + routed_model, provider, _, _ = get_llm_provider(model=model) + assert routed_model == "zai-org/GLM-5.3" + assert provider == "friendliai" From e9f1af84730ffd1c11b9f77a10c38d0bf20351be Mon Sep 17 00:00:00 2001 From: siyoon Date: Sun, 30 Aug 2026 14:29:28 +0900 Subject: [PATCH 073/120] chore(tests): drop redundant capability comment Per greptile review + CLAUDE.md comment policy: the comment restated the immediately following assertions without adding value. --- tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py b/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py index 5110307e5bb..0acd750e49a 100644 --- a/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py +++ b/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py @@ -27,7 +27,6 @@ def test_friendli_glm_5_3_flash_model_info(): assert info["supports_max_reasoning_effort"] is True assert info["supports_tool_choice"] is True assert info["supports_prompt_caching"] is True - # Friendli serves image and video input assert info["supports_vision"] is True assert info["supports_image_input"] is True assert info["supports_video_input"] is True From 3822ecc8ff1df4a0c41c2d3deb6282206eb537da Mon Sep 17 00:00:00 2001 From: siyoon Date: Sun, 30 Aug 2026 14:29:42 +0900 Subject: [PATCH 074/120] chore(tests): drop redundant capability comment Per greptile review + CLAUDE.md comment policy: the comment restated the immediately following assertion without adding value. --- tests/test_litellm/test_friendli_glm_5_3_model_metadata.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_litellm/test_friendli_glm_5_3_model_metadata.py b/tests/test_litellm/test_friendli_glm_5_3_model_metadata.py index c0b07358d23..5e7a1ede699 100644 --- a/tests/test_litellm/test_friendli_glm_5_3_model_metadata.py +++ b/tests/test_litellm/test_friendli_glm_5_3_model_metadata.py @@ -27,7 +27,6 @@ def test_friendli_glm_5_3_model_info(): assert info["supports_max_reasoning_effort"] is True assert info["supports_tool_choice"] is True assert info["supports_prompt_caching"] is True - # GLM-5.3 (non-flash) is text-only on Friendli's catalog assert info["supports_vision"] is False assert info["supports_image_input"] is False From 693279afb443d8a1775fb82b83ae8c2ad0c7a295 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 30 Aug 2026 07:57:16 +0000 Subject: [PATCH 075/120] chore(techdebt): clear fresh debt from the 2026-08-29 window Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- basedpyright-code-budget.json | 4 ++-- litellm/litellm_core_utils/model_response_utils.py | 2 -- litellm/llms/oci/chat/cohere.py | 4 ++-- .../llms/vertex_ai/image_edit/vertex_gemini_transformation.py | 1 - .../llms/vertex_ai/vector_stores/rag_api/transformation.py | 2 -- litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py | 4 ++-- .../proxy/guardrails/guardrail_hooks/repelloai/repelloai.py | 2 +- litellm/rag/ingestion/vertex_ai_ingestion.py | 2 -- ruff-strict-budget.json | 4 ++-- type-discipline-budget.json | 2 +- 10 files changed, 10 insertions(+), 17 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index f8ef0ce1732..529a36d2614 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 16389 + "limit": 16387 }, "reportArgumentType": { "limit": 2229 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 5242 + "limit": 5241 }, "reportFunctionMemberAccess": { "limit": 7 diff --git a/litellm/litellm_core_utils/model_response_utils.py b/litellm/litellm_core_utils/model_response_utils.py index 7bf667164ae..dc4f375daa7 100644 --- a/litellm/litellm_core_utils/model_response_utils.py +++ b/litellm/litellm_core_utils/model_response_utils.py @@ -144,7 +144,6 @@ def _is_choice_non_empty(choice: StreamingChoices) -> bool: # Check model_extra for dynamically added fields on the choice choice_extra_fields: Final[Mapping[str, object]] = choice.model_extra or {} for extra_field_name, extra_field_value in choice_extra_fields.items(): - # Skip certain structural fields that are just default/None placeholders if extra_field_name == "index" and extra_field_value == 0: continue if extra_field_name in {"finish_reason", "logprobs"} and extra_field_value is None: @@ -192,7 +191,6 @@ def _is_delta_non_empty(delta: Delta) -> bool: # Check model_extra for dynamically added fields (this is where Pydantic stores them) delta_extra_fields: Final[Mapping[str, object]] = delta.model_extra or {} for extra_field_value in delta_extra_fields.values(): - # Even structural fields are meaningful if they have actual content if _has_meaningful_content(extra_field_value): return True diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index 384e7ec4cf8..6e9bb83b0a0 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -9,7 +9,7 @@ response parsing, and streaming chunk parsing for models served with import datetime import json from collections.abc import Iterable, Mapping, Sequence -from typing import Any, Final +from typing import Final import httpx from pydantic import JsonValue, TypeAdapter, ValidationError @@ -76,7 +76,7 @@ def _content_text(content: str | Iterable[Mapping[str, object]] | None) -> str: return str(content) -def _extract_text_content(content: Any) -> str: +def _extract_text_content(content: str | Iterable[Mapping[str, object]] | None) -> str: """Return the plain-text representation of a message content value.""" return _content_text(content) diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index 5889a8eba06..725a7f39917 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -182,7 +182,6 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): else None ) - # Generation config with proper structure for image editing generation_config: Final[dict[str, object]] = { key: value for key, value in (("response_modalities", ["IMAGE"]), ("image_config", image_config)) if value } diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index eedd488ecdb..5c250fc1a7e 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -203,7 +203,6 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): if value is not None } - # Build the request body for Vertex AI RAG API query_body: Final[Mapping[str, object]] = { key: value for key, value in (("text", query), ("rag_retrieval_config", rag_retrieval_config or None)) @@ -294,7 +293,6 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): # Add metadata if provided metadata: Final = vector_store_create_optional_params.get("metadata") - # Build the request body for Vertex AI RAG Corpus creation request_body: Final[dict[str, object]] = { key: value for key, value in ( diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index daa15fb5e5f..d8c8c2f4974 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -433,7 +433,7 @@ class HeadroomGuardrail(CustomGuardrail): payload["model"] = model try: - raw_response: HttpxResponse = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] + raw_response: HttpxResponse = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped url=f"{self.headroom_api_base}/v1/compress", json=payload, headers=self._request_headers(), @@ -570,7 +570,7 @@ class HeadroomGuardrail(CustomGuardrail): params["query"] = query try: - raw_response: HttpxResponse = await self.async_handler.get( # pyright: ignore[reportUnknownMemberType] + raw_response: HttpxResponse = await self.async_handler.get( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.get is untyped url=f"{self.headroom_api_base}/v1/retrieve/{hash_value}", params=params, headers=self._request_headers(), diff --git a/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py b/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py index d842a9b9b6a..8925cc5b3a6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py @@ -197,7 +197,7 @@ class RepelloAIGuardrail(CustomGuardrail): repelloai_response: RepelloAIAnalyzeResponse | None = None try: verbose_proxy_logger.debug("RepelloAI Argus request: %s", request) - response: Final[HttpxResponse] = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] + response: Final[HttpxResponse] = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped url=endpoint, headers={"X-API-Key": self.repelloai_api_key}, json=request, diff --git a/litellm/rag/ingestion/vertex_ai_ingestion.py b/litellm/rag/ingestion/vertex_ai_ingestion.py index 07f9f346d08..eff8ad1b8cb 100644 --- a/litellm/rag/ingestion/vertex_ai_ingestion.py +++ b/litellm/rag/ingestion/vertex_ai_ingestion.py @@ -186,7 +186,6 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): base_url: Final = get_vertex_base_url(self.location) url: Final = f"{base_url}/v1beta1/projects/{self.project_id}/locations/{self.location}/ragCorpora" - # Build request body with camelCase keys (Vertex AI API format) vector_db_config: Final = self.vector_store_config.get("vector_db_config") embedding_model: Final = self.vector_store_config.get("embedding_model") embedding_model_config: Final = ( @@ -447,7 +446,6 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): # Add max embedding requests per minute if specified max_embedding_qpm: Final = self.vector_store_config.get("max_embedding_requests_per_min") - # Build request body with camelCase keys (Vertex AI API format) chunking_config: Final = ( {"chunkSize": chunk_size or 1024, "chunkOverlap": chunk_overlap or 200} if chunk_size or chunk_overlap diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index fce91d57d88..5d0e7b617cb 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 597 + "limit": 596 }, "ASYNC230": { "limit": 11 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1111 + "limit": 1110 }, "TRY002": { "limit": 524 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index bb02618b11e..9c805bd7b61 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -9,7 +9,7 @@ "limit": 269 }, "LIT004": { - "limit": 43 + "limit": 40 }, "LIT005": { "limit": 0 From 739f61df7dff60e6f37d847bcd6015a9346e29de Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 30 Aug 2026 12:26:06 -0700 Subject: [PATCH 076/120] test(model_management): drop docstring that restates the serialization path --- .../management_endpoints/test_model_management_endpoints.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 13c8a26377c..393953ccf68 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -4469,12 +4469,6 @@ class TestEnforceRpmTpmOnModelAdd: class TestBlockModelResponseSerialization: - """POST /model/block and /model/unblock return the raw prisma row through this - route's `LiteLLM_ProxyModelTable | None` response validation. The row is not a - dict, so the dict-assuming before-validator used to raise AttributeError inside - FastAPI's serialization layer: a 500 for the caller after the DB write already - landed. The routes must serialize the row to a 200 with the updated blocked flag.""" - @pytest.mark.parametrize( ("route", "blocked"), [("/model/block", True), ("/model/unblock", False)] ) From ce52e39052ea605377e2bea1848f8a18e3f36018 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 30 Aug 2026 12:28:41 -0700 Subject: [PATCH 077/120] fix(gigachat): honor ssl_verify config on router passthrough and type the request body --- .../pass_through_endpoints/llm_passthrough_endpoints.py | 6 +++--- .../test_llm_pass_through_endpoints.py | 1 + 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 0ddcf99a938..6c7d0c92e58 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2839,10 +2839,11 @@ async def gigachat_proxy_route( ) ## check for streaming - request_body: Final = await get_request_body(request) # pyright: ignore[reportUnknownVariableType] # get_request_body returns Unknown + request_body: Final[dict[str, object]] = await get_request_body(request) is_router_model = False # rebind-ok: conditionally set to True when model uses router - model: Final = request_body.get("model") # pyright: ignore[reportUnknownVariableType] # get_request_body returns Unknown + raw_model: Final = request_body.get("model") + model: Final = raw_model if isinstance(raw_model, str) else None if model: is_router_model = is_passthrough_request_using_router_model( request_body, llm_router @@ -2995,7 +2996,6 @@ async def handle_gigachat_passthrough_router_model( llm_provider=LlmProviders.GIGACHAT, params={ # mutable-ok: httpx client params "timeout": httpx.Timeout(timeout=600.0, connect=5.0), - "ssl_verify": False, }, ) 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 b969917c8ab..225e1b3d998 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 @@ -2228,6 +2228,7 @@ class TestGigachatProxyRoute: assert isinstance(result, Response) assert result.status_code == 200 assert captured_kwargs["target"] == "https://gigachat.devices.sberbank.ru/api/v1/models" + assert captured_kwargs["custom_headers"] == {"Authorization": "Bearer gigachat-test-token"} @pytest.mark.asyncio async def test_allm_passthrough_streaming_preserves_upstream_headers(self): From a5fa8ebfa73e2587ac7357bb6dafc93d5f38910a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 30 Aug 2026 12:59:18 -0700 Subject: [PATCH 078/120] fix(passthrough): keep upstream error body readable for streaming error status mapping --- ...odel_prices_and_context_window_backup.json | 11 +++++++- litellm/passthrough/main.py | 4 +++ .../test_async_streaming_error_propagation.py | 27 +++++++++++++++++++ 3 files changed, 41 insertions(+), 1 deletion(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 05c1cfd3179..1865d0b7f3a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -24342,7 +24342,7 @@ "supports_response_schema": true, "supports_vision": true }, - "gigachat/GigaChat-2-Lite": { + "gigachat/GigaChat-2": { "input_cost_per_token": 0.0, "litellm_provider": "gigachat", "max_input_tokens": 128000, @@ -24404,6 +24404,15 @@ "output_cost_per_token": 0.0, "output_vector_size": 2560 }, + "gigachat/GigaEmbeddings-3B-2025-09": { + "input_cost_per_token": 0.0, + "litellm_provider": "gigachat", + "max_input_tokens": 4096, + "max_tokens": 4096, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 2048 + }, "gmi/anthropic/claude-opus-4.5": { "input_cost_per_token": 5e-06, "litellm_provider": "gmi", diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 1715ec10f30..2780f510a76 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -86,6 +86,10 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): self._response.raise_for_status() self._iterator = _as_async_generator(self._response.aiter_bytes()) except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic + try: + await self._response.aread() + except Exception: # noqa: BLE001 S110 # Safe catch-all for cleanup logic + pass try: await self._response.aclose() except Exception: # noqa: BLE001 S110 # Safe catch-all for cleanup logic diff --git a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py b/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py index 7bf5dc874a8..9f2b436d2d8 100644 --- a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py +++ b/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py @@ -131,3 +131,30 @@ async def test_async_passthrough_wrapper_200_yields_chunks(): assert len(chunks) == 1 assert b"response.created" in chunks[0] mock_logging_obj.async_flush_passthrough_collected_chunks.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_error_body_readable_after_failed_await(): + """The upstream error body must stay readable so the proxy can map the real status and message.""" + from litellm.passthrough.main import AsyncPassthroughStreamingResponse + + error_body = b'{"message":"model not found"}' + + async def byte_stream(): + yield error_body + + request = httpx.Request("POST", "https://bedrock.example.com/model/x/converse-stream") + response = httpx.Response(400, content=byte_stream(), request=request) + + async def response_coro(): + return response + + with pytest.raises(httpx.HTTPStatusError) as exc_info: + await AsyncPassthroughStreamingResponse( + response=response_coro(), + litellm_logging_obj=_make_mock_logging_obj(), + provider_config=MagicMock(), + ) + + assert exc_info.value.response.status_code == 400 + assert await exc_info.value.response.aread() == error_body From de1f38820a2ff583d223027a468e74f4431e42b5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 30 Aug 2026 13:36:51 -0700 Subject: [PATCH 079/120] fix(passthrough): flush interrupted streams on client disconnect and reuse cached gigachat http clients --- litellm/llms/gigachat/authenticator.py | 3 +- litellm/llms/gigachat/file_handler.py | 6 +-- litellm/proxy/common_request_processing.py | 2 +- .../llms/gigachat/test_authenticator.py | 7 ++++ .../llms/gigachat/test_file_handler.py | 20 +++++----- .../proxy/test_common_request_processing.py | 40 ++++++++++++++++++- 6 files changed, 61 insertions(+), 17 deletions(-) diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index 293176d8931..d6b217d5746 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -17,6 +17,7 @@ from litellm.caching.caching import InMemoryCache from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, + _get_httpx_client, # pyright: ignore[reportPrivateUsage] # house cached-client factory has no public alias get_async_httpx_client, ) from litellm.secret_managers.main import get_secret_str @@ -56,7 +57,7 @@ def _get_scope() -> str: def _get_http_client() -> HTTPHandler: """Get cached httpx client with SSL verification disabled.""" - return HTTPHandler(ssl_verify=False) + return _get_httpx_client(params={"ssl_verify": False}) def get_access_token( diff --git a/litellm/llms/gigachat/file_handler.py b/litellm/llms/gigachat/file_handler.py index 900aa72f34c..163e944f124 100644 --- a/litellm/llms/gigachat/file_handler.py +++ b/litellm/llms/gigachat/file_handler.py @@ -14,7 +14,7 @@ from typing import Final from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( - HTTPHandler, + _get_httpx_client, get_async_httpx_client, ) from litellm.llms.gigachat.utils import get_api_base @@ -52,7 +52,7 @@ def _parse_data_url(data_url: str) -> tuple[bytes, str, str] | None: def _download_image_sync(url: str) -> tuple[bytes, str, str]: """Download image from URL synchronously.""" - client: Final = HTTPHandler(ssl_verify=False) + client: Final = _get_httpx_client(params={"ssl_verify": False}) response: Final = client.get(url) response.raise_for_status() @@ -120,7 +120,7 @@ def upload_file_sync( base_url: Final = get_api_base(api_base) upload_url: Final = f"{base_url}/files" - client: Final = HTTPHandler(ssl_verify=False) + client: Final = _get_httpx_client(params={"ssl_verify": False}) response: Final = client.post( upload_url, headers={"Authorization": f"Bearer {access_token}"}, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 8ffb259473c..67306ea7d46 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2496,7 +2496,7 @@ class ProxyBaseLLMRequestProcessing: # For passthrough routes, stream directly without error parsing # since we're dealing with raw binary data (e.g., AWS event streams) - return StreamingResponse( + return _UpstreamClosingStreamingResponse( content=generator, # pyright: ignore[reportArgumentType] # generator-configured StreamingResponse status_code=getattr(response, "status_code", status.HTTP_200_OK), media_type=self._passthrough_event_stream_media_type(), diff --git a/tests/test_litellm/llms/gigachat/test_authenticator.py b/tests/test_litellm/llms/gigachat/test_authenticator.py index 30b3dec52cb..0a2695dc21e 100644 --- a/tests/test_litellm/llms/gigachat/test_authenticator.py +++ b/tests/test_litellm/llms/gigachat/test_authenticator.py @@ -485,3 +485,10 @@ class TestParseTokenResponse: _parse_token_response(self._make_response({"exp": 1700000000000})) assert exc_info.value.status_code == 500 assert "Invalid token response" in exc_info.value.message + + +class TestGetHttpClient: + def test_reuses_cached_client_across_calls(self): + """Regression: the sync OAuth path must use the shared cached httpx client, + not construct a fresh HTTPHandler per token request.""" + assert authenticator._get_http_client() is authenticator._get_http_client() diff --git a/tests/test_litellm/llms/gigachat/test_file_handler.py b/tests/test_litellm/llms/gigachat/test_file_handler.py index ced7c30ed9d..ce9505f11f2 100644 --- a/tests/test_litellm/llms/gigachat/test_file_handler.py +++ b/tests/test_litellm/llms/gigachat/test_file_handler.py @@ -128,7 +128,7 @@ class TestParseDataUrl: class TestDownloadImageSync: - @patch(f"{FILE_MODULE}.HTTPHandler") + @patch(f"{FILE_MODULE}._get_httpx_client") def test_downloads_image_successfully(self, mock_http_handler_cls): mock_client = MagicMock() mock_response = MagicMock() @@ -144,7 +144,7 @@ class TestDownloadImageSync: assert ext == "jpeg" mock_client.get.assert_called_once_with("https://example.com/img.jpg") - @patch(f"{FILE_MODULE}.HTTPHandler") + @patch(f"{FILE_MODULE}._get_httpx_client") def test_raises_on_http_error(self, mock_http_handler_cls): mock_client = MagicMock() mock_client.get.side_effect = httpx.HTTPStatusError( @@ -157,7 +157,7 @@ class TestDownloadImageSync: with pytest.raises(httpx.HTTPStatusError): file_handler._download_image_sync("https://example.com/404") - @patch(f"{FILE_MODULE}.HTTPHandler") + @patch(f"{FILE_MODULE}._get_httpx_client") def test_parse_content_type_fallback(self, mock_http_handler_cls): mock_client = MagicMock() mock_response = MagicMock() @@ -171,7 +171,7 @@ class TestDownloadImageSync: assert content_type == "image/jpeg" assert ext == "jpeg" - @patch(f"{FILE_MODULE}.HTTPHandler") + @patch(f"{FILE_MODULE}._get_httpx_client") def test_extracts_extension_from_parametrized_type(self, mock_http_handler_cls): mock_client = MagicMock() mock_response = MagicMock() @@ -235,7 +235,7 @@ class TestDownloadImageAsync: class TestUploadFileSync: @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") - @patch(f"{FILE_MODULE}.HTTPHandler") + @patch(f"{FILE_MODULE}._get_httpx_client") def test_uploads_base64_image_and_caches( self, mock_http_handler_cls, mock_get_token, mock_get_api_base ): @@ -268,7 +268,7 @@ class TestUploadFileSync: @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") - @patch(f"{FILE_MODULE}.HTTPHandler") + @patch(f"{FILE_MODULE}._get_httpx_client") def test_returns_cached_file_id( self, mock_http_handler_cls, mock_get_token, mock_get_api_base ): @@ -282,7 +282,7 @@ class TestUploadFileSync: # No upload call was made mock_http_handler_cls.return_value.post.assert_not_called() - @patch(f"{FILE_MODULE}.HTTPHandler") + @patch(f"{FILE_MODULE}._get_httpx_client") @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") @patch(f"{FILE_MODULE}._download_image_sync") @@ -304,7 +304,7 @@ class TestUploadFileSync: assert result == "file-remote" mock_download.assert_called_once_with("https://example.com/remote.png") - @patch(f"{FILE_MODULE}.HTTPHandler") + @patch(f"{FILE_MODULE}._get_httpx_client") @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") def test_returns_none_on_upload_failure( @@ -325,7 +325,7 @@ class TestUploadFileSync: assert result is None - @patch(f"{FILE_MODULE}.HTTPHandler") + @patch(f"{FILE_MODULE}._get_httpx_client") @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") def test_returns_none_when_response_missing_id( @@ -346,7 +346,7 @@ class TestUploadFileSync: @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") - @patch(f"{FILE_MODULE}.HTTPHandler") + @patch(f"{FILE_MODULE}._get_httpx_client") def test_uploads_without_optional_args( self, mock_http_handler_cls, mock_get_token, mock_get_api_base ): diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 833c3754022..45e079c8595 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -4694,7 +4694,7 @@ class TestAllmPassthroughStreamingProviderGate: } return ProxyBaseLLMRequestProcessing(data=data) - async def _run(self, processing_obj, monkeypatch, chunks): + async def _run(self, processing_obj, monkeypatch, chunks, stream=None): import litellm.proxy.common_request_processing as crp from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth @@ -4702,9 +4702,11 @@ class TestAllmPassthroughStreamingProviderGate: for chunk in chunks: yield chunk + upstream_stream = stream if stream is not None else streaming_response() + async def fake_route_request(**kwargs): async def _llm_call(): - return streaming_response() + return upstream_stream return _llm_call() @@ -4729,6 +4731,40 @@ class TestAllmPassthroughStreamingProviderGate: skip_pre_call_logic=True, ) + @pytest.mark.asyncio + async def test_client_disconnect_closes_unbuffered_passthrough_stream(self, monkeypatch): + """Starlette abandons the body iterator when the client disconnects, so the + unbuffered passthrough branch must return _UpstreamClosingStreamingResponse, + whose shielded cleanup closes the upstream stream; that close is what flushes + buffered passthrough usage into spend logs.""" + processing_obj = self._build_processing_obj("gigachat") + monkeypatch.setattr(litellm, "callbacks", []) + upstream_closed = asyncio.Event() + + async def hanging_stream(): + try: + yield b"chunk-1" + await asyncio.Event().wait() + finally: + upstream_closed.set() + + result = await self._run(processing_obj, monkeypatch, [], stream=hanging_stream()) + + assert isinstance(result, _UpstreamClosingStreamingResponse) + + first_chunk_sent = asyncio.Event() + + async def receive(): + await first_chunk_sent.wait() + return {"type": "http.disconnect"} + + async def send(message): + if message["type"] == "http.response.body" and message.get("body"): + first_chunk_sent.set() + + await result({"type": "http"}, receive, send) + await asyncio.wait_for(upstream_closed.wait(), timeout=5) + @pytest.mark.asyncio async def test_non_bedrock_stream_is_not_buffered(self, monkeypatch): processing_obj = self._build_processing_obj("anthropic") From 1f80f93750a4c0c4c8717d9d34a69b6d93208caa Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 09:57:27 +0800 Subject: [PATCH 080/120] fix(guardrails): carry Anthropic url and file image sources through to guardrails _image_sources returned source["data"] only. An Anthropic image block has three shapes (types/llms/anthropic.py:259) and only the base64 one carries "data", so {"type": "url", "url": ...} yielded nothing and the image never reached any guardrail at all. This is not Bedrock-specific. Five guardrails consume GenericGuardrailAPIInputs["images"] (vigil_guard, custom_code, deepkeep, straiker, generic_guardrail_api) and every one of them was blind to url sources on /v1/messages. base64 now returns a data URI rather than the bare payload. A consumer otherwise has no way to recover media_type, and an API like Bedrock's ApplyGuardrail needs the format to build its request. The file shape stays unresolvable here: the bytes live behind the Files API and this extractor has no client to fetch them. Documented rather than silently dropped, so a consumer treating a missing entry as "no image to scan" is a known gap and not a surprise. Co-Authored-By: Claude Opus 5 (1M context) --- .../chat/guardrail_translation/handler.py | 32 +++++++++++++++++-- 1 file changed, 30 insertions(+), 2 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index eb6dbd3e5b3..4b763d11652 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -859,12 +859,40 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _image_sources(block: Mapping[str, object]) -> tuple[str, ...]: + """Normalize an Anthropic image block into strings a guardrail can read. + + `source` is one of three shapes (types/llms/anthropic.py:259): + + {"type": "base64", "media_type": "image/png", "data": ""} + {"type": "url", "url": "https://..."} + {"type": "file", "file_id": "..."} + + base64 is returned as a data URI rather than the bare payload: consumers of + ``GenericGuardrailAPIInputs["images"]`` otherwise have no way to know the + format, and an API like Bedrock's ApplyGuardrail requires it. url is passed + through so the consumer can fetch it under its own SSRF policy. + + file is not resolvable here (the bytes live behind the Files API), so it + yields nothing. That is a silent gap for any consumer that treats a missing + entry as "no image to scan"; scanning a file_id needs a fetch this extractor + has no client for. + """ source: Final = block.get("source") if not isinstance(source, Mapping): return () - # Could be base64 or url + + source_type: Final = source.get("type") + if source_type == "url": + url: Final = source.get("url") + return (url,) if isinstance(url, str) and url else () + data: Final = source.get("data") - return (data,) if data else () + if not isinstance(data, str) or not data: + return () + media_type: Final = source.get("media_type") + if isinstance(media_type, str) and media_type: + return (f"data:{media_type};base64,{data}",) + return (data,) async def _apply_guardrail_responses_to_input( self, From 0e7562dbc692812dfb5cb02ad72434d54a25b715 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 11:22:51 +0800 Subject: [PATCH 081/120] test(guardrails): cover every Anthropic image source shape in the extractor's own suite _image_sources had no test asserting what it extracts. The existing image tests live on the Bedrock side and all use base64 without a media_type, which is the one path the fix left unchanged, so both behaviors it does change went unverified: the url shape reaching the guardrail at all, and base64 arriving as a data URI. Against the pre-fix extractor the url case sees [] and the media_type case sees ['AAAA'] instead of ['data:image/png;base64,AAAA']. The remaining three assert behavior the fix deliberately preserves -- bare base64 passed through, a file source yielding nothing, a malformed source dropped rather than handed on for a consumer to choke on. Each message carries a text block because a message with no text never reaches the guardrail, which would make every source shape look equally dropped. Co-Authored-By: Claude Opus 5 (1M context) --- .../test_anthropic_guardrail_handler.py | 86 +++++++++++++++++++ 1 file changed, 86 insertions(+) diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index af3ccd65b11..2b606b9639a 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -1490,6 +1490,92 @@ class MockCanaryMaskingGuardrail(CustomGuardrail): return inputs +class TestAnthropicMessagesImageSources: + """An Anthropic image block has three source shapes (types/llms/anthropic.py:259). + + Only the base64 one carries "data", so reading that key alone drops url images + entirely -- for every guardrail consuming GenericGuardrailAPIInputs["images"], + not just Bedrock. + """ + + def _data(self, messages): + return {"model": "claude-sonnet-4-5", "messages": messages} + + async def _images_seen(self, content) -> list[str]: + handler = AnthropicMessagesHandler() + + class ImageRecordingGuardrail(MockCanaryMaskingGuardrail): + def __init__(self): + super().__init__() + self.seen_images: list[str] = [] # mutable-ok: accumulator for the assertion + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + self.seen_images.extend(inputs.get("images") or []) + return await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + + guardrail = ImageRecordingGuardrail() + # The text block is what gets the guardrail invoked at all: a message with + # no text gives the handler nothing to scan, so it never reaches the + # guardrail and every source shape would look equally "dropped". + await handler.process_input_messages( + data=self._data([{"role": "user", "content": [{"type": "text", "text": "describe it"}, *content]}]), + guardrail_to_apply=guardrail, + ) + return guardrail.seen_images + + @pytest.mark.asyncio + async def test_url_source_reaches_the_guardrail(self): + """A url source has no "data" key, so it used to yield nothing at all.""" + seen = await self._images_seen( + [{"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}] + ) + + assert seen == ["https://example.com/a.png"] + + @pytest.mark.asyncio + async def test_base64_source_carries_its_media_type(self): + """Bare base64 leaves the consumer no way to recover the format. + + An API like Bedrock's ApplyGuardrail needs it to build the request, so the + media_type travels with the payload as a data URI. + """ + seen = await self._images_seen( + [{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AAAA"}}] + ) + + assert seen == ["data:image/png;base64,AAAA"] + + @pytest.mark.asyncio + async def test_base64_source_without_a_media_type_is_passed_through(self): + """There is no format to attach, so the payload goes through unchanged.""" + seen = await self._images_seen([{"type": "image", "source": {"type": "base64", "data": "AAAA"}}]) + + assert seen == ["AAAA"] + + @pytest.mark.asyncio + async def test_file_source_yields_nothing(self): + """The bytes live behind the Files API and this extractor has no client. + + Documented as a known gap rather than silently handed on as a file_id string, + which a consumer would try to decode as an image. + """ + seen = await self._images_seen([{"type": "image", "source": {"type": "file", "file_id": "file_abc"}}]) + + assert seen == [] + + @pytest.mark.asyncio + async def test_a_malformed_source_is_dropped_rather_than_passed_on(self): + seen = await self._images_seen( + [ + {"type": "image", "source": {"type": "base64"}}, + {"type": "image", "source": {"type": "url"}}, + {"type": "image", "source": {"type": "base64", "data": ""}}, + ] + ) + + assert seen == [] + + class TestAnthropicMessagesToolResultScanning: """LIT-5251: tool_result blocks carry whatever a client's local tool fetched, so they are the request-path payload an indirect prompt injection actually arrives in. From bb51c121cf943b384fd00b0a2b106deca4dd5d53 Mon Sep 17 00:00:00 2001 From: "feng.tsai" Date: Mon, 31 Aug 2026 12:21:24 +0800 Subject: [PATCH 082/120] docs: reference the source union by type instead of a line number The line number went stale when the base moved. --- litellm/llms/anthropic/chat/guardrail_translation/handler.py | 2 +- .../guardrail_translation/test_anthropic_guardrail_handler.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 4b763d11652..5aa0dc94b89 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -861,7 +861,7 @@ class AnthropicMessagesHandler(BaseTranslation): def _image_sources(block: Mapping[str, object]) -> tuple[str, ...]: """Normalize an Anthropic image block into strings a guardrail can read. - `source` is one of three shapes (types/llms/anthropic.py:259): + `source` is one of three shapes (`AnthropicMessagesImageParam.source`): {"type": "base64", "media_type": "image/png", "data": ""} {"type": "url", "url": "https://..."} diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 2b606b9639a..0fe7730e91e 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -1491,7 +1491,7 @@ class MockCanaryMaskingGuardrail(CustomGuardrail): class TestAnthropicMessagesImageSources: - """An Anthropic image block has three source shapes (types/llms/anthropic.py:259). + """An Anthropic image block has three source shapes (`AnthropicMessagesImageParam.source`). Only the base64 one carries "data", so reading that key alone drops url images entirely -- for every guardrail consuming GenericGuardrailAPIInputs["images"], From dc034a086a17d349e005fc2b547ed8a6b9c2654e Mon Sep 17 00:00:00 2001 From: "feng.tsai" Date: Mon, 31 Aug 2026 13:26:02 +0800 Subject: [PATCH 083/120] docs: trim the _image_sources docstring It restated the source union that the type definition already carries. --- .../chat/guardrail_translation/handler.py | 18 +++--------------- 1 file changed, 3 insertions(+), 15 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 5aa0dc94b89..9395cc3d33e 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -861,21 +861,9 @@ class AnthropicMessagesHandler(BaseTranslation): def _image_sources(block: Mapping[str, object]) -> tuple[str, ...]: """Normalize an Anthropic image block into strings a guardrail can read. - `source` is one of three shapes (`AnthropicMessagesImageParam.source`): - - {"type": "base64", "media_type": "image/png", "data": ""} - {"type": "url", "url": "https://..."} - {"type": "file", "file_id": "..."} - - base64 is returned as a data URI rather than the bare payload: consumers of - ``GenericGuardrailAPIInputs["images"]`` otherwise have no way to know the - format, and an API like Bedrock's ApplyGuardrail requires it. url is passed - through so the consumer can fetch it under its own SSRF policy. - - file is not resolvable here (the bytes live behind the Files API), so it - yields nothing. That is a silent gap for any consumer that treats a missing - entry as "no image to scan"; scanning a file_id needs a fetch this extractor - has no client for. + base64 becomes a data URI so the format travels with the payload, which is what + the OpenAI path already puts in this field. A file source yields nothing: those + bytes live behind the Files API and this extractor has no client to fetch them. """ source: Final = block.get("source") if not isinstance(source, Mapping): From 69fc4496664f3b468bb4384cbcae91d0b0990c23 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 31 Aug 2026 07:54:31 +0000 Subject: [PATCH 084/120] chore(techdebt): drop restating comments from the 2026-08-30 window Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/google_genai/adapters/transformation.py | 1 - litellm/llms/runwayml/videos/transformation.py | 2 -- litellm/proxy/response_polling/background_streaming.py | 7 ++----- 3 files changed, 2 insertions(+), 8 deletions(-) diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index b73a1443330..8ea19deef4c 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -108,7 +108,6 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): def __init__(self, completion_stream: object): self.sent_first_chunk = False - # State tracking for accumulating partial tool calls self.accumulated_tool_calls = dict[int, _ToolCallAccumulator]() self._returned_response = False super().__init__(completion_stream) diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index c7289eb47f5..c7696a1cb29 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -160,14 +160,12 @@ class RunwayMLVideoConfig(BaseVideoConfig): **self._prompt_image_param(video_create_optional_params), **self._ratio_param(video_create_optional_params), **self._duration_param(video_create_optional_params), - # Pass through other parameters that aren't OpenAI-specific **{key: value for key, value in video_create_optional_params.items() if key not in supported_openai_params}, } @staticmethod def _prompt_image_param(video_create_optional_params: VideoCreateOptionalRequestParams) -> Mapping[str, object]: # Handle input_reference parameter - map to promptImage - # RunwayML supports URLs and data URIs directly if "input_reference" in video_create_optional_params: return {"promptImage": video_create_optional_params["input_reference"]} return {} diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index f755c6d478b..b03122966c0 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -144,10 +144,8 @@ async def background_streaming_task( # Process streaming response following OpenAI events format # https://platform.openai.com/docs/api-reference/responses-streaming - output_items: Final = dict[str, _OutputItem]() # Track output items by ID - accumulated_text: Final = dict[ - tuple[str, int], str - ]() # Track accumulated text deltas by (item_id, content_index) + output_items: Final = dict[str, _OutputItem]() + accumulated_text: Final = dict[tuple[str, int], str]() # ResponsesAPIResponse fields to extract from response.completed usage_data = None @@ -262,7 +260,6 @@ async def background_streaming_task( if "content" in delta_item: content_list = delta_item["content"] if content_index < len(content_list): - # Update existing content part with accumulated text content_entry = content_list[content_index] if isinstance(content_entry, dict): content_entry["text"] = accumulated_text[key] From b2df72f980c1f65dbf2e7200ecef4aa1a2a45aa5 Mon Sep 17 00:00:00 2001 From: nuernber Date: Mon, 31 Aug 2026 09:05:37 -0700 Subject: [PATCH 085/120] chore: ratchet down lint/type budgets after disconnect-billing merge fix --- basedpyright-code-budget.json | 10 +++++----- type-discipline-budget.json | 6 +++--- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index a542d025e74..f62d5f95256 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5658 }, "reportMissingTypeArgument": { - "limit": 15656 + "limit": 15655 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44831 + "limit": 44829 }, "reportUnknownLambdaType": { "limit": 109 }, "reportUnknownMemberType": { - "limit": 39267 + "limit": 39265 }, "reportUnknownParameterType": { - "limit": 19987 + "limit": 19986 }, "reportUnknownVariableType": { - "limit": 30922 + "limit": 30921 }, "reportUnnecessaryCast": { "limit": 117 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 1f2651f5f25..c5349689ec0 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23002 + "limit": 23001 }, "LIT002": { - "limit": 27145 + "limit": 27144 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16730 + "limit": 16729 }, "LIT011": { "limit": 5577 From e1fece511a20298f16926a8b595850a093cfe34f Mon Sep 17 00:00:00 2001 From: nuernber Date: Mon, 31 Aug 2026 09:11:51 -0700 Subject: [PATCH 086/120] test(anthropic): fix PT012 lint violation in upstream-error regression test --- .../messages/test_streaming_iterator.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index f0a3c7fdff0..6a1b6f417a2 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -501,10 +501,14 @@ async def test_async_sse_wrapper_reraises_upstream_error_to_connected_client(): ) received = [] - with pytest.raises(_ProviderStreamError) as excinfo: + + async def _drain(): async for chunk in iterator.async_sse_wrapper(_failing_stream()): received.append(chunk) + with pytest.raises(_ProviderStreamError) as excinfo: + await _drain() + assert excinfo.value.status_code == 529 assert received assert not any(c.startswith(b"event: error\n") for c in received) From 79fd2f4872ca142cc9a4df8ad9053e93905eb315 Mon Sep 17 00:00:00 2001 From: nuernber Date: Mon, 31 Aug 2026 09:15:36 -0700 Subject: [PATCH 087/120] test(anthropic): cover ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS=0 fallback to partial billing --- litellm/constants.py | 5 ++- .../messages/test_streaming_iterator.py | 42 +++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index 2c46153cff6..b03e122d09c 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -489,7 +489,10 @@ MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = int(os.getenv("MAX_LANGFUSE_INITIALIZE # bounded so a slow client throttles the upstream pump instead of letting it # buffer the whole response in memory; the detached-drain cap bounds how many # post-disconnect drains may run concurrently so client behavior can't create -# unbounded worker state. +# unbounded worker state. Setting the cap to 0 disables detached draining +# entirely: every post-disconnect pump bills whatever partial output it has +# already collected and aborts the upstream stream immediately, instead of +# continuing to drain for the real terminal usage. ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE: Final = int( os.getenv("ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", "1024") ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 6a1b6f417a2..8c3b2852345 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -635,6 +635,48 @@ async def test_async_sse_wrapper_bills_partial_when_detached_drain_cap_reached(m streaming_iterator_module._DETACHED_STREAM_DRAINS.discard(holder) +@pytest.mark.asyncio +async def test_async_sse_wrapper_bills_partial_when_detached_drains_disabled(monkeypatch): + """ + Regression: ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS=0 must disable + detached draining entirely, not just shrink the cap. With no slots ever + available, the very first post-disconnect chunk must fall back to partial + spend logging instead of hanging on a cap that's unreachable. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 0) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4) + + tail_reached = False + + async def _long_stream(): + nonlocal tail_reached + yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}} + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}} + for i in range(100): + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"more{i}"}} + tail_reached = True + yield {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 42}} + yield {"type": "message_stop"} + + iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("drains_disabled"), request_body={}) + gen = iterator.async_sse_wrapper(_long_stream()) + await gen.__anext__() # message_start + await gen.__anext__() # first delta + await gen.aclose() # client disconnects; 100+ chunks remain upstream + + for _ in range(200): + if iterator.logged_chunks: + break + await asyncio.sleep(0) + + assert iterator.logged_chunks, "pump never billed with detached drains disabled" + assert len(iterator.logged_chunks) <= 2 + streaming_iterator_module.ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE + assert len(iterator.logged_chunks) < 100 + assert not any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks) + assert tail_reached is False, "pump kept draining despite detached drains being disabled" + assert len(streaming_iterator_module._DETACHED_STREAM_DRAINS) == 0 + + @pytest.mark.asyncio async def test_async_sse_wrapper_aborts_upstream_when_detached_drain_cap_reached(monkeypatch): """ From 8e92f989051db0a518ef244e9961efc7b487a2fe Mon Sep 17 00:00:00 2001 From: nuernber Date: Mon, 31 Aug 2026 09:18:12 -0700 Subject: [PATCH 088/120] docs(constants): clarify which env var the cap=0 fallback note applies to --- litellm/constants.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index b03e122d09c..7a295c37010 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -489,13 +489,14 @@ MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = int(os.getenv("MAX_LANGFUSE_INITIALIZE # bounded so a slow client throttles the upstream pump instead of letting it # buffer the whole response in memory; the detached-drain cap bounds how many # post-disconnect drains may run concurrently so client behavior can't create -# unbounded worker state. Setting the cap to 0 disables detached draining -# entirely: every post-disconnect pump bills whatever partial output it has -# already collected and aborts the upstream stream immediately, instead of -# continuing to drain for the real terminal usage. +# unbounded worker state. ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE: Final = int( os.getenv("ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", "1024") ) +# Setting this to 0 disables detached draining entirely: every post-disconnect +# pump bills whatever partial output it has already collected and aborts the +# upstream stream immediately, instead of continuing to drain for the real +# terminal usage. ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: Final = int( os.getenv("ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", "100") ) From 2f50988fed00657dc1c0484828dde37590701930 Mon Sep 17 00:00:00 2001 From: nuernber Date: Mon, 31 Aug 2026 09:30:29 -0700 Subject: [PATCH 089/120] fix(anthropic): resolve TRY300 lint violation and ratchet budgets after litellm_internal_staging merge --- basedpyright-code-budget.json | 10 +++++----- .../messages/streaming_iterator.py | 3 ++- type-discipline-budget.json | 6 +++--- 3 files changed, 10 insertions(+), 9 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index d60c3e9c0af..7d105cdcfe0 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5611 }, "reportMissingTypeArgument": { - "limit": 15350 + "limit": 15349 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44368 + "limit": 44366 }, "reportUnknownLambdaType": { "limit": 109 }, "reportUnknownMemberType": { - "limit": 38468 + "limit": 38466 }, "reportUnknownParameterType": { - "limit": 19665 + "limit": 19664 }, "reportUnknownVariableType": { - "limit": 30066 + "limit": 30065 }, "reportUnnecessaryCast": { "limit": 111 diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 9e115c6624d..d822d09771b 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -541,9 +541,10 @@ class BaseAnthropicMessagesStreamingIterator: return False try: queue.put_nowait(item) - return True except asyncio.QueueFull: pass + else: + return True put_task: Final = asyncio.ensure_future(queue.put(item)) detached_task: Final = asyncio.ensure_future(client_detached.wait()) try: diff --git a/type-discipline-budget.json b/type-discipline-budget.json index ab34775c460..143cbc91787 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22521 + "limit": 22520 }, "LIT002": { - "limit": 26820 + "limit": 26819 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16546 + "limit": 16545 }, "LIT011": { "limit": 5575 From 0e78c5bff7cca880dca8d3c726df4606ffb71a95 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 10:17:27 -0700 Subject: [PATCH 090/120] fix(anthropic_messages): dispatch deferred spend logging when the client disconnects mid-relay When the pump finishes draining while the client is still connected, billing is deferred to the proxy's post-response hook, which only fires on a normally completed response. A client disconnect before the relay consumed the queued tail tore the generator down past that hook, so the request logged no spend at all. The relay teardown now dispatches the stored deferred billing whenever it never reached the end-of-stream sentinel. Also drops the live pass_through_tests script: that CI job runs against a fixed config with no Bedrock model or AWS credentials, so it could only fail there. The scenario is covered by unit tests on the relay/pump seam. --- .../messages/streaming_iterator.py | 24 ++- ..._v1_messages_streaming_disconnect_spend.py | 155 ------------------ .../messages/test_streaming_iterator.py | 52 ++++++ 3 files changed, 75 insertions(+), 156 deletions(-) delete mode 100644 tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 9e115c6624d..b54dac18c95 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -480,16 +480,37 @@ class BaseAnthropicMessagesStreamingIterator: _UPSTREAM_PUMP_TASKS.add(pump_task) pump_task.add_done_callback(_UPSTREAM_PUMP_TASKS.discard) + reached_end = False # rebind-ok: flipped once the relay consumes the end-of-stream sentinel try: while True: item = await queue.get() if item is None: + reached_end = True break if isinstance(item, BaseException): raise item yield item finally: client_detached.set() + if not reached_end: + self._dispatch_pending_deferred_logging() + + def _dispatch_pending_deferred_logging(self) -> None: + """Fire deferred billing that a torn-down response would otherwise drop. + + When the pump finishes draining while the client is still connected it + stores the logging coroutine for ProxyLogging._fire_deferred_stream_logging, + which the proxy only fires on a normally completed response: a client + disconnect (GeneratorExit / CancelledError) re-raises past it. Without + this dispatch that window loses the spend row entirely. + """ + deferred_cb: Final = getattr(self.litellm_logging_obj, "_on_deferred_stream_complete", None) + deferred_args: Final = getattr(self.litellm_logging_obj, "_deferred_stream_complete_args", None) + if deferred_cb is None or deferred_args is None: + return + self.litellm_logging_obj._on_deferred_stream_complete = None + self.litellm_logging_obj._deferred_stream_complete_args = None + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=deferred_cb(*deferred_args)) async def _bill_collected_chunks( self, @@ -541,9 +562,10 @@ class BaseAnthropicMessagesStreamingIterator: return False try: queue.put_nowait(item) - return True except asyncio.QueueFull: pass + else: + return True put_task: Final = asyncio.ensure_future(queue.put(item)) detached_task: Final = asyncio.ensure_future(client_detached.wait()) try: diff --git a/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py b/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py deleted file mode 100644 index e69de720ea4..00000000000 --- a/tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py +++ /dev/null @@ -1,155 +0,0 @@ -""" -Regression test: /v1/messages streaming interrupted mid-stream must still -produce a spend-log entry. - -On v1.79.1 the proxy records spend for the partially-streamed request. -A refactor on `main` broke that path, so the same scenario now produces -zero spend-log rows. - -Run against a live proxy (e.g. ``litellm --config proxy_server_config.yaml``): - - pytest tests/pass_through_tests/test_v1_messages_streaming_disconnect_spend.py -s -""" - -import asyncio -import json -import uuid - -import aiohttp -import pytest - - -BASE_URL = "http://127.0.0.1:4000" # change appropriately -ADMIN_KEY = "sk-1234" - - -async def _generate_key(session: aiohttp.ClientSession) -> str: - """Create a fresh virtual key so spend is isolated.""" - url = f"{BASE_URL}/key/generate" - headers = {"Authorization": f"Bearer {ADMIN_KEY}", "Content-Type": "application/json"} - async with session.post(url, headers=headers, json={"models": []}) as resp: - assert resp.status == 200, f"key/generate failed: {await resp.text()}" - data = await resp.json() - return data["key"] - - -async def _get_spend_logs_by_spend_id(session: aiohttp.ClientSession, api_key: str, spend_id: str): - """Query /spend/logs by api_key then filter by spend_id in metadata.""" - url = f"{BASE_URL}/spend/logs?api_key={api_key}" - headers = {"Authorization": f"Bearer {ADMIN_KEY}", "Content-Type": "application/json"} - async with session.get(url, headers=headers) as resp: - assert resp.status == 200, f"spend/logs failed: {await resp.text()}" - all_logs = await resp.json() - if not isinstance(all_logs, list): - return [] - matched = [] - for log in all_logs: - meta = log.get("metadata") - if isinstance(meta, str): - meta = json.loads(meta) - if isinstance(meta, dict): - slm = meta.get("spend_logs_metadata") or {} - if slm.get("spend_id") == spend_id: - matched.append(log) - return matched - - -@pytest.mark.asyncio -@pytest.mark.flaky(retries=3, delay=2) -async def test_v1_messages_streaming_disconnect_has_spend_log(): - """ - 1. Send a streaming POST to /v1/messages. - 2. Read a few SSE chunks, then close the connection (simulating a client - disconnect / interruption). - 3. Wait for the proxy's async spend-tracking pipeline to flush. - 4. Assert that at least one spend-log row exists for the request. - - This PASSES on v1.79.1 and FAILS on the latest main branch. - """ - async with aiohttp.ClientSession( - timeout=aiohttp.ClientTimeout(total=60) - ) as session: - key = await _generate_key(session) - - spend_id = str(uuid.uuid4()) - - headers = { - "Authorization": f"Bearer {key}", - "Content-Type": "application/json", - "x-litellm-spend-logs-metadata": '{"spend_id": "' + spend_id + '"}', - } - - payload = { - "model": "us.anthropic.claude-sonnet-4-5-20250929-v1:0", - "max_tokens": 3000, - "stream": True, - "messages": [ - { - "role": "user", - "content": ( - f"Write several detailed paragraphs (at least 500 words) about the " - f"history of the Roman Empire. Unique id: {uuid.uuid4()}" - ), - } - ], - } - - chunks_read = 0 - - async with session.post( - f"{BASE_URL}/v1/messages", json=payload, headers=headers - ) as resp: - assert resp.status == 200, f"/v1/messages failed: {await resp.text()}" - - async for raw_line in resp.content: - line = raw_line.decode("utf-8", errors="replace").strip() - if not line: - continue - chunks_read += 1 - print(f" chunk #{chunks_read}: {line[:120]}") - if chunks_read >= 5: - break - - assert chunks_read >= 3, ( - f"Expected at least 3 chunks before disconnect, got {chunks_read}" - ) - - print( - f"\nDisconnected after {chunks_read} chunks. " - f"Waiting for spend pipeline to flush …" - ) - - spend_data = None - max_retries = 4 - for attempt in range(1, max_retries + 1): - await asyncio.sleep(10) - print(f" spend-log poll attempt {attempt}/{max_retries}") - spend_data = await _get_spend_logs_by_spend_id(session, key, spend_id) - if spend_data and len(spend_data) > 0: - print(f" ✓ found {len(spend_data)} spend-log row(s)") - break - print(" … not found yet") - - assert spend_data is not None and len(spend_data) > 0, ( - f"No spend-log entry found for spend_id={spend_id} " - f"after streaming disconnect. " - f"This is the regression: interrupted /v1/messages streams must " - f"still record spend." - ) - - log_entry = spend_data[0] - print( - f"\nSpend-log entry:\n{json.dumps(log_entry, indent=2, default=str)}" - ) - - prompt_tokens = log_entry.get("prompt_tokens", 0) - completion_tokens = log_entry.get("completion_tokens", 0) - assert prompt_tokens > 0, ( - "Spend-log row exists but has zero prompt tokens, so usage was not recorded." - ) - assert completion_tokens >= 100, ( - f"Spend-log completion_tokens={completion_tokens} is far below the full " - f"response Bedrock generated and billed. The interrupted stream was billed " - f"on the few chunks the client drained, not the full upstream output. " - f"chunks_read={chunks_read}, prompt_tokens={prompt_tokens}" - ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 8c3b2852345..2135397302e 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -472,6 +472,58 @@ async def test_async_sse_wrapper_bills_full_stream_when_client_reads_all(): assert not any(c.startswith(b"event: error\n") for c in iterator.logged_chunks) +@pytest.mark.asyncio +async def test_async_sse_wrapper_dispatches_deferred_logging_when_client_disconnects_mid_tail(): + """ + Regression: when the pump finishes draining while the client is still + connected, ``_handle_streaming_logging`` defers billing for the proxy's + post-response hook (``ProxyLogging._fire_deferred_stream_logging``), which + only fires on a normally completed response. If the client then disconnects + before consuming the queued tail, the response generator tears down via + GeneratorExit and that hook never runs. The relay teardown must dispatch + the stored deferred billing itself, or the request logs no spend at all. + """ + dispatched = [] + deferred_fired = asyncio.Event() + + def _deferred_stream_complete(logging_coroutine): + dispatched.append(logging_coroutine) + + async def _consume(): + logging_coroutine.close() + deferred_fired.set() + + return _consume() + + logging_obj = _make_logging_obj("test_deferred_dispatch_on_disconnect_mid_tail") + logging_obj._on_deferred_stream_complete = _deferred_stream_complete + iterator = BaseAnthropicMessagesStreamingIterator(litellm_logging_obj=logging_obj, request_body={}) + + async def _full_stream(): + for event in (*_STREAM_PREFIX, *_STREAM_TAIL): + yield event + + gen = iterator.async_sse_wrapper(_full_stream()) + client_chunks = [] + async for chunk in gen: + client_chunks.append(chunk) + if len(client_chunks) == len(_STREAM_PREFIX): + break + + for _ in range(100): + if getattr(logging_obj, "_deferred_stream_complete_args", None) is not None: + break + await asyncio.sleep(0.01) + assert getattr(logging_obj, "_deferred_stream_complete_args", None) is not None, "pump never deferred billing" + + await gen.aclose() + + assert len(dispatched) == 1, "relay teardown did not dispatch the deferred billing" + assert logging_obj._on_deferred_stream_complete is None + assert logging_obj._deferred_stream_complete_args is None + await asyncio.wait_for(deferred_fired.wait(), timeout=5) + + class _ProviderStreamError(Exception): """Stand-in for a provider-specific streaming failure carrying a status code.""" From 6b2ada2a780e1ebd4f4899833fbc7be9a3214123 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 11:49:09 -0700 Subject: [PATCH 091/120] fix(bedrock): per-response realtime usage deltas, spend-log event filter, single transcript completed --- litellm/llms/bedrock/realtime/handler.py | 70 +++++++--- .../llms/bedrock/realtime/transformation.py | 81 ++++++++---- .../realtime/test_bedrock_realtime_handler.py | 125 ++++++++++++++---- .../test_bedrock_realtime_transformation.py | 106 +++++++++++++++ 4 files changed, 313 insertions(+), 69 deletions(-) diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index ecd143dd087..42fe8941443 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -7,14 +7,17 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic. import asyncio import contextlib import json +from collections.abc import AsyncIterator, Mapping from typing import Final, Protocol from pydantic import JsonValue, TypeAdapter +import litellm from litellm._logging import _redact_string, verbose_proxy_logger from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeEventTypes from litellm.types.llms.openai import OpenAIRealtimeEvents from litellm.types.realtime import RealtimeResponseTransformInput @@ -34,6 +37,17 @@ def _json_str(value: JsonValue) -> str | None: return value if isinstance(value, str) else None +def _should_log_event(openai_message: Mapping[str, object]) -> bool: + logged_types: Final = ( + litellm.logged_real_time_event_types + if litellm.logged_real_time_event_types is not None + else DefaultLoggedRealTimeEventTypes + ) + if logged_types == "*": + return True + return openai_message.get("type") in logged_types + + class RealtimeClientWebSocket(Protocol): """The client-facing websocket surface the realtime bridge talks to.""" @@ -207,18 +221,22 @@ class BedrockRealtime(BaseAWSLLM): ) ) - logged_events: Final[list[OpenAIRealtimeEvents]] = [] # mutable-ok: filled across the stream loop - bedrock_to_client_task: Final = asyncio.create_task( - self._forward_bedrock_to_client( - bedrock_stream, - websocket, - transformation_config, - model, - logging_obj, - session_state, - logged_events, + async def forward_bedrock_and_collect_logged_events() -> tuple[OpenAIRealtimeEvents, ...]: + return tuple( + [ + event + async for event in self._forward_bedrock_to_client( + bedrock_stream, + websocket, + transformation_config, + model, + logging_obj, + session_state, + ) + ] ) - ) + + bedrock_to_client_task: Final = asyncio.create_task(forward_bedrock_and_collect_logged_events()) # Wait for both tasks to complete await asyncio.gather( @@ -227,9 +245,25 @@ class BedrockRealtime(BaseAWSLLM): return_exceptions=True, ) + forwarded_logged_events: Final = ( + bedrock_to_client_task.result() + if not bedrock_to_client_task.cancelled() and bedrock_to_client_task.exception() is None + else () + ) + logged_events: Final = ( + *forwarded_logged_events, + *( + leftover_event + for leftover_event in transformation_config.leftover_usage_done_events() + if _should_log_event(leftover_event) + ), + ) if logged_events: GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - logging_obj.dispatch_success_handlers(logged_events, prefer_async_handlers=True) + logging_obj.dispatch_success_handlers( + list(logged_events), # mutable-ok: realtime spend logging requires a list result + prefer_async_handlers=True, + ) ) except Exception as e: @@ -313,9 +347,8 @@ class BedrockRealtime(BaseAWSLLM): model: str, logging_obj: LiteLLMLogging, session_state: RealtimeResponseTransformInput, - logged_events: "list[OpenAIRealtimeEvents] | None" = None, # mutable-ok: caller-owned spend log accumulator - ): - """Forward messages from Bedrock stream to client WebSocket.""" + ) -> AsyncIterator[OpenAIRealtimeEvents]: + """Forward messages from Bedrock to the client, yielding the ones to record for spend logging.""" try: while True: # Receive from Bedrock @@ -363,13 +396,14 @@ class BedrockRealtime(BaseAWSLLM): ) # Send transformed messages to client - openai_messages = transformed_response.get("response", []) + response_value = transformed_response["response"] + openai_messages = response_value if isinstance(response_value, list) else (response_value,) for openai_message in openai_messages: - if logged_events is not None and isinstance(openai_message, dict): - logged_events.append(openai_message) message_json = json.dumps(openai_message) await client_ws.send_text(message_json) verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200]) + if _should_log_event(openai_message): + yield openai_message except Exception as e: verbose_proxy_logger.debug("Bedrock to client forwarding ended: %s", e, exc_info=True) diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index ec5eff0a6fc..28c2e446d10 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -41,7 +41,6 @@ from litellm.types.realtime import ( RealtimeResponseTransformInput, RealtimeResponseTypedDict, ) -from litellm.utils import get_empty_usage class BedrockContentEnd(BaseModel): @@ -118,7 +117,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self._user_transcript_active = False self._user_transcript_generation_stage: str | None = None self._user_item_id: str | None = None - self._latest_usage: OpenAIRealtimeResponseUsage | None = None + self._user_transcript_buffer = "" + self._cumulative_usage = BedrockUsageEvent() + self._reported_usage = BedrockUsageEvent() def validate_environment(self, headers: dict, model: str, api_key: str | None = None) -> dict: """Validate environment - no special validation needed for Bedrock.""" @@ -836,47 +837,79 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): return (speech_event,) def transform_usage_event(self, usage_event: BedrockUsageEvent) -> None: - """Record Bedrock usageEvent token totals for the next response.done.""" + """Record Bedrock's session-cumulative usage totals for the next response.done.""" verbose_logger.debug("Handling usageEvent") + self._cumulative_usage = usage_event + + def _take_usage_delta(self) -> OpenAIRealtimeResponseUsage: + """Usage for the response now completing: cumulative totals minus what prior response.done events reported.""" + prior: Final = self._reported_usage + latest: Final = self._cumulative_usage + self._reported_usage = latest input_details: Final[OpenAIRealtimeUsageTokenDetails] = { - "audio_tokens": usage_event.details.total.input.speechTokens, - "text_tokens": usage_event.details.total.input.textTokens, + "audio_tokens": latest.details.total.input.speechTokens - prior.details.total.input.speechTokens, + "text_tokens": latest.details.total.input.textTokens - prior.details.total.input.textTokens, "cached_tokens": 0, } output_details: Final[OpenAIRealtimeUsageTokenDetails] = { - "audio_tokens": usage_event.details.total.output.speechTokens, - "text_tokens": usage_event.details.total.output.textTokens, + "audio_tokens": latest.details.total.output.speechTokens - prior.details.total.output.speechTokens, + "text_tokens": latest.details.total.output.textTokens - prior.details.total.output.textTokens, } - latest_usage: Final[OpenAIRealtimeResponseUsage] = { - "input_tokens": usage_event.totalInputTokens, - "output_tokens": usage_event.totalOutputTokens, - "total_tokens": usage_event.totalTokens, + usage_delta: Final[OpenAIRealtimeResponseUsage] = { + "input_tokens": latest.totalInputTokens - prior.totalInputTokens, + "output_tokens": latest.totalOutputTokens - prior.totalOutputTokens, + "total_tokens": latest.totalTokens - prior.totalTokens, "input_token_details": input_details, "output_token_details": output_details, } - self._latest_usage = latest_usage + return usage_delta + + def leftover_usage_done_events(self) -> tuple[OpenAIRealtimeEvents, ...]: + """Logged-only response.done for usage Bedrock reports after the final turn's contentEnd.""" + if self._cumulative_usage == self._reported_usage: + return () + usage: Final = self._take_usage_delta() + leftover_done: Final = OpenAIRealtimeDoneEvent( + type="response.done", + event_id=f"event_{uuid.uuid4()}", + response=OpenAIRealtimeResponseDoneObject( + object="realtime.response", + id=f"resp_{uuid.uuid4()}", + status="completed", + conversation_id=f"conv_{uuid.uuid4()}", + usage=dict(usage), # mutable-ok: OpenAIRealtimeResponseDoneObject types usage as plain dict + ), + ) + return (leftover_done,) def transform_user_transcript_event(self, transcript: str) -> tuple[OpenAIRealtimeEvents, ...]: - """Transform a USER-role Bedrock textOutput (ASR transcript) to OpenAI transcription events.""" + """Transform a USER-role Bedrock textOutput (ASR transcript) to an OpenAI transcription delta.""" verbose_logger.debug("Handling USER textOutput (ASR transcript)") - item_id: Final = self._current_user_item_id() delta_event: Final[OpenAIRealtimeInputAudioTranscriptionDelta] = { "type": "conversation.item.input_audio_transcription.delta", "event_id": f"event_{uuid.uuid4()}", - "item_id": item_id, + "item_id": self._current_user_item_id(), "content_index": 0, "delta": transcript, } - if self._user_transcript_generation_stage == "SPECULATIVE": - return (delta_event,) + if self._user_transcript_generation_stage != "SPECULATIVE": + self._user_transcript_buffer += transcript + return (delta_event,) + + def user_transcript_completed_events(self) -> tuple[OpenAIRealtimeEvents, ...]: + """One completed event with the full transcript once the FINAL user content block ends.""" + transcript: Final = self._user_transcript_buffer + if not transcript: + return () + self._user_transcript_buffer = "" completed_event: Final[OpenAIRealtimeInputAudioTranscriptionCompleted] = { "type": "conversation.item.input_audio_transcription.completed", "event_id": f"event_{uuid.uuid4()}", - "item_id": item_id, + "item_id": self._current_user_item_id(), "content_index": 0, "transcript": transcript, } - return (delta_event, completed_event) + return (completed_event,) def transform_text_output_event( self, @@ -1096,14 +1129,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): if not current_response_id or not current_conversation_id: return [], None, None, None - empty_usage: Final = get_empty_usage() - zero_usage: Final[OpenAIRealtimeResponseUsage] = { - "input_tokens": empty_usage.prompt_tokens, - "output_tokens": empty_usage.completion_tokens, - "total_tokens": empty_usage.total_tokens, - } - usage: Final = self._latest_usage or zero_usage - self._latest_usage = None + usage: Final = self._take_usage_delta() response_done: Final = OpenAIRealtimeDoneEvent( type="response.done", event_id=f"event_{uuid.uuid4()}", @@ -1324,6 +1350,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): elif "contentEnd" in event and self._user_transcript_active: self._user_transcript_active = False self._user_transcript_generation_stage = None + returned_messages.extend(self.user_transcript_completed_events()) elif "contentEnd" in event: events, current_delta_chunks = self.transform_content_end_event( diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index aa6573884e2..0ea5b7ad4a1 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -6,7 +6,7 @@ from unittest.mock import MagicMock import pytest - +import litellm from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.handler import BedrockRealtime from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig @@ -124,14 +124,6 @@ class ScriptedBedrockStream: return (None, self._receiver) -class ImmediatelyEndingBedrockStream: - def __init__(self): - self.input_stream = FakeInputStream() - - async def await_output(self): - return (None, EndedBedrockReceiver()) - - class FakeStaticCredentialsResolver: pass @@ -171,7 +163,7 @@ def stub_aws_sdk_client(monkeypatch): async def invoke_model_with_bidirectional_stream(self, operation_input): captured["operation_input"] = operation_input - return ImmediatelyEndingBedrockStream() + return ScriptedBedrockStream(captured.get("scripted_payloads", [])) package = types.ModuleType("aws_sdk_bedrock_runtime") client_module = types.ModuleType("aws_sdk_bedrock_runtime.client") @@ -292,7 +284,40 @@ class TestBedrockRealtimeHandler: assert stream.input_stream.closed @pytest.mark.asyncio - async def test_forwarded_events_are_collected_for_spend_logging(self): + async def test_forwarded_events_are_filtered_to_logged_types_for_spend_logging(self): + handler = BedrockRealtime() + stream = ScriptedBedrockStream( + [ + json.dumps({"event": {"userSpeechStart": {}}}), + json.dumps({"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}), + json.dumps({"event": {"textOutput": {"content": "Hi"}}}), + json.dumps({"event": {"contentEnd": {"stopReason": "END_TURN"}}}), + ] + ) + client_ws = RealtimeClientWS() + + logged_events = [ + event + async for event in handler._forward_bedrock_to_client( + stream, + client_ws, + BedrockRealtimeConfig(), + "amazon.nova-sonic-v1:0", + FakeLogging(), + {}, + ) + ] + + assert [event["type"] for event in logged_events] == ["response.done"] + sent_types = [json.loads(message)["type"] for message in client_ws.sent_to_client] + assert "input_audio_buffer.speech_started" in sent_types + assert "response.text.delta" in sent_types + assert "response.done" in sent_types + assert client_ws.closed + + @pytest.mark.asyncio + async def test_logged_event_types_star_collects_every_forwarded_event(self, monkeypatch): + monkeypatch.setattr(litellm, "logged_real_time_event_types", "*") handler = BedrockRealtime() stream = ScriptedBedrockStream( [ @@ -301,37 +326,89 @@ class TestBedrockRealtimeHandler: ] ) client_ws = RealtimeClientWS() - logged_events = [] - await handler._forward_bedrock_to_client( - stream, - client_ws, - BedrockRealtimeConfig(), - "amazon.nova-sonic-v1:0", - FakeLogging(), - {}, - logged_events, - ) + logged_events = [ + event + async for event in handler._forward_bedrock_to_client( + stream, + client_ws, + BedrockRealtimeConfig(), + "amazon.nova-sonic-v1:0", + FakeLogging(), + {}, + ) + ] assert [event["type"] for event in logged_events] == [ "input_audio_buffer.speech_started", "input_audio_buffer.speech_stopped", ] - assert client_ws.closed + + @pytest.mark.asyncio + async def test_trailing_usage_after_last_done_is_dispatched_for_spend(self, stub_aws_sdk_client, monkeypatch): + import litellm.llms.bedrock.realtime.handler as handler_module + + dispatched = {} + + class RecordingLogging(FakeLogging): + async def dispatch_success_handlers(self, result=None, prefer_async_handlers=False, **kwargs): + dispatched["events"] = result + + class RecordingLoggingWorker: + def ensure_initialized_and_enqueue(self, coro): + dispatched["coro"] = coro + + monkeypatch.setattr(handler_module, "GLOBAL_LOGGING_WORKER", RecordingLoggingWorker()) + stub_aws_sdk_client["scripted_payloads"] = [ + json.dumps( + { + "event": { + "usageEvent": { + "totalInputTokens": 3, + "totalOutputTokens": 6, + "totalTokens": 9, + "details": { + "total": { + "input": {"speechTokens": 3, "textTokens": 0}, + "output": {"speechTokens": 0, "textTokens": 6}, + } + }, + } + } + } + ) + ] + + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=RealtimeClientWS(), + logging_obj=RecordingLogging(), + aws_region_name="us-east-1", + aws_access_key_id="k", + aws_secret_access_key="s", + ) + await dispatched["coro"] + + assert [event["type"] for event in dispatched["events"]] == ["response.done"] + usage = dispatched["events"][0]["response"]["usage"] + assert (usage["input_tokens"], usage["output_tokens"], usage["total_tokens"]) == (3, 6, 9) + assert usage["input_token_details"] == {"audio_tokens": 3, "text_tokens": 0, "cached_tokens": 0} + assert usage["output_token_details"] == {"audio_tokens": 0, "text_tokens": 6} @pytest.mark.asyncio async def test_bedrock_stream_end_closes_client_websocket(self): handler = BedrockRealtime() client_ws = ClosableClientWS() - await handler._forward_bedrock_to_client( + async for _ in handler._forward_bedrock_to_client( EndedBedrockStream(), client_ws, BedrockRealtimeConfig(), "amazon.nova-sonic-v1:0", MagicMock(), {}, - ) + ): + pass assert client_ws.closed diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py index dbe84cacb29..a74f03449a1 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py @@ -1025,6 +1025,112 @@ class TestBedrockRealtimeUserEventsAndUsage: assert usage["output_tokens"] == 0 assert usage["total_tokens"] == 0 + @staticmethod + def _usage_event(total_input, total_output, in_speech, in_text, out_speech, out_text): + return { + "event": { + "usageEvent": { + "totalInputTokens": total_input, + "totalOutputTokens": total_output, + "totalTokens": total_input + total_output, + "details": { + "total": { + "input": {"speechTokens": in_speech, "textTokens": in_text}, + "output": {"speechTokens": out_speech, "textTokens": out_text}, + } + }, + } + } + } + + _ASSISTANT_TURN = ( + {"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}, + {"event": {"textOutput": {"content": "Hi"}}}, + {"event": {"contentEnd": {"stopReason": "END_TURN"}}}, + ) + + def test_multi_turn_usage_reports_per_response_deltas_not_cumulative_totals(self): + events = self._run( + BedrockRealtimeConfig(), + [ + self._usage_event(25, 40, in_speech=20, in_text=5, out_speech=30, out_text=10), + *self._ASSISTANT_TURN, + self._usage_event(40, 100, in_speech=30, in_text=10, out_speech=75, out_text=25), + *self._ASSISTANT_TURN, + ], + ) + usages = [e["response"]["usage"] for e in events if e["type"] == "response.done"] + assert len(usages) == 2 + assert (usages[0]["input_tokens"], usages[0]["output_tokens"], usages[0]["total_tokens"]) == (25, 40, 65) + assert (usages[1]["input_tokens"], usages[1]["output_tokens"], usages[1]["total_tokens"]) == (15, 60, 75) + assert usages[1]["input_token_details"] == {"audio_tokens": 10, "text_tokens": 5, "cached_tokens": 0} + assert usages[1]["output_token_details"] == {"audio_tokens": 45, "text_tokens": 15} + assert sum(u["total_tokens"] for u in usages) == 140 + + def test_usage_reported_after_last_response_done_flushes_as_logged_only_done(self): + config = BedrockRealtimeConfig() + self._run( + config, + [ + self._usage_event(25, 40, in_speech=20, in_text=5, out_speech=30, out_text=10), + *self._ASSISTANT_TURN, + ], + ) + assert config.leftover_usage_done_events() == () + + self._run(config, [self._usage_event(25, 46, in_speech=20, in_text=5, out_speech=30, out_text=16)]) + leftover = config.leftover_usage_done_events() + assert len(leftover) == 1 + assert leftover[0]["type"] == "response.done" + usage = leftover[0]["response"]["usage"] + assert (usage["input_tokens"], usage["output_tokens"], usage["total_tokens"]) == (0, 6, 6) + assert usage["output_token_details"] == {"audio_tokens": 0, "text_tokens": 6} + assert config.leftover_usage_done_events() == () + + def test_final_transcript_fragments_emit_one_completed_with_full_transcript(self): + events = self._run( + BedrockRealtimeConfig(), + [ + { + "event": { + "contentStart": { + "role": "USER", + "type": "TEXT", + "additionalModelFields": json.dumps({"generationStage": "FINAL"}), + } + } + }, + {"event": {"textOutput": {"content": "What is the "}}}, + {"event": {"textOutput": {"content": "capital of France?"}}}, + {"event": {"contentEnd": {"stopReason": "PARTIAL_TURN"}}}, + ], + ) + deltas = [e for e in events if e["type"] == "conversation.item.input_audio_transcription.delta"] + completed = [e for e in events if e["type"] == "conversation.item.input_audio_transcription.completed"] + assert [d["delta"] for d in deltas] == ["What is the ", "capital of France?"] + assert len(completed) == 1 + assert completed[0]["transcript"] == "What is the capital of France?" + assert {e["item_id"] for e in deltas + completed} == {completed[0]["item_id"]} + + def test_speculative_transcript_block_end_emits_no_completed(self): + events = self._run( + BedrockRealtimeConfig(), + [ + { + "event": { + "contentStart": { + "role": "USER", + "type": "TEXT", + "additionalModelFields": json.dumps({"generationStage": "SPECULATIVE"}), + } + } + }, + {"event": {"textOutput": {"content": "rea"}}}, + {"event": {"contentEnd": {"stopReason": "PARTIAL_TURN"}}}, + ], + ) + assert [e["type"] for e in events] == ["conversation.item.input_audio_transcription.delta"] + if __name__ == "__main__": pytest.main([__file__, "-v"]) From 666d8737aa8766ac293995e565d91c1faf82069e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:19:41 -0700 Subject: [PATCH 092/120] fix(init): ignore pydantic ReadOnly TypedDict warning that floods proxy boot --- litellm/__init__.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/litellm/__init__.py b/litellm/__init__.py index c83e72a78b4..1447e05fdf7 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -7,6 +7,9 @@ warnings.filterwarnings("ignore", message=".*conflict with protected namespace.* # Suppress Pydantic 2.11+ deprecation warning about accessing model_fields on instances # This warning can accumulate during streaming and cause memory leaks warnings.filterwarnings("ignore", message=".*Accessing the.*attribute on the instance is deprecated.*") +# ReadOnly on TypedDict fields is repo-wide static discipline (LIT012); pydantic warns it +# cannot enforce it at runtime, which floods proxy boot once such a type is schema-walked +warnings.filterwarnings("ignore", message=".*`ReadOnly` qualifier.*") ### INIT VARIABLES ######################### import threading import os From 46e090d2f333e9974e0240cac4a4ce8fdb8840e4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:47:26 -0700 Subject: [PATCH 093/120] fix(anthropic_messages): bill partial spend when a queued pump error is never consumed When the upstream errors while the client is still connected, the pump forwards the exception through the relay queue so the proxy's failure handling re-raises it. If the client disconnects before consuming that queued exception, neither the failure hook nor billing ran and the spend row was lost. The pump now waits for client detach and, if the exception was never consumed, salvages partial spend like the post-disconnect error path. Also rewrites the bedrock disconnect logging test to the detached-pump contract: billing fires after the upstream drain completes, not synchronously at aclose(). --- .../messages/streaming_iterator.py | 20 ++++++++-- .../messages/test_streaming_iterator.py | 37 +++++++++++++++++++ .../test_anthropic_claude3_transformation.py | 23 ++++++++---- 3 files changed, 70 insertions(+), 10 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index b54dac18c95..55a64dedee4 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -157,6 +157,17 @@ def _try_claim_detached_drain_slot() -> bool: return True +def _exception_left_unconsumed(queue: "asyncio.Queue[bytes | None | BaseException]", exc: BaseException) -> bool: + """After client detach the relay never reads the queue again, so drain it here. + + The forwarded exception still sitting in the queue means the relay tore + down before re-raising it, so the proxy's failure handling never ran and + the caller must salvage spend itself. + """ + remaining: Final = tuple(queue.get_nowait() for _ in range(queue.qsize())) + return any(item is exc for item in remaining) + + def _sse_event(event_type: str, payload: Mapping[str, object]) -> bytes: return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode() @@ -637,13 +648,16 @@ class BaseAnthropicMessagesStreamingIterator: Handing the original exception to the client-facing generator lets it re-raise so the proxy's failure handling keeps the provider status and - owns logging (no success-bill). If the client already went away, no - failure hook runs, so bill the partial instead of dropping the request. + owns logging (no success-bill). If the client already went away, or + disconnects before ever consuming the queued exception, no failure hook + runs, so bill the partial instead of dropping the request. """ from litellm._logging import verbose_proxy_logger if not client_detached.is_set() and await self._enqueue_for_client(queue, client_detached, exc): - return + await client_detached.wait() + if not _exception_left_unconsumed(queue, exc): + return verbose_proxy_logger.warning( "async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)", len(collected_chunks), diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 2135397302e..e8099a4217b 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -601,6 +601,43 @@ async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_ assert iterator.logged_chunks == received +@pytest.mark.asyncio +async def test_async_sse_wrapper_salvages_spend_when_queued_error_is_never_consumed(): + """ + When the upstream errors while the client is still connected, the pump + forwards the exception through the queue expecting the relay to re-raise it + into the proxy's failure handling. If the client disconnects before + consuming that queued exception, the handoff never happens and no failure + hook runs, so the pump must notice the unconsumed exception at teardown and + salvage partial spend instead of dropping the row entirely. + """ + upstream_errored = asyncio.Event() + + async def _failing_stream(): + yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}} + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}} + upstream_errored.set() + raise _ProviderStreamError("mid-stream failure", status_code=500) + + iterator = _RecordingLoggingIterator( + litellm_logging_obj=_make_logging_obj("test_salvage_on_unconsumed_queued_error"), + request_body={}, + ) + + gen = iterator.async_sse_wrapper(_failing_stream()) + received = [await gen.__anext__(), await gen.__anext__()] + await upstream_errored.wait() # exception is now queued behind the consumed chunks + await gen.aclose() # client disconnects without ever consuming the queued exception + + for _ in range(100): + if iterator.logged_chunks: + break + await asyncio.sleep(0.01) + + assert iterator.logging_call_count == 1 + assert iterator.logged_chunks == received + + @pytest.mark.asyncio async def test_async_sse_wrapper_applies_backpressure_to_slow_client(monkeypatch): """ diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 1e09afd6919..8d07d38b1b6 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -3066,17 +3066,21 @@ def test_bedrock_invoke_messages_allows_converted_websearch_function_tool(): async def test_bedrock_sse_wrapper_dispatches_logging_on_client_disconnect(): """ Regression test for LIT-5839: closing the outer bedrock_sse_wrapper - mid-stream (what the proxy does on a client disconnect) must close the - inner async_sse_wrapper deterministically so the partial-stream logging - fires. `completion_start_time` is only stamped on the logging object by - that dispatch, so it observing a value proves the whole chain ran. + mid-stream (what the proxy does on a client disconnect) must not lose the + stream's spend logging. Since the detached-pump relay, the upstream read + survives the disconnect and billing fires once the provider stream ends, + so the dispatch is awaited after releasing the upstream instead of being + observed synchronously at aclose(). `completion_start_time` is only + stamped on the logging object by that dispatch, so it observing a value + proves the whole chain ran. """ cfg = AmazonAnthropicClaudeMessagesConfig() + release_upstream = asyncio.Event() - async def _hanging_stream(): + async def _gated_stream(): yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 25, "output_tokens": 1}}} yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}} - await asyncio.Event().wait() + await release_upstream.wait() logging_obj = LiteLLMLoggingObj( model="bedrock/invoke/anthropic.claude-3-sonnet-20240229-v1:0", @@ -3087,11 +3091,16 @@ async def test_bedrock_sse_wrapper_dispatches_logging_on_client_disconnect(): litellm_call_id="test_bedrock_sse_wrapper_disconnect_logging", function_id="test_bedrock_sse_wrapper_disconnect_logging", ) - wrapped = cfg.bedrock_sse_wrapper(_hanging_stream(), litellm_logging_obj=logging_obj, request_body={}) + wrapped = cfg.bedrock_sse_wrapper(_gated_stream(), litellm_logging_obj=logging_obj, request_body={}) await wrapped.__anext__() await wrapped.__anext__() assert logging_obj.completion_start_time is None await wrapped.aclose() + release_upstream.set() + for _ in range(500): + if logging_obj.completion_start_time is not None: + break + await asyncio.sleep(0.01) assert logging_obj.completion_start_time is not None From 25c8d58400a48d1010e73d98987dd44b035764b5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:09:39 -0700 Subject: [PATCH 094/120] fix(docker): install file and make in wolfi builders so the uvloop sdist can build --- Dockerfile | 2 ++ docker/Dockerfile.database | 2 ++ docker/Dockerfile.non_root | 2 ++ 3 files changed, 6 insertions(+) diff --git a/Dockerfile b/Dockerfile index 700b0d6525e..4688a733331 100644 --- a/Dockerfile +++ b/Dockerfile @@ -39,7 +39,9 @@ COPY --from=uvbin /uvx /usr/local/bin/uvx RUN apk add --no-cache \ bash \ + file \ gcc \ + make \ python3 \ python3-dev \ rust \ diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index f0d6d02fccf..9a243c68dca 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -38,7 +38,9 @@ COPY --from=uvbin /uvx /usr/local/bin/uvx RUN apk add --no-cache \ bash \ + file \ gcc \ + make \ python3 \ python3-dev \ openssl \ diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 4a5df6ecd69..2cf70f5b01d 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -39,7 +39,9 @@ RUN for i in 1 2 3; do \ apk add --no-cache \ python3 \ python3-dev \ + file \ gcc \ + make \ rust \ bash \ coreutils \ From ed5ee51dd225fc268166fcfd233a96a0daffbea2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:16:40 -0700 Subject: [PATCH 095/120] fix(passthrough): map sync streaming errors, keep router streaming responses unwrapped, and resolve gigachat from api base - sync llm_passthrough_route: read and close an error-status streaming response before mapping it, so upstream 4xx/5xx surface as the provider error instead of httpx.ResponseNotRead - AsyncPassthroughStreamingResponse: expose aiter_bytes() and carry _hidden_params so the router attaches headers in place instead of wrapping the stream in HiddenParamsAsyncIteratorWrapper, which 500'd every streaming azure router-model passthrough request - logging: swap the passthrough httpx result for the transformed ModelResponse/EmbeddingResponse when firing success callbacks - get_llm_provider: resolve gigachat from its api base and drop the dead gigachat_models elif branch - constants: register the gigachat api base in openai_compatible_endpoints --- litellm/constants.py | 1 + .../get_llm_provider_logic.py | 2 - litellm/litellm_core_utils/litellm_logging.py | 2 +- litellm/passthrough/main.py | 17 +++- .../test_get_llm_provider_endpoint_match.py | 23 ++++++ .../test_litellm_logging.py | 61 ++++++++++++++ .../passthrough/test_passthrough_main.py | 81 +++++++++++++++++++ .../test_llm_pass_through_endpoints.py | 73 +++++++++++++++++ 8 files changed, 256 insertions(+), 4 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index cc6db6c10cc..0f1fa2ee3a7 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -806,6 +806,7 @@ openai_compatible_endpoints: Final[list] = [ "https://api.meta.ai/v1", "https://api.cognition.ai/v1", "https://api.scx.ai/v1", + "https://gigachat.devices.sberbank.ru/api/v1", ] diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index b35ef659d67..9b53b79bbe6 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -496,8 +496,6 @@ def get_llm_provider( custom_llm_provider = "amazon_nova" elif model.startswith("sap/"): custom_llm_provider = "sap" - elif model in litellm.gigachat_models or model.startswith("gigachat/"): - custom_llm_provider = "gigachat" # Last resort for an otherwise-unknown model: a declarative # fallback-generalization routing rule (e.g. routes future claude-* to anthropic). diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e34c647efc2..350e5403e4c 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2141,7 +2141,7 @@ class Logging(LiteLLMLoggingBaseClass): logging_result: Final = self.normalize_logging_result(result=result) - if isinstance(result, Response) and isinstance(logging_result, ModelResponse): + if isinstance(result, Response) and isinstance(logging_result, (ModelResponse, EmbeddingResponse)): result = logging_result if standard_logging_object is None and result is not None and self.stream is not True: diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 2780f510a76..9095cee15a9 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -54,6 +54,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): self._raw_bytes: list[bytes] = [] # mutable-ok: instance buffer for streaming chunks self._flush_scheduled = False self._background_tasks: set[asyncio.Task] = set() # mutable-ok: instance set for background task tracking + self._hidden_params: dict[str, object] = {} # mutable-ok: router attaches response headers here in place @property def status_code(self) -> int: @@ -127,6 +128,9 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): def __aiter__(self) -> AsyncPassthroughStreamingResponse: return self + def aiter_bytes(self) -> AsyncPassthroughStreamingResponse: + return self + async def __anext__(self) -> bytes: if not self._initialized: await self # pyright: ignore[reportGeneralTypeIssues] # structural type check misses __await__ @@ -556,7 +560,18 @@ def llm_passthrough_route( else: # Sync path - client.client.send returns Response directly response: httpx.Response = client.client.send(request=request, stream=is_streaming_request) - response.raise_for_status() + try: + response.raise_for_status() + except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic + try: + response.read() + except Exception: # noqa: BLE001 S110 # Safe catch-all for cleanup logic + pass + try: + response.close() + except Exception: # noqa: BLE001 S110 # Safe catch-all for cleanup logic + pass + raise if hasattr(response, "iter_bytes") and is_streaming_request: return PassthroughStreamingResponse(response, litellm_logging_obj, provider_config) diff --git a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py index 6cacd119030..419ca104bb1 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py +++ b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py @@ -184,3 +184,26 @@ class TestTogetherApiBaseResolvesProvider: assert provider == "together_ai" assert api_base == "https://api.together.ai/v1" + + +class TestGigachatApiBaseResolvesProvider: + """ + Regression for the GigaChat api_base branch: the provider-mapping chain + carried an ``endpoint == "https://gigachat.devices.sberbank.ru/api/v1"`` + elif, but the URL was never added to ``openai_compatible_endpoints``, so + the endpoint loop never fired the branch and a caller-supplied GigaChat + api_base raised BadRequestError instead of resolving to ``gigachat``. + """ + + def test_gigachat_api_base_resolves_to_gigachat(self, monkeypatch): + monkeypatch.setenv("GIGACHAT_API_KEY", "gigachat-key-from-env") + + model, provider, dynamic_api_key, returned_api_base = get_llm_provider( + model="GigaChat-2", + api_base="https://gigachat.devices.sberbank.ru/api/v1", + ) + + assert provider == "gigachat" + assert dynamic_api_key == "gigachat-key-from-env" + assert returned_api_base == "https://gigachat.devices.sberbank.ru/api/v1" + assert model == "GigaChat-2" diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 947a55410ef..c7328adb0b3 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -6101,3 +6101,64 @@ def test_response_timing_metrics_survive_deepcopy(logging_obj): logging_obj.set_response_timing_metrics({"_response_ms": 12.5}) assert copy.deepcopy(logging_obj).response_timing_metrics == {"_response_ms": 12.5} + + +def test_passthrough_embeddings_result_swapped_for_callbacks(): + """ + Regression: for gigachat passthrough /embeddings, normalize_logging_result + produces an EmbeddingResponse, but the result swap only accepted + ModelResponse, so callbacks kept receiving the raw httpx.Response (which + crashes attribute readers like OTEL). The swap must cover + EmbeddingResponse too. + """ + import datetime as dt + + from litellm.types.utils import EmbeddingResponse + + logging_obj = LitellmLogging( + model="EmbeddingsGigaR", + messages=[], + stream=False, + call_type="allm_passthrough_route", + start_time=time.time(), + litellm_call_id="passthrough-embed-call-id", + function_id="passthrough-embed-fn-id", + ) + logging_obj.update_environment_variables( + litellm_params={}, + optional_params={}, + model="EmbeddingsGigaR", + custom_llm_provider="gigachat", + endpoint="/embeddings", + request_data={"model": "EmbeddingsGigaR", "input": ["hello"]}, + input=["hello"], + ) + + httpx_response = httpx.Response( + 200, + json={ + "object": "list", + "data": [ + { + "object": "embedding", + "embedding": [0.1, 0.2, 0.3], + "index": 0, + "usage": {"prompt_tokens": 5}, + } + ], + "model": "EmbeddingsGigaR", + }, + request=httpx.Request( + "POST", "https://gigachat.devices.sberbank.ru/api/v1/embeddings" + ), + ) + + _, _, swapped_result = logging_obj._success_handler_helper_fn( + result=httpx_response, + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + cache_hit=False, + ) + + assert isinstance(swapped_result, EmbeddingResponse) + assert swapped_result.data[0]["embedding"] == [0.1, 0.2, 0.3] diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/test_litellm/passthrough/test_passthrough_main.py index fb67aa6f3e4..1950c37a12e 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/test_litellm/passthrough/test_passthrough_main.py @@ -719,6 +719,87 @@ async def test_allm_passthrough_route_429_streaming_raises(): assert exc_info.value.response.status_code == 429 +def test_llm_passthrough_route_sync_streaming_error_maps_upstream_status(): + """ + Regression test: a sync streaming passthrough whose upstream answers an + error status must surface the mapped provider error, not + httpx.ResponseNotRead. + + Before the fix, raise_for_status() raised on the still-unread streamed + response, and _handle_error then touched e.response.text, which raises + ResponseNotRead on a streamed-but-unread body, masking the real upstream + error entirely. + """ + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + error_body = json.dumps( + { + "error": { + "code": "429", + "message": "Rate limit exceeded. Retry after 10 seconds.", + } + } + ).encode() + + class _UnreadErrorStream(httpx.SyncByteStream): + def __iter__(self): + yield error_body + + def _handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 429, + stream=_UnreadErrorStream(), + headers={"content-type": "application/json"}, + ) + + sync_client = HTTPHandler( + client=httpx.Client(transport=httpx.MockTransport(_handler)) + ) + + mock_provider_config = MagicMock() + mock_provider_config.get_complete_url.return_value = ( + httpx.URL("https://gigachat.devices.sberbank.ru/api/v1/chat/completions"), + "https://gigachat.devices.sberbank.ru/api/v1", + ) + mock_provider_config.get_api_key.return_value = "fake-key" + mock_provider_config.validate_environment.return_value = { + "Authorization": "Bearer fake-key" + } + mock_provider_config.sign_request.return_value = ( + {"Authorization": "Bearer fake-key"}, + None, + ) + mock_provider_config.is_streaming_request.return_value = True + mock_provider_config.get_error_class.side_effect = ( + lambda error_message, status_code, headers: BaseLLMException( + status_code=status_code, message=error_message, headers=headers + ) + ) + + mock_logging_obj = MagicMock() + + with pytest.raises(BaseLLMException) as exc_info: + llm_passthrough_route( + model="gigachat/GigaChat-2", + endpoint="chat/completions", + method="POST", + custom_llm_provider="gigachat", + api_base="https://gigachat.devices.sberbank.ru/api/v1", + api_key="fake-key", + json={ + "model": "GigaChat-2", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + }, + client=sync_client, + litellm_logging_obj=mock_logging_obj, + provider_config=mock_provider_config, + ) + + assert exc_info.value.status_code == 429 + assert "Rate limit exceeded" in str(exc_info.value) + + def test_llm_passthrough_route_propagates_allm_passthrough_route_to_logging_obj(): """ Regression guard for LIT-4192: `allm_passthrough_route` sets 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 225e1b3d998..303c71a1630 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 @@ -4865,3 +4865,76 @@ class TestPassthroughRouterModelBudgetReservation: ) self._assert_metadata_carries_attribution(captured, user_api_key_dict) + + +class TestAzureRouterModelStreamingDispatch: + """ + Regression: ``llm_router.allm_passthrough_route`` returns an awaited + ``AsyncPassthroughStreamingResponse`` for streaming calls, which is no + longer an async generator under ``inspect.isasyncgen``. The dispatch's + else branch therefore calls ``.aiter_bytes()`` / ``.status_code`` / + ``.headers`` on it. The router's ``set_response_headers`` also runs the + result through ``prepare_response_for_header_attachment``, which used to + wrap it in ``HiddenParamsAsyncIteratorWrapper`` (no ``aiter_bytes``), so + every streaming Azure router-model request 500'd with + ``AttributeError: aiter_bytes``; ``_hidden_params`` on the streaming + response keeps it unwrapped. + """ + + @pytest.mark.asyncio + async def test_azure_router_model_streaming_returns_streaming_response(self, monkeypatch): + import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.passthrough.main import AsyncPassthroughStreamingResponse + + upstream_body = b"data: hello\n\n" + + async def _upstream_response() -> httpx.Response: + upstream_request = httpx.Request( + "POST", + "https://my-azure.openai.azure.com/openai/deployments/gpt-5/chat/completions", + ) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=upstream_body, + request=upstream_request, + ) + + logging_obj = MagicMock() + logging_obj.async_flush_passthrough_collected_chunks = AsyncMock() + + from litellm.router_utils.add_retry_fallback_headers import prepare_response_for_header_attachment + + class StreamingRouter: + async def allm_passthrough_route(self, **kwargs): + streaming_response = await AsyncPassthroughStreamingResponse( + response=_upstream_response(), + litellm_logging_obj=logging_obj, + provider_config=MagicMock(), + ) + return prepare_response_for_header_attachment(streaming_response) + + async def fake_get_request_body(_request): + return {"model": "gpt-5", "stream": True} + + monkeypatch.setattr(proxy_server, "llm_router", StreamingRouter()) + monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) + monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + + request = MagicMock(spec=Request) + request.method = "POST" + request.headers = {"content-type": "application/json"} + request.query_params = {} + + result = await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=request, + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + + assert isinstance(result, StreamingResponse) + assert result.status_code == 200 + body = b"".join([chunk async for chunk in result.body_iterator]) + assert body == upstream_body From 15aa51a88ae15c46696d687fdda25d8a49567b49 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:24:38 -0700 Subject: [PATCH 096/120] ci(osv): ignore GHSA-h7x2-h6g9-p789 until mlflow ships a fixed release --- osv-scanner.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/osv-scanner.toml b/osv-scanner.toml index 7ab450945f5..205e9a39358 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -2,3 +2,8 @@ id = "GHSA-w8v5-vhqr-4h9v" ignoreUntil = 2026-09-09 reason = "diskcache has no fixed release published; remove this entry once one exists" + +[[IgnoredVulns]] +id = "GHSA-h7x2-h6g9-p789" +ignoreUntil = 2026-09-14 +reason = "mlflow 3.15.0 has no fixed release published; remove this entry once one exists" From a90fb538bfd0a71e39d82b344c5c666200d01a63 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:25:51 -0700 Subject: [PATCH 097/120] fix(friendli): declare GLM-5.3-Flash reasoning efforts as explicit levels --- model_prices_and_context_window.json | 7 +++++-- .../test_friendli_glm_5_3_flash_model_metadata.py | 3 +-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index fd6b9e11adf..f0e95d77018 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19575,8 +19575,11 @@ "output_cost_per_token": 5e-07, "cache_read_input_token_cost": 3e-08, "supports_prompt_caching": true, - "supports_max_reasoning_effort": true, - "supports_low_reasoning_effort": true, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], "supports_parallel_function_calling": true, "supports_response_schema": true, "supports_native_structured_output": true, diff --git a/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py b/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py index 0acd750e49a..7e94205fb09 100644 --- a/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py +++ b/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py @@ -23,8 +23,7 @@ def test_friendli_glm_5_3_flash_model_info(): assert info["max_output_tokens"] == 1048576 assert info["supports_function_calling"] is True assert info["supports_reasoning"] is True - assert info["supports_low_reasoning_effort"] is True - assert info["supports_max_reasoning_effort"] is True + assert info["reasoning_effort_levels"] == ["low", "high", "max"] assert info["supports_tool_choice"] is True assert info["supports_prompt_caching"] is True assert info["supports_vision"] is True From bd794f9f18c7422824a36e7d02086ebea25e8ec2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:26:37 -0700 Subject: [PATCH 098/120] fix(friendli): track GLM-5.3 discounted live pricing and declare effort levels --- model_prices_and_context_window.json | 13 ++++++++----- .../test_friendli_glm_5_3_model_metadata.py | 9 ++++----- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 24aad7c90c8..f70a0ac63d9 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19571,12 +19571,15 @@ "max_input_tokens": 1048576, "max_tokens": 1048576, "max_output_tokens": 1048576, - "input_cost_per_token": 1.4e-06, - "output_cost_per_token": 4.4e-06, - "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.26e-06, + "output_cost_per_token": 3.96e-06, + "cache_read_input_token_cost": 2.34e-07, "supports_prompt_caching": true, - "supports_max_reasoning_effort": true, - "supports_low_reasoning_effort": true, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], "supports_parallel_function_calling": true, "supports_response_schema": true, "supports_native_structured_output": true, diff --git a/tests/test_litellm/test_friendli_glm_5_3_model_metadata.py b/tests/test_litellm/test_friendli_glm_5_3_model_metadata.py index 5e7a1ede699..5282b0f589e 100644 --- a/tests/test_litellm/test_friendli_glm_5_3_model_metadata.py +++ b/tests/test_litellm/test_friendli_glm_5_3_model_metadata.py @@ -16,15 +16,14 @@ def test_friendli_glm_5_3_model_info(): ), f"{model} not found in model_prices_and_context_window.json" assert info["litellm_provider"] == "friendliai" assert info["mode"] == "chat" - assert info["input_cost_per_token"] == 1.4e-06 - assert info["output_cost_per_token"] == 4.4e-06 - assert info["cache_read_input_token_cost"] == 2.6e-07 + assert info["input_cost_per_token"] == 1.26e-06 + assert info["output_cost_per_token"] == 3.96e-06 + assert info["cache_read_input_token_cost"] == 2.34e-07 assert info["max_input_tokens"] == 1048576 assert info["max_output_tokens"] == 1048576 assert info["supports_function_calling"] is True assert info["supports_reasoning"] is True - assert info["supports_low_reasoning_effort"] is True - assert info["supports_max_reasoning_effort"] is True + assert info["reasoning_effort_levels"] == ["low", "high", "max"] assert info["supports_tool_choice"] is True assert info["supports_prompt_caching"] is True assert info["supports_vision"] is False From ebdb54d4f0c51c7c8df038dd5b189f4e84577203 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:32:19 -0700 Subject: [PATCH 099/120] fix(docker): keep image venvs on the apk python and bump the wolfi digest Since the requires-python cap moved to <3.15, uv resolved the project python to 3.14, downloaded a managed interpreter under /root/.local/share/uv that the runtime stage never receives, and every layer-cache-miss image build broke: first at uvloop's cp314 sdist configure step, then, with file/make added, at the runtime stage where the copied venv's python symlink dangles and prisma imports fall through to the system python. UV_PYTHON_DOWNLOADS=0 (already the convention in migrations/backend/gateway) roots the venv on the apk python3. The wolfi-base digest bump is required alongside it: the pinned 08-22 base ships glibc-2.43 while the current apk python-3.13 needs GLIBC_2.44, and wolfi version-names glibc packages so apk upgrade cannot cross that boundary. With the venv on system 3.13 every dependency installs from wheels again, so the file and make packages added for the sdist build are reverted. --- Dockerfile | 11 +++++++---- backend/Dockerfile | 4 ++-- docker/Dockerfile.database | 11 +++++++---- docker/Dockerfile.non_root | 11 +++++++---- gateway/Dockerfile | 4 ++-- migrations/Dockerfile | 4 ++-- 6 files changed, 27 insertions(+), 18 deletions(-) diff --git a/Dockerfile b/Dockerfile index 4688a733331..675a79b0686 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,10 +1,10 @@ # syntax=docker/dockerfile:1.7 # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:57108e597a8cf3bd376b810f1c3539c21942daefa242cb9dddaae30f8aac735d # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:57108e597a8cf3bd376b810f1c3539c21942daefa242cb9dddaae30f8aac735d ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 @@ -39,9 +39,7 @@ COPY --from=uvbin /uvx /usr/local/bin/uvx RUN apk add --no-cache \ bash \ - file \ gcc \ - make \ python3 \ python3-dev \ rust \ @@ -51,8 +49,13 @@ RUN apk add --no-cache \ npm \ libsndfile +# UV_PYTHON_DOWNLOADS=0 keeps the venv on the apk python3 above. Without it, +# uv resolves requires-python to the newest allowed minor, downloads a managed +# interpreter under /root/.local/share/uv that the runtime stage never +# receives, and the copied venv's python symlink dangles at runtime. ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ UV_LINK_MODE=copy \ + UV_PYTHON_DOWNLOADS=0 \ PATH="/app/.venv/bin:${PATH}" # Copy dependency metadata first for layer caching diff --git a/backend/Dockerfile b/backend/Dockerfile index 4ca40944606..8aea8312df9 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:57108e597a8cf3bd376b810f1c3539c21942daefa242cb9dddaae30f8aac735d +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:57108e597a8cf3bd376b810f1c3539c21942daefa242cb9dddaae30f8aac735d ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 9a243c68dca..4243eea5796 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -1,10 +1,10 @@ # syntax=docker/dockerfile:1.7 # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:57108e597a8cf3bd376b810f1c3539c21942daefa242cb9dddaae30f8aac735d # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:57108e597a8cf3bd376b810f1c3539c21942daefa242cb9dddaae30f8aac735d ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 @@ -38,9 +38,7 @@ COPY --from=uvbin /uvx /usr/local/bin/uvx RUN apk add --no-cache \ bash \ - file \ gcc \ - make \ python3 \ python3-dev \ openssl \ @@ -49,8 +47,13 @@ RUN apk add --no-cache \ npm \ libsndfile +# UV_PYTHON_DOWNLOADS=0 keeps the venv on the apk python3 above. Without it, +# uv resolves requires-python to the newest allowed minor, downloads a managed +# interpreter under /root/.local/share/uv that the runtime stage never +# receives, and the copied venv's python symlink dangles at runtime. ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ UV_LINK_MODE=copy \ + UV_PYTHON_DOWNLOADS=0 \ PATH="/app/.venv/bin:${PATH}" # Copy dependency metadata first for layer caching diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 2cf70f5b01d..19ef97a4f46 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -1,8 +1,8 @@ # syntax=docker/dockerfile:1.7 # Base images -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:57108e597a8cf3bd376b810f1c3539c21942daefa242cb9dddaae30f8aac735d +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:57108e597a8cf3bd376b810f1c3539c21942daefa242cb9dddaae30f8aac735d ARG PROXY_EXTRAS_SOURCE=published ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. @@ -39,9 +39,7 @@ RUN for i in 1 2 3; do \ apk add --no-cache \ python3 \ python3-dev \ - file \ gcc \ - make \ rust \ bash \ coreutils \ @@ -52,8 +50,13 @@ RUN for i in 1 2 3; do \ npm && break || sleep 5; \ done +# UV_PYTHON_DOWNLOADS=0 keeps the venv on the apk python3 above. Without it, +# uv resolves requires-python to the newest allowed minor, downloads a managed +# interpreter under /root/.local/share/uv that the runtime stage never +# receives, and the copied venv's python symlink dangles at runtime. ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ UV_LINK_MODE=copy \ + UV_PYTHON_DOWNLOADS=0 \ PATH="/app/.venv/bin:${PATH}" \ LITELLM_NON_ROOT=true \ XDG_CACHE_HOME=/app/.cache diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 4a2e32e186e..235c535f9f1 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:57108e597a8cf3bd376b810f1c3539c21942daefa242cb9dddaae30f8aac735d +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:57108e597a8cf3bd376b810f1c3539c21942daefa242cb9dddaae30f8aac735d ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin diff --git a/migrations/Dockerfile b/migrations/Dockerfile index 6335e6f6bd8..524450024e0 100644 --- a/migrations/Dockerfile +++ b/migrations/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:57108e597a8cf3bd376b810f1c3539c21942daefa242cb9dddaae30f8aac735d +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:57108e597a8cf3bd376b810f1c3539c21942daefa242cb9dddaae30f8aac735d ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin From 2bbf5135a5782a07e62d51d6df1fc2a774ec501e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:55:39 -0700 Subject: [PATCH 100/120] fix(friendli): ship GLM-5.3-Flash in the bundled backup cost map --- ...odel_prices_and_context_window_backup.json | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 05c1cfd3179..f0e95d77018 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -19564,6 +19564,34 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "friendliai/zai-org/GLM-5.3-Flash": { + "litellm_provider": "friendliai", + "supports_reasoning": true, + "supports_function_calling": true, + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "max_output_tokens": 1048576, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 3e-08, + "supports_prompt_caching": true, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "mode": "chat", + "comment": "Native multimodal GLM model for efficient coding and long-horizon agent tasks", + "source": "https://api.friendli.ai/serverless/v1/models", + "supports_vision": true, + "supports_image_input": true, + "supports_video_input": true + }, "ft:babbage-002": { "deprecation_date": "2026-10-23", "input_cost_per_token": 1.6e-06, From 0ac328910266e1edfab6279c6d39312cf9ced0f8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:55:54 -0700 Subject: [PATCH 101/120] fix(friendli): ship GLM-5.3 in the bundled backup cost map --- ...odel_prices_and_context_window_backup.json | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 05c1cfd3179..f70a0ac63d9 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -19564,6 +19564,33 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "friendliai/zai-org/GLM-5.3": { + "litellm_provider": "friendliai", + "supports_reasoning": true, + "supports_function_calling": true, + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "max_output_tokens": 1048576, + "input_cost_per_token": 1.26e-06, + "output_cost_per_token": 3.96e-06, + "cache_read_input_token_cost": 2.34e-07, + "supports_prompt_caching": true, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "mode": "chat", + "comment": "Flagship GLM model for long-horizon coding, agents, and complex project delivery", + "source": "https://api.friendli.ai/serverless/v1/models", + "supports_vision": false, + "supports_image_input": false + }, "ft:babbage-002": { "deprecation_date": "2026-10-23", "input_cost_per_token": 1.6e-06, From db7eb641cbc0a81dcd640350501b357df53abf94 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:58:53 -0700 Subject: [PATCH 102/120] ci(osv): ignore GHSA-h7x2-h6g9-p789 until mlflow ships a fix The advisory was modified 2026-08-31 and flags mlflow 3.13.0 through 3.15.2 with no fixed release published, so every osv-scan run fails with nothing to bump. Same treatment as the existing diskcache entry. --- osv-scanner.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/osv-scanner.toml b/osv-scanner.toml index 7ab450945f5..5b0339bdcd0 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -2,3 +2,8 @@ id = "GHSA-w8v5-vhqr-4h9v" ignoreUntil = 2026-09-09 reason = "diskcache has no fixed release published; remove this entry once one exists" + +[[IgnoredVulns]] +id = "GHSA-h7x2-h6g9-p789" +ignoreUntil = 2026-09-14 +reason = "mlflow has no fixed release published; remove this entry once one exists" From 9c577c6045b01c0002e68d584da69b8bad996ed4 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 14:04:15 -0700 Subject: [PATCH 103/120] test(e2e): assert user-observable behavior instead of DOM structure MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The UI e2e suite had a class of assertions that pin how the dashboard is built rather than what it does, so an ordinary refactor turns them red without any user-visible change. Geometry. The auto-router template select had two tests made of pixel arithmetic plus a data-side="bottom" check, which is Base UI's own positioner signal. The regression they guard (#38554) is a popup opening on top of the control that spawned it, so both cases collapse to one invariant: the options never cover the trigger. It now runs at both viewport heights and reads the popup as role=listbox. The models header test compared the tabs and refresh centers within 2px, which a padding change flips; it now asserts the two share a row. Structure. The logs drawer test walked xpath=../../.. from a text node and read collapsed state off chevron icon classes. SectionHeader now renders a real disclosure button with aria-expanded, and its two copy buttons carry distinct names instead of both being "Copy". Sidebar group toggles expose aria-expanded too, so the migration spec can ask for a collapsed group by state rather than by nesting depth. Positional lookups. keyRow.locator("button").first(), row.locator("td") .first() and getByTestId(grid).locator("div").first() all named a position where they meant an action; they now name the control. Table scoping moves from "table tbody" to role=row. Timing. Nine waitForTimeout calls are gone. Every assertion that followed them already retried to its own timeout, so the sleeps only slowed the run down. Both files under tests/users/ were wrapped in test.skip("...", () => {}), which registers one skipped test and never runs the body, so the four tests inside had never executed and were written against a UI that has since changed (the search placeholder is "Search by email…", the ID filters moved into a drawer, pagination is labelled "Go to previous page"). Rewritten against the current surface: the suite goes from 104 collected tests to 107. Left in place deliberately: the chip and dialog-footer data-slot selectors, because the accessible names they work around live in components/ui/, which is shadcn CLI-managed and not hand-edited. --- .../tests/internal-user/internalUser.spec.ts | 7 +- .../internal-user/internalUserNoTeam.spec.ts | 7 +- .../internalUserWithTeams.spec.ts | 9 +- .../internal-viewer/internalViewer.spec.ts | 4 +- tests/e2e/ui/tests/logs/logs.spec.ts | 38 +++--- tests/e2e/ui/tests/mcp/mcpServerEdit.spec.ts | 2 +- .../ui/tests/migration/migratedPages.spec.ts | 10 +- .../e2e/ui/tests/modelsPage/addModel.spec.ts | 29 ++--- .../autoRouterTemplateSelect.spec.ts | 53 +++------ .../tests/modelsPage/responsiveHeader.spec.ts | 9 +- tests/e2e/ui/tests/proxy-admin/keys.spec.ts | 15 ++- .../e2e/ui/tests/team-admin/teamAdmin.spec.ts | 2 +- tests/e2e/ui/tests/usage/usagePage.spec.ts | 7 +- tests/e2e/ui/tests/users/searchUsers.spec.ts | 110 ++++++------------ .../ui/tests/users/viewInternalUsers.spec.ts | 59 +++------- .../src/components/leftnav.test.tsx | 14 +++ .../src/components/leftnav.tsx | 1 + .../LogDetailsDrawer/SectionHeader.test.tsx | 18 +++ .../LogDetailsDrawer/SectionHeader.tsx | 73 +++++++----- 19 files changed, 198 insertions(+), 269 deletions(-) diff --git a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts index b8424b06115..26e34dd2fe5 100644 --- a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts @@ -22,8 +22,7 @@ test.describe("Internal User", () => { const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); await page.keyboard.type(E2E_TEAM_CRUD_ALIAS); - const dropdown = page.locator('[data-slot="combobox-content"]:visible'); - await expect(dropdown.getByText(E2E_TEAM_CRUD_ALIAS).first()).toBeVisible({ timeout: 5_000 }); + await expect(page.getByRole("option", { name: E2E_TEAM_CRUD_ALIAS }).first()).toBeVisible({ timeout: 5_000 }); }); test("Team info page omits the Settings tab for non-admin members", async ({ page }) => { @@ -43,12 +42,12 @@ test.describe("Internal User", () => { // Anchor on the user's own seeded key so the absence check below cannot // pass vacuously against an empty table. - await expect(page.locator("table tbody").getByText(E2E_INTERNAL_USER_KEY_ALIAS).first()).toBeVisible({ + await expect(page.getByRole("row").filter({ hasText: E2E_INTERNAL_USER_KEY_ALIAS }).first()).toBeVisible({ timeout: 10_000, }); // The litellm-dashboard team is the proxy's internal bookkeeping team — // its keys must never leak into an internal user's Virtual Keys table. - await expect(page.locator("table tbody").getByText("litellm-dashboard")).toHaveCount(0); + await expect(page.getByRole("row").filter({ hasText: "litellm-dashboard" })).toHaveCount(0); }); }); diff --git a/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts index c44305187f1..653e096b713 100644 --- a/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts @@ -30,16 +30,13 @@ test.describe("Internal User with no team memberships", () => { const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); - const dropdown = page.locator('[data-slot="combobox-content"]:visible').first(); - await expect(dropdown).toBeVisible({ timeout: 5_000 }); - // Wait for the settled-empty state, not a transient one. The dropdown shows // "Loading teams…" while teams load and only swaps in "No teams found" once // the request resolves with nothing (team_dropdown.tsx passes both copies to // PaginatedSearchSelect). Asserting on it means a regression where teams DO // load for this user fails here instead of racing a one-shot count() against // an in-flight request. - await expect(dropdown.getByText("No teams found")).toBeVisible({ timeout: 10_000 }); - await expect(dropdown.getByRole("option")).toHaveCount(0); + await expect(page.getByText("No teams found")).toBeVisible({ timeout: 10_000 }); + await expect(page.getByRole("option")).toHaveCount(0); }); }); diff --git a/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts index 68319154554..49e27a36673 100644 --- a/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts @@ -21,13 +21,10 @@ test.describe("Internal User with team memberships", () => { const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); - const dropdown = page.locator('[data-slot="combobox-content"]:visible').first(); - await expect(dropdown).toBeVisible({ timeout: 5_000 }); - // Both seeded memberships render, and nothing else does — proving the // dropdown is scoped to the user's teams rather than empty or unfiltered. - await expect(dropdown.getByText(E2E_TEAM_CRUD_ALIAS, { exact: true })).toBeVisible({ timeout: 10_000 }); - await expect(dropdown.getByText(E2E_TEAM_ORG_ALIAS, { exact: true })).toBeVisible(); - await expect(dropdown.getByRole("option")).toHaveCount(2); + await expect(page.getByRole("option", { name: E2E_TEAM_CRUD_ALIAS })).toBeVisible({ timeout: 10_000 }); + await expect(page.getByRole("option", { name: E2E_TEAM_ORG_ALIAS })).toBeVisible(); + await expect(page.getByRole("option")).toHaveCount(2); }); }); diff --git a/tests/e2e/ui/tests/internal-viewer/internalViewer.spec.ts b/tests/e2e/ui/tests/internal-viewer/internalViewer.spec.ts index 4de86c46398..dd40976341d 100644 --- a/tests/e2e/ui/tests/internal-viewer/internalViewer.spec.ts +++ b/tests/e2e/ui/tests/internal-viewer/internalViewer.spec.ts @@ -59,9 +59,9 @@ test.describe("Internal Viewer", () => { await expect(page.getByRole("button", { name: /Create New Key/i })).toHaveCount(0); // Open the viewer's own key info page - const keyRow = page.locator("tr", { hasText: E2E_VIEWER_KEY_ALIAS }); + const keyRow = page.getByRole("row").filter({ hasText: E2E_VIEWER_KEY_ALIAS }); await expect(keyRow).toBeVisible({ timeout: 10_000 }); - await keyRow.locator("button").first().click(); + await keyRow.getByRole("button", { name: E2E_VIEWER_KEY_ALIAS }).click(); await expect(page.getByText("Back to Keys")).toBeVisible({ timeout: 10_000 }); // None of the destructive / mutating actions should render diff --git a/tests/e2e/ui/tests/logs/logs.spec.ts b/tests/e2e/ui/tests/logs/logs.spec.ts index 56a3d0f0109..610a88c6cd0 100644 --- a/tests/e2e/ui/tests/logs/logs.spec.ts +++ b/tests/e2e/ui/tests/logs/logs.spec.ts @@ -11,12 +11,11 @@ import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, sendChatCompletion, waitForSpendLog } const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; -/** - * Walking up from the label is the only stable handle: the header carries no role, test id or class, - * and its copy button is icon-only with a hover-only tooltip. - */ -const sectionHeader = (drawer: Locator, label: "Input" | "Output"): Locator => - drawer.getByText(label, { exact: true }).locator("xpath=../../.."); +const sectionToggle = (drawer: Locator, label: "Input" | "Output"): Locator => + drawer.getByRole("button", { name: new RegExp(`^${label}\\b`) }); + +const sectionCopy = (drawer: Locator, label: "Input" | "Output"): Locator => + drawer.getByRole("button", { name: `Copy ${label.toLowerCase()}` }); /** Every tab stays mounted, so the DOM holds four tables at once; scope to the visible one. */ const requestLogsRows = (page: PlaywrightPage): Locator => @@ -95,14 +94,14 @@ test.describe("Logs page", () => { await expect(drawer).toBeVisible({ timeout: 20_000 }); // Copy request: the Input card's copy button puts the prompt on the clipboard. - await sectionHeader(drawer, "Input").getByRole("button").click(); + await sectionCopy(drawer, "Input").click(); await expect(page.getByText("Input copied")).toBeVisible({ timeout: 10_000, }); expect(await page.evaluate(() => navigator.clipboard.readText())).toContain(prompt); // Copy response: the Output card's copy button puts the completion on it. - await sectionHeader(drawer, "Output").getByRole("button").click(); + await sectionCopy(drawer, "Output").click(); await expect(page.getByText("Output copied")).toBeVisible({ timeout: 10_000, }); @@ -125,24 +124,15 @@ test.describe("Logs page", () => { timeout: 20_000, }); - // The body collapses via `max-height: 0; overflow: hidden`, which zeroes its own bounding - // box, so the wrapper reads as hidden while the clipped text node inside it does not. - const header = sectionHeader(drawer, "Input"); - const body = header.locator("xpath=following-sibling::div[1]"); - await expect(header.locator(".lucide-chevron-up")).toBeVisible(); - await expect(body).toBeVisible(); + const toggle = sectionToggle(drawer, "Input"); + await expect(toggle).toHaveAttribute("aria-expanded", "true"); + await expect(drawer.getByText(prompt, { exact: false })).toBeVisible(); - await header.click(); - await expect(header.locator(".lucide-chevron-down")).toBeVisible({ - timeout: 10_000, - }); - await expect(body).toBeHidden({ timeout: 10_000 }); + await toggle.click(); + await expect(toggle).toHaveAttribute("aria-expanded", "false", { timeout: 10_000 }); - await header.click(); - await expect(header.locator(".lucide-chevron-up")).toBeVisible({ - timeout: 10_000, - }); - await expect(body).toBeVisible({ timeout: 10_000 }); + await toggle.click(); + await expect(toggle).toHaveAttribute("aria-expanded", "true", { timeout: 10_000 }); await expect(drawer.getByText(prompt, { exact: false })).toBeVisible({ timeout: 10_000, }); diff --git a/tests/e2e/ui/tests/mcp/mcpServerEdit.spec.ts b/tests/e2e/ui/tests/mcp/mcpServerEdit.spec.ts index 46799c8a18f..aa7cdf82498 100644 --- a/tests/e2e/ui/tests/mcp/mcpServerEdit.spec.ts +++ b/tests/e2e/ui/tests/mcp/mcpServerEdit.spec.ts @@ -73,7 +73,7 @@ test.describe("MCP Servers - edit and delete", () => { test("Deleting a server removes it", async ({ page }) => { expect(await findServerByName(page, serverName), `created server ${serverName} exists`).toBeTruthy(); - const card = page.getByTestId("mcp-servers-grid").locator("div").filter({ hasText: serverName }).first(); + const card = page.getByTestId("mcp-servers-grid").getByRole("button", { name: serverName }); await card.getByRole("button", { name: "Server actions" }).click(); await page.getByRole("menuitem", { name: "Delete" }).click(); diff --git a/tests/e2e/ui/tests/migration/migratedPages.spec.ts b/tests/e2e/ui/tests/migration/migratedPages.spec.ts index 3ad4b217d08..473d0b795f1 100644 --- a/tests/e2e/ui/tests/migration/migratedPages.spec.ts +++ b/tests/e2e/ui/tests/migration/migratedPages.spec.ts @@ -36,16 +36,10 @@ async function expectRendered(page: Page) { async function clickSidebar(page: Page, segment: string) { const link = sidebar(page).locator(`a[href$="/ui/${segment}"]`).first(); for (let i = 0; i < 8 && !(await link.isVisible().catch(() => false)); i++) { - // A collapsed group is a menu item with a group-toggle button but no - // rendered submenu yet; clicking the toggle expands it. - const collapsedGroup = sidebar(page) - .locator( - '[data-slot="sidebar-menu-item"]:has(> [data-slot="sidebar-menu-button"]):not(:has(> [data-slot="sidebar-menu-sub"])) > [data-slot="sidebar-menu-button"]', - ) - .first(); + const collapsedGroup = sidebar(page).getByRole("button", { expanded: false }).first(); if (!(await collapsedGroup.isVisible().catch(() => false))) break; await collapsedGroup.click(); - await page.waitForTimeout(250); + await expect(collapsedGroup).toHaveAttribute("aria-expanded", "true"); } await link.click(); } diff --git a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts index dad716b4c83..84d1c01b452 100644 --- a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts @@ -188,7 +188,7 @@ test.describe("Add Model", () => { await expect(resultsModal).toBeHidden({ timeout: 5_000 }); const created = await captureRequestBody(page, { method: "POST", urlIncludes: "/model/new" }, async () => { - await page.getByRole("button", { name: "Add Model" }).last().click(); + await page.getByTestId("add-model-btn").click(); }); expect(created.model_name, "the model is created under the name that was typed").toBe(publicName); expect(created.litellm_params?.api_base, "the api base survives the form").toBe(MOCK_LLM_BASE); @@ -254,7 +254,7 @@ test.describe("Add Model", () => { // Click Add Model button by its text const created = await captureRequestBody(page, { method: "POST", urlIncludes: "/model/new" }, async () => { - await page.getByRole("button", { name: "Add Model" }).last().click(); + await page.getByTestId("add-model-btn").click(); }); // The form sends custom_llm_provider separately from the name, so both halves have to arrive. expect(created.model_name, "the selected model is what goes on the wire").toBe("claude-haiku-4-5"); @@ -267,11 +267,9 @@ test.describe("Add Model", () => { // Navigate to All Models tab await page.getByRole("tab", { name: "All Models" }).click(); await page.waitForLoadState("networkidle"); - await page.waitForTimeout(2000); // Search for the model we just added await page.getByPlaceholder("Search model names").fill("claude-haiku-4-5"); - await page.waitForTimeout(1000); // Verify the model appears in the results count (not "Showing 0 results") await expect(page.getByTestId("pagination-range")).toHaveText(/Showing \d+-\d+ of \d+/, { @@ -279,8 +277,9 @@ test.describe("Add Model", () => { }); // Verify the model name appears in the table body - const tableBody = page.locator("table tbody"); - await expect(tableBody.getByText("claude-haiku-4-5").first()).toBeVisible({ timeout: 15_000 }); + await expect(page.getByRole("row").filter({ hasText: "claude-haiku-4-5" })).not.toHaveCount(0, { + timeout: 15_000, + }); // A row proves the name is there, not what the deployment routes to. const stored = await findDeploymentByName(page, "claude-haiku-4-5"); @@ -333,11 +332,11 @@ test.describe("Add Model", () => { const teamDropdown = page.getByTestId("team-dropdown").getByRole("combobox"); await expect(teamDropdown).toBeVisible({ timeout: 5_000 }); await teamDropdown.click(); - const teamOption = page.locator('[data-slot="combobox-content"]:visible').getByText(E2E_TEAM_CRUD_ID).first(); + const teamOption = page.getByRole("option", { name: E2E_TEAM_CRUD_ID }).first(); await expect(teamOption).toBeVisible({ timeout: 5_000 }); await teamOption.click(); - await page.getByRole("button", { name: "Add Model" }).last().click(); + await page.getByTestId("add-model-btn").click(); // Scope to the toast container so a stale toast can't satisfy this. await expect(page.locator("[data-sonner-toast]").getByText("created successfully").last()).toBeVisible({ @@ -347,12 +346,9 @@ test.describe("Add Model", () => { // The Models table renders team-scoped models with the team id in the row. await page.getByRole("tab", { name: "All Models" }).click(); await page.waitForLoadState("networkidle"); - // networkidle fires before the table finishes re-rendering. - await page.waitForTimeout(2000); await page.getByPlaceholder("Search model names").fill("cohere"); - await page.waitForTimeout(1000); - + // Clearer failure than timing out on a row assertion when the table is empty. await expect(page.getByTestId("pagination-range")).toHaveText(/Showing \d+-\d+ of \d+/, { timeout: 15_000, @@ -361,7 +357,7 @@ test.describe("Add Model", () => { // Pin to one row carrying both the name and the team, so the sibling test's // team-less cohere row can't satisfy it. const teamCohereRow = page - .locator("table tbody tr") + .getByRole("row") .filter({ hasText: "cohere/" }) .filter({ hasText: E2E_TEAM_CRUD_ID }); await expect(teamCohereRow).toHaveCount(1, { timeout: 15_000 }); @@ -387,7 +383,7 @@ test.describe("Add Model", () => { // Click Add Model button by its text const created = await captureRequestBody(page, { method: "POST", urlIncludes: "/model/new" }, async () => { - await page.getByRole("button", { name: "Add Model" }).last().click(); + await page.getByTestId("add-model-btn").click(); }); // A wildcard with the star stripped becomes a plain "cohere" deployment that matches nothing. expect(created.model_name, "the wildcard route goes on the wire intact").toBe("cohere/*"); @@ -398,11 +394,9 @@ test.describe("Add Model", () => { // Navigate to All Models tab await page.getByRole("tab", { name: "All Models" }).click(); await page.waitForLoadState("networkidle"); - await page.waitForTimeout(2000); // Search for the wildcard model await page.getByPlaceholder("Search model names").fill("cohere"); - await page.waitForTimeout(1000); // Verify the model appears in the results count (not "Showing 0 results") await expect(page.getByTestId("pagination-range")).toHaveText(/Showing \d+-\d+ of \d+/, { @@ -410,8 +404,7 @@ test.describe("Add Model", () => { }); // Verify the wildcard model appears in the table body (wildcard models show as "cohere/*") - const tableBody = page.locator("table tbody"); - await expect(tableBody.getByText("cohere/").first()).toBeVisible({ timeout: 15_000 }); + await expect(page.getByRole("row").filter({ hasText: "cohere/" })).not.toHaveCount(0, { timeout: 15_000 }); // "cohere/" in the table also matches a plain cohere deployment; require the wildcard exactly. const stored = await findDeploymentByName(page, "cohere/*"); diff --git a/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts b/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts index 1d080ec82b8..fe168267a54 100644 --- a/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts @@ -17,49 +17,32 @@ async function openTemplateSelect(page: PlaywrightPage) { return trigger; } -function pollPixelsBelowTrigger(trigger: Locator, popup: Locator) { +function pollOptionsCoverTrigger(trigger: Locator, options: Locator) { return expect.poll(async () => { const triggerBox = await trigger.boundingBox(); - const popupBox = await popup.boundingBox(); - if (!triggerBox || !popupBox) return null; - return popupBox.y - (triggerBox.y + triggerBox.height); - }); -} - -function pollPopupOverlapsTrigger(trigger: Locator, popup: Locator) { - return expect.poll(async () => { - const triggerBox = await trigger.boundingBox(); - const popupBox = await popup.boundingBox(); - if (!triggerBox || !popupBox) return null; - return popupBox.y < triggerBox.y + triggerBox.height && popupBox.y + popupBox.height > triggerBox.y; + const optionsBox = await options.boundingBox(); + if (!triggerBox || !optionsBox) return null; + return optionsBox.y < triggerBox.y + triggerBox.height && optionsBox.y + optionsBox.height > triggerBox.y; }); } test.describe("Auto Router template select anchoring", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); - test("opens the options below the trigger rather than over it", async ({ page }) => { - await page.setViewportSize({ width: 1280, height: 900 }); - const trigger = await openTemplateSelect(page); + for (const { room, height } of [ + { room: "with room below it", height: 900 }, + { room: "with no room below it", height: 560 }, + ]) { + test(`keeps the trigger uncovered when the options open ${room}`, async ({ page }) => { + await page.setViewportSize({ width: 1280, height }); + const trigger = await openTemplateSelect(page); + await trigger.scrollIntoViewIfNeeded(); - await trigger.click(); - const popup = page.locator('[data-slot="select-content"]'); - await expect(popup).toBeVisible(); + await trigger.click(); + const options = page.getByRole("listbox"); + await expect(options).toBeVisible(); - // Item-aligned mode reports "none" and puts the active item over the trigger. - await expect(popup).toHaveAttribute("data-side", "bottom"); - await pollPixelsBelowTrigger(trigger, popup).toBeGreaterThanOrEqual(0); - }); - - test("flips above the trigger instead of covering it when there is no room below", async ({ page }) => { - await page.setViewportSize({ width: 1280, height: 560 }); - const trigger = await openTemplateSelect(page); - await trigger.scrollIntoViewIfNeeded(); - - await trigger.click(); - const popup = page.locator('[data-slot="select-content"]'); - await expect(popup).toBeVisible(); - - await pollPopupOverlapsTrigger(trigger, popup).toBe(false); - }); + await pollOptionsCoverTrigger(trigger, options).toBe(false); + }); + } }); diff --git a/tests/e2e/ui/tests/modelsPage/responsiveHeader.spec.ts b/tests/e2e/ui/tests/modelsPage/responsiveHeader.spec.ts index 6ad1ccb8451..366c371208f 100644 --- a/tests/e2e/ui/tests/modelsPage/responsiveHeader.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/responsiveHeader.spec.ts @@ -7,9 +7,7 @@ test.describe("Models and Endpoints responsive header", () => { viewport: { width: 900, height: 720 }, }); - test("keeps the refresh action on the same row as the tabs", async ({ - page, - }) => { + test("keeps the refresh action on the same row as the tabs", async ({ page }) => { await page.goto("/ui"); await page .getByRole("complementary") @@ -26,8 +24,7 @@ test.describe("Models and Endpoints responsive header", () => { expect(tabsBox).not.toBeNull(); expect(refreshBox).not.toBeNull(); - const tabsCenterY = tabsBox!.y + tabsBox!.height / 2; - const refreshCenterY = refreshBox!.y + refreshBox!.height / 2; - expect(Math.abs(tabsCenterY - refreshCenterY)).toBeLessThanOrEqual(2); + const sharesARow = refreshBox!.y < tabsBox!.y + tabsBox!.height && refreshBox!.y + refreshBox!.height > tabsBox!.y; + expect(sharesARow, "refresh wrapped onto its own row below the tabs").toBe(true); }); }); diff --git a/tests/e2e/ui/tests/proxy-admin/keys.spec.ts b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts index 0c38641dcc7..f5bee68f245 100644 --- a/tests/e2e/ui/tests/proxy-admin/keys.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts @@ -43,7 +43,7 @@ test.describe("Proxy Admin - Keys", () => { const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); await page.keyboard.type(E2E_TEAM_CRUD_ALIAS); - await page.locator('[data-slot="combobox-content"]:visible').getByText(E2E_TEAM_CRUD_ALIAS).first().click(); + await page.getByRole("option", { name: E2E_TEAM_CRUD_ALIAS }).first().click(); // Select models — the popup is portaled to the body, so scope options to the page. await page.getByRole("combobox", { name: "Select models" }).click(); @@ -74,10 +74,9 @@ test.describe("Proxy Admin - Keys", () => { const before = await findKeyByAlias(page, E2E_REGENERATE_KEY_ALIAS); expect(before?.token, `seeded key ${E2E_REGENERATE_KEY_ALIAS} has a token`).toBeTruthy(); - // Key IDs are rendered as buttons in the table - const keyRow = page.locator("tr", { hasText: E2E_REGENERATE_KEY_ALIAS }); + const keyRow = page.getByRole("row").filter({ hasText: E2E_REGENERATE_KEY_ALIAS }); await expect(keyRow).toBeVisible({ timeout: 10_000 }); - await keyRow.locator("button").first().click(); + await keyRow.getByRole("button", { name: E2E_REGENERATE_KEY_ALIAS }).click(); await expect(page.getByText("Back to Keys")).toBeVisible({ timeout: 10_000 }); @@ -109,9 +108,9 @@ test.describe("Proxy Admin - Keys", () => { const before = await findKeyByAlias(page, E2E_UPDATE_LIMITS_KEY_ALIAS); expect(before, `seeded key ${E2E_UPDATE_LIMITS_KEY_ALIAS} exists`).toBeTruthy(); - const keyRow = page.locator("tr", { hasText: E2E_UPDATE_LIMITS_KEY_ALIAS }); + const keyRow = page.getByRole("row").filter({ hasText: E2E_UPDATE_LIMITS_KEY_ALIAS }); await expect(keyRow).toBeVisible({ timeout: 10_000 }); - await keyRow.locator("button").first().click(); + await keyRow.getByRole("button", { name: E2E_UPDATE_LIMITS_KEY_ALIAS }).click(); await expect(page.getByText("Back to Keys")).toBeVisible({ timeout: 10_000 }); @@ -147,9 +146,9 @@ test.describe("Proxy Admin - Keys", () => { await navigateToPage(page, Page.ApiKeys); await dismissFeedbackPopup(page); - const keyRow = page.locator("tr", { hasText: E2E_DELETE_KEY_ALIAS }); + const keyRow = page.getByRole("row").filter({ hasText: E2E_DELETE_KEY_ALIAS }); await expect(keyRow).toBeVisible({ timeout: 10_000 }); - await keyRow.locator("button").first().click(); + await keyRow.getByRole("button", { name: E2E_DELETE_KEY_ALIAS }).click(); await expect(page.getByText("Back to Keys")).toBeVisible({ timeout: 10_000 }); diff --git a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts index f93cca75347..5116cb5df19 100644 --- a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts +++ b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts @@ -142,7 +142,7 @@ test.describe("Team Admin", () => { const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); await page.keyboard.type(E2E_TEAM_CRUD_ALIAS); - await page.locator('[data-slot="combobox-content"]:visible').getByText(E2E_TEAM_CRUD_ALIAS).first().click(); + await page.getByRole("option", { name: E2E_TEAM_CRUD_ALIAS }).first().click(); // Models — pick "All Team Models". The popup is portaled to the body, so // scope the option lookup to the page. diff --git a/tests/e2e/ui/tests/usage/usagePage.spec.ts b/tests/e2e/ui/tests/usage/usagePage.spec.ts index 8fa59beb905..f61ab018e1b 100644 --- a/tests/e2e/ui/tests/usage/usagePage.spec.ts +++ b/tests/e2e/ui/tests/usage/usagePage.spec.ts @@ -51,20 +51,19 @@ test.describe("Usage page", () => { const card = await openUsage(page); // Table view (the default): the key is listed by its alias. - const row = card.locator("tbody tr").filter({ hasText: alias }); + const row = card.getByRole("row").filter({ hasText: alias }); await expect(row, `${alias} missing from Top Virtual Keys`).toHaveCount(1, { timeout: 30_000, }); // Chart view swaps the table out for the bar chart, and back. await card.getByText("Chart View", { exact: true }).click(); - await expect(card.locator("tbody tr")).toHaveCount(0, { timeout: 10_000 }); + await expect(card.getByRole("table")).toHaveCount(0, { timeout: 10_000 }); await card.getByText("Table View", { exact: true }).click(); await expect(row).toHaveCount(1, { timeout: 10_000 }); - // Clicking the Key ID cell fetches key info and opens the detail panel. // The alias is already in the row behind the modal, so match the panel's own controls. - await row.locator("td").first().click(); + await row.getByRole("button", { name: token }).click(); const keyInfo = page.getByRole("tab", { name: "Overview", exact: true }); await expect(keyInfo, "key info panel did not open").toBeVisible({ timeout: 20_000, diff --git a/tests/e2e/ui/tests/users/searchUsers.spec.ts b/tests/e2e/ui/tests/users/searchUsers.spec.ts index e87218b5a5e..5e7e3e35b91 100644 --- a/tests/e2e/ui/tests/users/searchUsers.spec.ts +++ b/tests/e2e/ui/tests/users/searchUsers.spec.ts @@ -1,91 +1,51 @@ -import { test, expect, Page } from "@playwright/test"; +import { test, expect, Page as PlaywrightPage } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; -test.skip("Internal Users Search", () => { +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; + +const userRows = (page: PlaywrightPage) => page.getByRole("row").filter({ has: page.getByRole("cell") }); + +async function goToInternalUsers(page: PlaywrightPage) { + await navigateToPage(page, Page.Users); + await expect(page.getByRole("columnheader", { name: "User ID" })).toBeVisible({ timeout: 30_000 }); + await expect(userRows(page)).not.toHaveCount(0, { timeout: 30_000 }); +} + +test.describe("Internal Users Search", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); - async function goToInternalUsers(page: Page) { - await page.goto("/ui"); - - const tab = page.getByRole("menuitem", { name: "Internal User" }); - await expect(tab).toBeVisible(); - await tab.click(); - - await expect(page.locator("tbody tr").first()).toBeVisible(); - await expect(page.locator('[data-slot="skeleton"]')).toHaveCount(0); - } - - test("can search users by email", async ({ page }) => { + test("narrows the table to the matching email, and restores it when cleared", async ({ page }) => { await goToInternalUsers(page); - const rows = page.locator("tbody tr"); - const searchInput = page.getByPlaceholder("Search by email..."); + const search = page.getByPlaceholder("Search by email…"); + await expect(search).toBeVisible(); - await expect(searchInput).toBeVisible(); + await search.fill("noteam@"); + await expect(userRows(page)).toHaveCount(1, { timeout: 30_000 }); + await expect(userRows(page).first()).toContainText("noteam@test.local"); - // Ensure initial data is loaded - const initialCount = await rows.count(); - expect(initialCount).toBeGreaterThan(0); - - // 🔹 Apply filter + wait for backend response - await Promise.all([ - page.waitForResponse( - (res) => - res.url().includes("/user/list") && - res.url().includes("user_email=test%40") && // encoded "test@" - res.status() === 200, - ), - searchInput.fill("test@"), - ]); - await page.waitForTimeout(5000); - const filteredCount = await rows.count(); - await expect(filteredCount).toBeLessThan(initialCount); - - // 🔹 Clear filter + wait for unfiltered request - await Promise.all([ - page.waitForResponse( - (res) => res.url().includes("/user/list") && !res.url().includes("user_email=") && res.status() === 200, - ), - searchInput.clear(), - ]); - - const resetCount = await rows.count(); - await expect(resetCount).toBe(initialCount); + await search.clear(); + await expect(userRows(page).filter({ hasText: "admin@test.local" })).not.toHaveCount(0, { timeout: 30_000 }); }); - test("can filter users by user ID and SSO ID", async ({ page }) => { + test("filters the table down to one user by user ID", async ({ page }) => { await goToInternalUsers(page); - const rows = page.locator("tbody tr"); - // Ensure initial data is loaded - const initialCount = await rows.count(); - expect(initialCount).toBeGreaterThan(0); + await page.getByRole("button", { name: "Filters" }).click(); + await page.getByTestId("users-filter-user-id").fill("e2e-internal-noteam"); + await page.getByTestId("filter-drawer-apply").click(); - const filtersButton = page.getByRole("button", { - name: "Filters", - exact: true, - }); - await filtersButton.click(); + await expect(userRows(page)).toHaveCount(1, { timeout: 30_000 }); + await expect(userRows(page).first()).toContainText("noteam@test.local"); + }); - const userIdInput = page.getByPlaceholder("Filter by User ID"); - const ssoIdInput = page.getByPlaceholder("Filter by SSO ID"); - await Promise.all([ - page.waitForResponse( - (res) => res.url().includes("/user/list") && res.url().includes("user_ids=user") && res.status() === 200, - ), - userIdInput.fill("user"), - ]); + test("shows no users when the SSO ID matches nobody", async ({ page }) => { + await goToInternalUsers(page); - await Promise.all([ - page.waitForResponse( - (res) => - res.url().includes("/user/list") && - res.url().includes("user_ids=user") && - res.url().includes("sso_user_ids=sso") && - res.status() === 200, - ), - ssoIdInput.fill("sso"), - ]); - const combinedFilteredCount = await rows.count(); - await expect(combinedFilteredCount).toBeLessThan(initialCount); + await page.getByRole("button", { name: "Filters" }).click(); + await page.getByTestId("users-filter-sso-id").fill("e2e-sso-id-that-matches-nobody"); + await page.getByTestId("filter-drawer-apply").click(); + + await expect(userRows(page)).toHaveCount(0, { timeout: 30_000 }); }); }); diff --git a/tests/e2e/ui/tests/users/viewInternalUsers.spec.ts b/tests/e2e/ui/tests/users/viewInternalUsers.spec.ts index 614191372d0..b46fb4d112a 100644 --- a/tests/e2e/ui/tests/users/viewInternalUsers.spec.ts +++ b/tests/e2e/ui/tests/users/viewInternalUsers.spec.ts @@ -1,54 +1,29 @@ -import { test, expect, Page } from "@playwright/test"; +import { test, expect, Page as PlaywrightPage } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; -test.skip("Internal Users Page", () => { +async function goToInternalUsers(page: PlaywrightPage) { + await navigateToPage(page, Page.Users); + await expect(page.getByRole("columnheader", { name: "User ID" })).toBeVisible({ timeout: 30_000 }); + await expect(userRows(page)).not.toHaveCount(0, { timeout: 30_000 }); +} + +const userRows = (page: PlaywrightPage) => page.getByRole("row").filter({ has: page.getByRole("cell") }); + +test.describe("Internal Users Page", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); - async function goToInternalUsers(page: Page) { - await page.goto("/ui"); - - const internalUserTab = page.getByRole("menuitem", { name: "Internal User" }); - await expect(internalUserTab).toBeVisible(); - await internalUserTab.click(); - - const firstRow = page.locator("tbody tr").first(); - await expect(firstRow).toBeVisible(); - await expect(page.locator('[data-slot="skeleton"]')).toHaveCount(0); - } - - test("renders internal users table correctly", async ({ page }) => { + test("lists the seeded users under the identifying columns", async ({ page }) => { await goToInternalUsers(page); - const rows = page.locator("tbody tr"); - const rowCount = await rows.count(); - expect(rowCount).toBeGreaterThan(0); - - const userIdHeader = page.getByRole("columnheader", { name: "User ID" }); - await expect(userIdHeader).toBeVisible(); - - const virtualKeysHeader = page.getByRole("columnheader", { name: "Virtual Keys" }); - await expect(virtualKeysHeader).toBeVisible(); + await expect(page.getByRole("columnheader", { name: "User ID" })).toBeVisible(); + await expect(page.getByRole("columnheader", { name: "Virtual Keys" })).toBeVisible(); }); - test("pagination controls work correctly", async ({ page }) => { + test("cannot page backwards off the first page", async ({ page }) => { await goToInternalUsers(page); - const paginationInfo = page.locator(".text-sm.text-gray-700"); - const prevButton = page.getByRole("button", { name: "Previous" }); - const nextButton = page.getByRole("button", { name: "Next" }); - - const infoText = (await paginationInfo.textContent()) || ""; - - // On first page, Previous should be disabled - if (infoText.includes("1 -")) { - await expect(prevButton).toBeDisabled(); - } - - await page.waitForTimeout(1000); - // Check if there are more pages - const hasMorePages = infoText.includes("of") && !infoText.endsWith("25 of 25"); - if (hasMorePages) { - await expect(nextButton).toBeEnabled(); - } + await expect(page.getByRole("button", { name: "Go to previous page" })).toBeDisabled(); }); }); diff --git a/ui/litellm-dashboard/src/components/leftnav.test.tsx b/ui/litellm-dashboard/src/components/leftnav.test.tsx index b0bb0e8a5b5..c3e1f924d09 100644 --- a/ui/litellm-dashboard/src/components/leftnav.test.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.test.tsx @@ -213,6 +213,20 @@ describe("Sidebar (leftnav)", () => { expect(screen.getByText("Search Tools")).toBeInTheDocument(); }); }); + it("reports whether a nested tab is expanded", async () => { + renderWithProviders(); + + const toggle = screen.getByText("Tools").closest("button")!; + expect(toggle).toHaveAttribute("aria-expanded", "false"); + + act(() => { + fireEvent.click(toggle); + }); + await waitFor(() => { + expect(toggle).toHaveAttribute("aria-expanded", "true"); + }); + }); + it("keeps Router Settings as a single Settings child", () => { // Router Settings is admin-only, so getAvailablePages() filters it out entirely and the // page_utils duplicate-key guard cannot see it. Walk menuGroups directly, otherwise a diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 42389facac2..51ba36348e1 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -570,6 +570,7 @@ const Sidebar_: React.FC = ({ toggleGroup(item.key)} title={collapsed ? labelText(item) : undefined} > diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SectionHeader.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SectionHeader.test.tsx index 5aee6b33ec5..6cd9f476628 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SectionHeader.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SectionHeader.test.tsx @@ -53,6 +53,24 @@ describe("SectionHeader", () => { expect(onToggleCollapse).toHaveBeenCalledTimes(1); }); + it("reports its collapsed state to assistive technology", () => { + const { rerender } = render( + , + ); + + expect(screen.getByRole("button", { name: /^Input/ })).toHaveAttribute("aria-expanded", "true"); + + rerender(); + + expect(screen.getByRole("button", { name: /^Input/ })).toHaveAttribute("aria-expanded", "false"); + }); + + it("names each copy button for the section it belongs to", () => { + render(); + + expect(screen.getByRole("button", { name: "Copy output" })).toBeInTheDocument(); + }); + it("stays inert when no toggle handler is given", async () => { const onCopy = vi.fn(); render(); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SectionHeader.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SectionHeader.tsx index 93e9953b2ee..6a18aaf8642 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SectionHeader.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SectionHeader.tsx @@ -17,6 +17,8 @@ interface SectionHeaderProps { turnCount?: number; } +const SUMMARY_CLASSES = "flex flex-1 items-center gap-4"; + export function SectionHeader({ type, tokens, @@ -26,42 +28,53 @@ export function SectionHeader({ onToggleCollapse, turnCount, }: SectionHeaderProps) { + const summary = ( + <> + {onToggleCollapse && + (isCollapsed ? ( + + ) : ( + + ))} + +
+ {type === "input" ? ( + + ) : ( + ✨ + )} + {type === "input" ? "Input" : "Output"} +
+ + {tokens !== undefined && Tokens: {tokens.toLocaleString()}} + + {cost !== undefined && Cost: ${cost.toFixed(6)}} + + {turnCount !== undefined && turnCount > 0 && ( + Turns: {turnCount} + )} + + ); + return (
-
- {onToggleCollapse && - (isCollapsed ? ( - - ) : ( - - ))} - -
- {type === "input" ? ( - - ) : ( - ✨ - )} - {type === "input" ? "Input" : "Output"} -
- - {tokens !== undefined && ( - Tokens: {tokens.toLocaleString()} - )} - - {cost !== undefined && Cost: ${cost.toFixed(6)}} - - {turnCount !== undefined && turnCount > 0 && ( - Turns: {turnCount} - )} -
+ {onToggleCollapse ? ( + + ) : ( +
{summary}
+ )} { e.stopPropagation(); onCopy(); From 98b1e2e7b4a71d9f77a42f6c9631188ebbfd9eb5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 14:28:30 -0700 Subject: [PATCH 104/120] fix(gigachat): correct cached-token accounting, stream usage on all finish reasons, stop mutating cached request body 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. --- litellm/llms/gigachat/chat/streaming.py | 35 ++++---- litellm/llms/gigachat/utils.py | 21 ++--- .../llm_passthrough_endpoints.py | 25 +++--- .../chat/test_gigachat_chat_streaming.py | 85 +++++++++++++++++++ .../test_litellm/llms/gigachat/test_utils.py | 6 +- .../test_llm_pass_through_endpoints.py | 71 ++++++++++++++++ 6 files changed, 196 insertions(+), 47 deletions(-) create mode 100644 tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_streaming.py 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", From 78e1c658b4828ac5595d1bdabb259d873691148d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 14:45:11 -0700 Subject: [PATCH 105/120] fix(e2e): assert sidebar expansion without a self-resolving locator The migration smoke waited on `getByRole("button", { expanded: false })` after clicking it. Playwright re-resolves that locator on every retry, so once the clicked group flipped to expanded it matched the next collapsed group instead, and the assertion could never pass. Count the remaining collapsed groups and wait for that count to drop by one. --- tests/e2e/ui/tests/migration/migratedPages.spec.ts | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/e2e/ui/tests/migration/migratedPages.spec.ts b/tests/e2e/ui/tests/migration/migratedPages.spec.ts index 473d0b795f1..547330190bd 100644 --- a/tests/e2e/ui/tests/migration/migratedPages.spec.ts +++ b/tests/e2e/ui/tests/migration/migratedPages.spec.ts @@ -35,11 +35,12 @@ async function expectRendered(page: Page) { */ async function clickSidebar(page: Page, segment: string) { const link = sidebar(page).locator(`a[href$="/ui/${segment}"]`).first(); + const collapsedGroups = sidebar(page).getByRole("button", { expanded: false }); for (let i = 0; i < 8 && !(await link.isVisible().catch(() => false)); i++) { - const collapsedGroup = sidebar(page).getByRole("button", { expanded: false }).first(); - if (!(await collapsedGroup.isVisible().catch(() => false))) break; - await collapsedGroup.click(); - await expect(collapsedGroup).toHaveAttribute("aria-expanded", "true"); + const stillCollapsed = await collapsedGroups.count(); + if (stillCollapsed === 0) break; + await collapsedGroups.first().click(); + await expect(collapsedGroups).toHaveCount(stillCollapsed - 1); } await link.click(); } From cc258b5473932c939903d589604f83f2ca260469 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 14:54:11 -0700 Subject: [PATCH 106/120] test(e2e): keep the placement guarantees the geometry rewrites dropped The consolidated popup test only asserted the options never cover the trigger, so opening above the trigger with room below it, the regression PR #38554 fixed, would have passed. Split it back into a below-trigger case and a cramped-viewport case. The header test accepted a single pixel of vertical intersection; require the refresh control's centre to sit within the tab row instead. --- .../autoRouterTemplateSelect.spec.ts | 58 +++++++++++++------ .../tests/modelsPage/responsiveHeader.spec.ts | 3 +- 2 files changed, 42 insertions(+), 19 deletions(-) diff --git a/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts b/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts index fe168267a54..7fc20104f20 100644 --- a/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts @@ -17,32 +17,54 @@ async function openTemplateSelect(page: PlaywrightPage) { return trigger; } +async function boxes(trigger: Locator, options: Locator) { + const triggerBox = await trigger.boundingBox(); + const optionsBox = await options.boundingBox(); + return triggerBox && optionsBox ? { triggerBox, optionsBox } : null; +} + +function pollOptionsOpenBelowTrigger(trigger: Locator, options: Locator) { + return expect.poll(async () => { + const box = await boxes(trigger, options); + return box && box.optionsBox.y >= box.triggerBox.y + box.triggerBox.height; + }); +} + function pollOptionsCoverTrigger(trigger: Locator, options: Locator) { return expect.poll(async () => { - const triggerBox = await trigger.boundingBox(); - const optionsBox = await options.boundingBox(); - if (!triggerBox || !optionsBox) return null; - return optionsBox.y < triggerBox.y + triggerBox.height && optionsBox.y + optionsBox.height > triggerBox.y; + const box = await boxes(trigger, options); + return ( + box && + box.optionsBox.y < box.triggerBox.y + box.triggerBox.height && + box.optionsBox.y + box.optionsBox.height > box.triggerBox.y + ); }); } test.describe("Auto Router template select anchoring", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); - for (const { room, height } of [ - { room: "with room below it", height: 900 }, - { room: "with no room below it", height: 560 }, - ]) { - test(`keeps the trigger uncovered when the options open ${room}`, async ({ page }) => { - await page.setViewportSize({ width: 1280, height }); - const trigger = await openTemplateSelect(page); - await trigger.scrollIntoViewIfNeeded(); + test("opens the options below the trigger when there is room below it", async ({ page }) => { + await page.setViewportSize({ width: 1280, height: 900 }); + const trigger = await openTemplateSelect(page); + await trigger.scrollIntoViewIfNeeded(); - await trigger.click(); - const options = page.getByRole("listbox"); - await expect(options).toBeVisible(); + await trigger.click(); + const options = page.getByRole("listbox"); + await expect(options).toBeVisible(); - await pollOptionsCoverTrigger(trigger, options).toBe(false); - }); - } + await pollOptionsOpenBelowTrigger(trigger, options).toBe(true); + }); + + test("keeps the trigger uncovered when the options open with no room below it", async ({ page }) => { + await page.setViewportSize({ width: 1280, height: 560 }); + const trigger = await openTemplateSelect(page); + await trigger.scrollIntoViewIfNeeded(); + + await trigger.click(); + const options = page.getByRole("listbox"); + await expect(options).toBeVisible(); + + await pollOptionsCoverTrigger(trigger, options).toBe(false); + }); }); diff --git a/tests/e2e/ui/tests/modelsPage/responsiveHeader.spec.ts b/tests/e2e/ui/tests/modelsPage/responsiveHeader.spec.ts index 366c371208f..aabdf18d427 100644 --- a/tests/e2e/ui/tests/modelsPage/responsiveHeader.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/responsiveHeader.spec.ts @@ -24,7 +24,8 @@ test.describe("Models and Endpoints responsive header", () => { expect(tabsBox).not.toBeNull(); expect(refreshBox).not.toBeNull(); - const sharesARow = refreshBox!.y < tabsBox!.y + tabsBox!.height && refreshBox!.y + refreshBox!.height > tabsBox!.y; + const refreshCenterY = refreshBox!.y + refreshBox!.height / 2; + const sharesARow = refreshCenterY > tabsBox!.y && refreshCenterY < tabsBox!.y + tabsBox!.height; expect(sharesARow, "refresh wrapped onto its own row below the tabs").toBe(true); }); }); From cc078edd1fb2a29e2c5c2826803ae399e6078143 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 31 Aug 2026 21:56:57 +0000 Subject: [PATCH 107/120] test(mcp): isolate global MCP server registry in discoverable endpoints tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/test_discoverable_endpoints.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 3279c59acd4..598e9276423 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -35,6 +35,22 @@ def mock_mcp_client_ip(): yield +@pytest.fixture(autouse=True) +def isolate_global_mcp_registry(): + """Restore the module-global MCP server registry after each test. + + Tests here register servers on ``global_mcp_server_manager`` directly; without a + restore, entries leak into other test modules sharing the same worker and break + assertions over the full registry contents. + """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + snapshot = dict(global_mcp_server_manager.registry) + yield + global_mcp_server_manager.registry.clear() + global_mcp_server_manager.registry.update(snapshot) + + def _mock_callback_request(base_url: str = "http://localhost:3000/"): """Return a MagicMock Request for callback/authorize same-origin tests. From 613d0ef3fa7153e5485698ae6fae47481f23b027 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 15:04:07 -0700 Subject: [PATCH 108/120] test(gcs_pub_sub): expect router_metadata in the spend logs payload The base added router_metadata to SpendLogsMetadata in #39001 without updating this fixture, and its CI run never executed logging_testing, so the job now fails on every branch merged with current staging. --- .../gcs_pub_sub_body/spend_logs_payload.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json index 1838fb16e91..28912a27501 100644 --- a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json +++ b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json @@ -11,7 +11,7 @@ "user": "", "team_id": "", "organization_id": "", - "metadata": "{\"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", + "metadata": "{\"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"router_metadata\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, From e7b7a2276fb1dcc3fd2381630412b8b076978ffa Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 15:11:50 -0700 Subject: [PATCH 109/120] test(e2e): cover SCIM token creation and SCIM API auth in the Admin UI suite --- tests/e2e/ui/tests/settings/scim.spec.ts | 66 ++++++++++++++++++++++++ 1 file changed, 66 insertions(+) create mode 100644 tests/e2e/ui/tests/settings/scim.spec.ts diff --git a/tests/e2e/ui/tests/settings/scim.spec.ts b/tests/e2e/ui/tests/settings/scim.spec.ts new file mode 100644 index 00000000000..5489111a580 --- /dev/null +++ b/tests/e2e/ui/tests/settings/scim.spec.ts @@ -0,0 +1,66 @@ +import { test, expect, Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { masterKey } from "../../helpers/traffic"; + +const rootPath = (): string => process.env.SERVER_ROOT_PATH ?? ""; + +async function createScimTokenViaUi(page: PlaywrightPage, alias: string): Promise { + await navigateToPage(page, Page.AdminPanel); + await page.getByRole("tab", { name: "SCIM" }).click(); + + await expect(page.getByText("SCIM Tenant URL")).toBeVisible(); + await expect(page.locator("input[disabled]").first()).toHaveValue(/\/scim\/v2$/); + + await page.getByLabel("Token Name").fill(alias); + await page.getByRole("button", { name: "Create SCIM Token" }).click(); + + await expect(page.getByText(/copy this token now/i)).toBeVisible({ timeout: 15_000 }); + const token = await page.locator('input[type="password"]').inputValue(); + expect(token, "the one-time token panel shows a usable virtual key").toMatch(/^sk-/); + return token; +} + +test.describe("Admin Settings - SCIM", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Create SCIM Token shows the token once and offers to create another", async ({ page }) => { + const token = await createScimTokenViaUi(page, `e2e-scim-ui-${Date.now()}`); + + await page.getByRole("button", { name: "Create Another Token" }).click(); + await expect(page.getByRole("button", { name: "Create SCIM Token" })).toBeVisible(); + await expect(page.getByText(/copy this token now/i)).toBeHidden(); + + await deleteKey(page, token); + }); + + test("a UI-minted SCIM token authorizes the SCIM API", async ({ page, request }) => { + test.skip(!process.env.LITELLM_LICENSE, "LITELLM_LICENSE not set in test env — /scim/v2 is premium-gated"); + + const token = await createScimTokenViaUi(page, `e2e-scim-api-${Date.now()}`); + + const denied = await request.get(`${rootPath()}/scim/v2/Groups`, { + headers: { Authorization: "Bearer sk-not-a-real-key" }, + }); + expect(denied.status(), "an unknown key must not reach SCIM").toBe(401); + + const res = await request.get(`${rootPath()}/scim/v2/Groups`, { + headers: { Authorization: `Bearer ${token}` }, + }); + expect(res.status(), `SCIM Groups listing failed: ${await res.text()}`).toBe(200); + const body = await res.json(); + expect(body.schemas, "SCIM answers with a ListResponse").toContain("urn:ietf:params:scim:api:messages:2.0:ListResponse"); + expect(Array.isArray(body.Resources), "SCIM ListResponse carries a Resources array").toBe(true); + + await deleteKey(page, token); + }); +}); + +async function deleteKey(page: PlaywrightPage, key: string): Promise { + const res = await page.request.post(`${rootPath()}/key/delete`, { + headers: { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json" }, + data: { keys: [key] }, + }); + expect(res.ok(), `cleanup of the SCIM token failed (${res.status()})`).toBe(true); +} From 91595780ecd8e475361691f6876917e3963d1994 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 15:16:44 -0700 Subject: [PATCH 110/120] fix(gigachat): fold cached tokens back into prompt and total token counts GigaChat reports prompt_tokens and total_tokens after subtracting cached tokens (the docs example is prompt_tokens=1, precached_prompt_tokens=37, total_tokens=5, so the fields are disjoint, not a subset). Map to the OpenAI convention by adding precached_prompt_tokens back onto prompt and total while still surfacing it as prompt_tokens_details.cached_tokens. --- litellm/llms/gigachat/utils.py | 4 ++-- .../llms/gigachat/chat/test_gigachat_chat_streaming.py | 10 ++++++---- tests/test_litellm/llms/gigachat/test_utils.py | 8 +++++--- 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/litellm/llms/gigachat/utils.py b/litellm/llms/gigachat/utils.py index 2072e6003e3..cbb35cd1b57 100644 --- a/litellm/llms/gigachat/utils.py +++ b/litellm/llms/gigachat/utils.py @@ -15,10 +15,10 @@ def convert_usage(usage_data: Mapping[str, int]) -> Usage: ) return Usage( - prompt_tokens=usage_data.get("prompt_tokens", 0), + prompt_tokens=usage_data.get("prompt_tokens", 0) + precached_prompt_tokens, completion_tokens=usage_data.get("completion_tokens", 0), prompt_tokens_details=prompt_tokens_details, - total_tokens=usage_data.get("total_tokens", 0), + total_tokens=usage_data.get("total_tokens", 0) + precached_prompt_tokens, ) 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 index 500aff264f2..35ca93319f5 100644 --- a/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_streaming.py +++ b/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_streaming.py @@ -65,8 +65,10 @@ class TestChunkParserUsage: 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.""" + def test_cache_hit_usage_folds_cached_tokens_back_in(self): + """GigaChat reports prompt_tokens and total_tokens after subtracting cached tokens + (docs example: prompt_tokens=1, precached_prompt_tokens=37, total_tokens=5), so the + OpenAI-convention usage must add them back and surface them as cached_tokens.""" parsed = _parse( { "choices": [{"delta": {"content": ""}, "index": 0, "finish_reason": "stop"}], @@ -80,6 +82,6 @@ class TestChunkParserUsage: ) assert parsed["usage"] is not None - assert parsed["usage"]["prompt_tokens"] == 25 - assert parsed["usage"]["total_tokens"] == 32 + assert parsed["usage"]["prompt_tokens"] == 45 + assert parsed["usage"]["total_tokens"] == 52 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 3d6b7d842e7..71a193d7b29 100644 --- a/tests/test_litellm/llms/gigachat/test_utils.py +++ b/tests/test_litellm/llms/gigachat/test_utils.py @@ -26,7 +26,9 @@ class TestConvertUsage: ) def test_usage_with_precached_prompt_tokens(self): - """precached_prompt_tokens is a subset of prompt_tokens (OpenAI cached_tokens semantics), never additive.""" + """GigaChat's prompt_tokens and total_tokens exclude cached tokens (docs example: + prompt_tokens=1, precached_prompt_tokens=37, total_tokens=5), so OpenAI-convention + usage adds precached back in and surfaces it as cached_tokens.""" result = convert_usage( { "prompt_tokens": 10, @@ -37,9 +39,9 @@ class TestConvertUsage: ) assert result == Usage( - prompt_tokens=10, + prompt_tokens=13, completion_tokens=5, - total_tokens=15, + total_tokens=18, prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=3), ) From abd8beec018eb8faf4e40530f3063b9ed6cb15fb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 15:28:34 -0700 Subject: [PATCH 111/120] fix(e2e): measure the clipped popup and assert the table's empty state Two assertions were checking the wrong thing. The anchoring tests read getByRole("listbox"), which resolves to SelectPrimitive.List; that sits at full content height inside the popup that clips and scrolls it, so the box overlapped the trigger even when nothing visible did. Measure the popup. The SSO-ID search expected zero rows, but DataTable renders a "No results" message row when a filter matches nothing, so the count is one. Assert the empty state the user actually sees. --- .../modelsPage/autoRouterTemplateSelect.spec.ts | 12 ++++++------ tests/e2e/ui/tests/users/searchUsers.spec.ts | 2 +- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts b/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts index 7fc20104f20..51df50a2e68 100644 --- a/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/autoRouterTemplateSelect.spec.ts @@ -23,6 +23,8 @@ async function boxes(trigger: Locator, options: Locator) { return triggerBox && optionsBox ? { triggerBox, optionsBox } : null; } +const clippedPopup = (page: PlaywrightPage) => page.locator('[data-slot="select-content"]'); + function pollOptionsOpenBelowTrigger(trigger: Locator, options: Locator) { return expect.poll(async () => { const box = await boxes(trigger, options); @@ -50,10 +52,9 @@ test.describe("Auto Router template select anchoring", () => { await trigger.scrollIntoViewIfNeeded(); await trigger.click(); - const options = page.getByRole("listbox"); - await expect(options).toBeVisible(); + await expect(page.getByRole("listbox")).toBeVisible(); - await pollOptionsOpenBelowTrigger(trigger, options).toBe(true); + await pollOptionsOpenBelowTrigger(trigger, clippedPopup(page)).toBe(true); }); test("keeps the trigger uncovered when the options open with no room below it", async ({ page }) => { @@ -62,9 +63,8 @@ test.describe("Auto Router template select anchoring", () => { await trigger.scrollIntoViewIfNeeded(); await trigger.click(); - const options = page.getByRole("listbox"); - await expect(options).toBeVisible(); + await expect(page.getByRole("listbox")).toBeVisible(); - await pollOptionsCoverTrigger(trigger, options).toBe(false); + await pollOptionsCoverTrigger(trigger, clippedPopup(page)).toBe(false); }); }); diff --git a/tests/e2e/ui/tests/users/searchUsers.spec.ts b/tests/e2e/ui/tests/users/searchUsers.spec.ts index 5e7e3e35b91..ee1a3f18f69 100644 --- a/tests/e2e/ui/tests/users/searchUsers.spec.ts +++ b/tests/e2e/ui/tests/users/searchUsers.spec.ts @@ -46,6 +46,6 @@ test.describe("Internal Users Search", () => { await page.getByTestId("users-filter-sso-id").fill("e2e-sso-id-that-matches-nobody"); await page.getByTestId("filter-drawer-apply").click(); - await expect(userRows(page)).toHaveCount(0, { timeout: 30_000 }); + await expect(page.getByRole("row").filter({ hasText: "No results" })).toBeVisible({ timeout: 30_000 }); }); }); From 94f6827530b6469069ee322a51d1958678b3bd31 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 15:36:35 -0700 Subject: [PATCH 112/120] test(e2e): drop redundant SCIM key cleanup, the throwaway db is the teardown --- tests/e2e/ui/tests/settings/scim.spec.ts | 15 +-------------- 1 file changed, 1 insertion(+), 14 deletions(-) diff --git a/tests/e2e/ui/tests/settings/scim.spec.ts b/tests/e2e/ui/tests/settings/scim.spec.ts index 5489111a580..d7dd4248f50 100644 --- a/tests/e2e/ui/tests/settings/scim.spec.ts +++ b/tests/e2e/ui/tests/settings/scim.spec.ts @@ -2,7 +2,6 @@ import { test, expect, Page as PlaywrightPage } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; import { Page } from "../../fixtures/pages"; import { navigateToPage } from "../../helpers/navigation"; -import { masterKey } from "../../helpers/traffic"; const rootPath = (): string => process.env.SERVER_ROOT_PATH ?? ""; @@ -26,13 +25,11 @@ test.describe("Admin Settings - SCIM", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); test("Create SCIM Token shows the token once and offers to create another", async ({ page }) => { - const token = await createScimTokenViaUi(page, `e2e-scim-ui-${Date.now()}`); + await createScimTokenViaUi(page, `e2e-scim-ui-${Date.now()}`); await page.getByRole("button", { name: "Create Another Token" }).click(); await expect(page.getByRole("button", { name: "Create SCIM Token" })).toBeVisible(); await expect(page.getByText(/copy this token now/i)).toBeHidden(); - - await deleteKey(page, token); }); test("a UI-minted SCIM token authorizes the SCIM API", async ({ page, request }) => { @@ -52,15 +49,5 @@ test.describe("Admin Settings - SCIM", () => { const body = await res.json(); expect(body.schemas, "SCIM answers with a ListResponse").toContain("urn:ietf:params:scim:api:messages:2.0:ListResponse"); expect(Array.isArray(body.Resources), "SCIM ListResponse carries a Resources array").toBe(true); - - await deleteKey(page, token); }); }); - -async function deleteKey(page: PlaywrightPage, key: string): Promise { - const res = await page.request.post(`${rootPath()}/key/delete`, { - headers: { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json" }, - data: { keys: [key] }, - }); - expect(res.ok(), `cleanup of the SCIM token failed (${res.status()})`).toBe(true); -} From 5eff708d0fd53fc474627b674ffe143f494d3aa1 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 31 Aug 2026 15:49:03 -0700 Subject: [PATCH 113/120] fix(proxy): keep persisted spend another pod has not incremented in the window seed The batch start told the seed which LiteLLM_SpendLogs rows were its own, but using it as a hard cutoff also dropped rows another pod had already persisted. Those rows are only repaid by that pod's own increment, so if it died first the window row stayed permanently under the recorded spend. The seed now reads both sums in one scan and takes off this batch's own spend, flooring at the pre-batch total for the case where its log rows have not landed yet. Redis payloads keep an empty request_ids so a leader from before the field was dropped can still merge what it pops during a rolling deploy. Claude-Session: https://claude.ai/code/session_01QvQzYztinxj8ZuD5YxbVdL --- .../proxy/db/budget_window_spend_writer.py | 95 +++++++++------ .../redis_update_buffer.py | 10 +- .../window_spend_update_queue.py | 37 +++++- .../test_redis_update_buffer.py | 34 ++++++ .../db/test_budget_window_spend_writer.py | 108 ++++++++++-------- 5 files changed, 197 insertions(+), 87 deletions(-) diff --git a/litellm/proxy/db/budget_window_spend_writer.py b/litellm/proxy/db/budget_window_spend_writer.py index 8b1ad0e24c7..8cf2f737063 100644 --- a/litellm/proxy/db/budget_window_spend_writer.py +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -7,17 +7,15 @@ instead of aggregating LiteLLM_SpendLogs every time a window counter goes cold (issue #35766). Raw SQL rather than the Prisma upsert helper because the conditional roll cannot be expressed through the query builder. -Seeding a row that does not exist yet reads LiteLLM_SpendLogs once, summing -only rows that started before the batch being flushed so neither source counts -the same request twice. Anything at or after that cutoff is owed by an -increment that still reaches the row, on this pod's next flush or another -pod's, so a row lags real spend by at most one flush interval of queued -increments: the same lag the SpendLogs aggregate it replaces (and every other -spend column) already has. A request whose increment is lost before it flushes, -which today means the pod dying, is missed by both sources and stays missing. +Seeding a row that does not exist yet reads LiteLLM_SpendLogs once and takes +off what the increments being flushed will add, so neither source counts the +same request twice. A row therefore lags real spend by at most one flush +interval of increments queued elsewhere: the same lag the SpendLogs aggregate +it replaces (and every other spend column) already has. """ from collections.abc import Sequence +from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Final, Protocol @@ -64,32 +62,45 @@ _ROLL_WINDOW_SPEND_SQL: Final = ( ) _SEED_FROM_SPEND_LOGS_KEY_SQL: Final = ( - 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' - "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') " - "AND \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')" + "SELECT COALESCE(SUM(spend), 0.0) AS total, " + "COALESCE(SUM(spend) FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) AS before_batch " + 'FROM "LiteLLM_SpendLogs" ' + "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" ) _SEED_FROM_SPEND_LOGS_TEAM_SQL: Final = ( - 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' - "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') " - "AND \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')" + "SELECT COALESCE(SUM(spend), 0.0) AS total, " + "COALESCE(SUM(spend) FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) AS before_batch " + 'FROM "LiteLLM_SpendLogs" ' + "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" ) _SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL: Final = ( - 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' + "SELECT COALESCE(SUM(spend), 0.0) AS total, COALESCE(SUM(spend), 0.0) AS before_batch " + 'FROM "LiteLLM_SpendLogs" ' "WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" ) _SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL: Final = ( - 'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" ' + "SELECT COALESCE(SUM(spend), 0.0) AS total, COALESCE(SUM(spend), 0.0) AS before_batch " + 'FROM "LiteLLM_SpendLogs" ' "WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" ) _UPSERT_TRANSACTION_TIMEOUT: Final = timedelta(seconds=60) +@dataclass(frozen=True, slots=True) +class WindowSeedTotals: + """The two sums a seed needs: everything persisted for the window, and the + part of it that predates the batch being flushed.""" + + total: float + before_batch: float + + class WindowSpendLogsAggregate(Protocol): - """Sums LiteLLM_SpendLogs for one entity between window_start and the + """Sums LiteLLM_SpendLogs for one entity since window_start, split at the batch's earliest request. Injected so the flush can be exercised without a database and so the @@ -103,18 +114,18 @@ class WindowSpendLogsAggregate(Protocol): entity_id: str, window_start: datetime, batch_started_at: datetime | None, - ) -> float | None: ... + ) -> WindowSeedTotals | None: ... -async def spend_logs_total_before_batch( +async def spend_logs_seed_totals( prisma_client: "PrismaClient", entity_type: str, entity_id: str, window_start: datetime, batch_started_at: datetime | None, -) -> float | None: - """LiteLLM_SpendLogs spend for one entity since window_start, stopping - before the requests the increments being flushed already cover. +) -> WindowSeedTotals | None: + """LiteLLM_SpendLogs spend for one entity since window_start, both in full + and up to the start of the batch being flushed, in one scan. The spend log writer drains its own queue on a ~2s poll whenever anything is queued, while window increments flush on the much slower batch tick, so @@ -122,11 +133,12 @@ async def spend_logs_total_before_batch( already in the table. Counting them in the seed and again in the increment is what made a fresh row land at twice the true spend. - Every log row at or after the cutoff belongs to a request whose own - increment still reaches this row, on this pod's next flush or another pod's, - so bounding the sum by time needs nothing from the request itself. Without a - known start the whole window is summed: that can only over-count once, which - enforcement tolerates, whereas under-counting is a budget bypass. + Both halves are needed because neither is safe alone: the full sum + double-counts this batch, and the sum before the batch drops spend another + pod has already persisted but not yet incremented. _seed_base picks between + them. Without a known batch start the two are the same sum, so the seed + counts everything: that can only over-count once, which enforcement + tolerates, whereas under-counting is a budget bypass. """ if entity_type == Litellm_EntityType.KEY.value: bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_KEY_SQL, _SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL @@ -145,8 +157,11 @@ async def spend_logs_total_before_batch( ) ) if not rows: - return 0.0 - return float(rows[0].get("total") or 0.0) + return WindowSeedTotals(total=0.0, before_batch=0.0) + return WindowSeedTotals( + total=float(rows[0].get("total") or 0.0), + before_batch=float(rows[0].get("before_batch") or 0.0), + ) def _exclusion_upper_bound(started_at: datetime) -> datetime: @@ -186,19 +201,33 @@ async def _seed_base_for_missing_row( This is the LiteLLM_SpendLogs aggregate the window counter reseed runs on every cold counter today, but here it runs once per window lifetime and off - the request path, and it stops before the queued increments so they are + the request path, and it discounts the queued increments so they are counted once. """ if _primary_key(transaction) in existing_primary_keys: return 0.0 - base: Final = await spend_logs_aggregate( + totals: Final = await spend_logs_aggregate( prisma_client=prisma_client, entity_type=transaction["entity_type"], entity_id=transaction["entity_id"], window_start=datetime.fromisoformat(transaction["window_start"]).replace(tzinfo=timezone.utc), batch_started_at=_transaction_started_at(transaction), ) - return float(base or 0.0) + if totals is None: + return 0.0 + return _seed_base(totals=totals, batch_spend=transaction["spend"]) + + +def _seed_base(totals: WindowSeedTotals, batch_spend: float) -> float: + """What the window already held before the increments about to be applied. + + Subtracting the batch's own spend from the full sum keeps every other + request in the seed, including the ones another pod persisted and has not + incremented yet, which a plain cutoff would drop for good if that pod died. + When this batch's own log rows have not landed yet the subtraction takes + spend that was never counted, so the sum before the batch is the floor. + """ + return max(totals.total - batch_spend, totals.before_batch) def _transaction_started_at(transaction: WindowSpendTransaction) -> datetime | None: @@ -232,7 +261,7 @@ def _upsert_params( async def commit_window_spend_updates( prisma_client: "PrismaClient", transactions: Sequence[WindowSpendTransaction], - spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_total_before_batch, + spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_seed_totals, ) -> None: """Apply aggregated window increments to LiteLLM_BudgetWindowSpend. diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 4dd23270bf8..c06f2e04aca 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -46,6 +46,7 @@ from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdate from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WindowSpendTransaction, WindowSpendUpdateQueue, + to_wire_payload, ) from litellm.secret_managers.main import str_to_bool from litellm.types.caching import ( @@ -298,7 +299,7 @@ class RedisUpdateBuffer: ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE, ), ( - window_spend_update_transactions, + tuple(map(to_wire_payload, window_spend_update_transactions)), REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_WINDOW_SPEND_UPDATE_QUEUE, ), @@ -484,7 +485,12 @@ class RedisUpdateBuffer: (daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY), (daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY), (daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY), - (window_spend_update_transactions, REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY), + ( + None + if window_spend_update_transactions is None + else tuple(map(to_wire_payload, window_spend_update_transactions)), + REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, + ), ) rpush_list: Final = tuple( diff --git a/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py index 43b069fd8c2..372a6666c02 100644 --- a/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py @@ -27,10 +27,11 @@ class WindowSpendTransaction(TypedDict): transaction survives the JSON round trip through the Redis buffer. started_at is the earliest request start in the batch. The one-time seed for - a window that has no row yet sums only LiteLLM_SpendLogs rows that started - before it, because the spend log writer flushes on its own ~2s poll and will - usually have persisted this batch's rows before the window queue flushes; - without the bound the seed and the increment would each count them. + a window that has no row yet uses it to tell this batch's own + LiteLLM_SpendLogs rows from everything else, because the spend log writer + flushes on its own ~2s poll and will usually have persisted this batch's + rows before the window queue flushes; without that split the seed and the + increment would each count them. """ entity_type: ReadOnly[str] @@ -41,6 +42,34 @@ class WindowSpendTransaction(TypedDict): started_at: ReadOnly[str | None] +class WindowSpendWirePayload(WindowSpendTransaction): + """How an increment is encoded in the shared Redis buffer. + + request_ids is dead weight here: workers built before this field was + dropped index it while merging whatever they pop, and the pop is + destructive, so a leader still running one of those during a rolling deploy + would raise on a payload without the key and lose those increments. It is + always empty, which only makes such a leader seed without exclusions. + + TODO: remove once no supported version reads it, i.e. one release after the + field stopped being written. + """ + + request_ids: ReadOnly[Sequence[str]] + + +def to_wire_payload(transaction: WindowSpendTransaction) -> WindowSpendWirePayload: + return WindowSpendWirePayload( + entity_type=transaction["entity_type"], + entity_id=transaction["entity_id"], + window_duration=transaction["window_duration"], + window_start=transaction["window_start"], + spend=transaction["spend"], + started_at=transaction.get("started_at"), + request_ids=(), + ) + + def to_naive_utc(value: datetime) -> datetime: """LiteLLM_BudgetWindowSpend.window_start is TIMESTAMP(3), which holds naive UTC.""" if value.tzinfo is None: diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 99ac1fe8b50..8f3508fc4e9 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -523,10 +523,44 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_up "window_start": "2026-08-01T00:00:00.000000", "spend": 1.25, "started_at": "2026-08-10T12:00:00.000000", + "request_ids": [], } ] +@pytest.mark.asyncio +async def test_budget_window_payloads_keep_request_ids_for_older_workers(redis_update_buffer, mock_redis_cache): + """A leader from before the field was dropped indexes request_ids while + merging what it popped, and the pop is destructive, so a payload without + the key would cost a rolling deploy those increments.""" + from datetime import datetime, timezone + + from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendUpdateQueue, + build_window_spend_transaction, + ) + + mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1]) + window_queue = WindowSpendUpdateQueue() + await window_queue.add_update( + build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=1.25, + ) + ) + + await redis_update_buffer.restore_transactions_to_redis( + window_spend_update_transactions=await window_queue.flush_and_get_aggregated_window_spend_transactions(), + ) + + rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] + restored = json.loads(rpush_list[0]["values"][0]) + assert [payload["request_ids"] for payload in restored] == [[]] + + @pytest.mark.asyncio async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpush_failure( redis_update_buffer, mock_redis_cache diff --git a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py index 74b71f63401..130f0c56ccf 100644 --- a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py +++ b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py @@ -6,9 +6,10 @@ from typing import Any import pytest from litellm.proxy.db.budget_window_spend_writer import ( + WindowSeedTotals, commit_window_spend_updates, roll_window_spend_row, - spend_logs_total_before_batch, + spend_logs_seed_totals, ) from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( build_window_spend_transaction, @@ -70,10 +71,15 @@ class _FakePrismaClient: class _RecordingAggregate: - """Stands in for the LiteLLM_SpendLogs seed aggregate.""" + """Stands in for the LiteLLM_SpendLogs seed aggregate. before_batch + defaults to the full total, the state where none of this batch's own log + rows have been persisted yet.""" - def __init__(self, value: float = 5.0) -> None: - self.value = value + def __init__(self, total: float = 5.0, before_batch: float | None = None) -> None: + self.totals = WindowSeedTotals( + total=total, + before_batch=total if before_batch is None else before_batch, + ) self.calls: list[dict[str, Any]] = [] async def __call__( @@ -83,7 +89,7 @@ class _RecordingAggregate: entity_id: str, window_start: datetime, batch_started_at: datetime | None, - ) -> float | None: + ) -> WindowSeedTotals | None: self.calls.append( { "entity_type": entity_type, @@ -92,13 +98,13 @@ class _RecordingAggregate: "batch_started_at": batch_started_at, } ) - return self.value + return self.totals class _SpendLogsFake: """Sums the LiteLLM_SpendLogs rows (request_id, spend, startTime) it holds, - honouring the cutoff exactly as the real aggregate's - startTime < bound does.""" + splitting them at the batch start exactly as the real aggregate's + SUM(...) FILTER (WHERE startTime < bound) does.""" def __init__(self, rows: tuple[tuple[str, float, datetime], ...]) -> None: self.rows = rows @@ -110,11 +116,14 @@ class _SpendLogsFake: entity_id: str, window_start: datetime, batch_started_at: datetime | None, - ) -> float | None: - return math.fsum( - spend - for _request_id, spend, started_at in self.rows - if batch_started_at is None or started_at < batch_started_at + ) -> WindowSeedTotals | None: + return WindowSeedTotals( + total=math.fsum(spend for _request_id, spend, _started_at in self.rows), + before_batch=math.fsum( + spend + for _request_id, spend, started_at in self.rows + if batch_started_at is None or started_at < batch_started_at + ), ) @@ -151,7 +160,7 @@ async def test_missing_row_is_seeded_from_spend_logs_once(): existed, so a brand new primary key inserts the SpendLogs total plus this increment.""" db = _FakeDB(existing_rows=[]) - aggregate = _RecordingAggregate(value=5.0) + aggregate = _RecordingAggregate(total=5.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -178,7 +187,7 @@ async def test_existing_row_is_never_reseeded(): """The seed is a full LiteLLM_SpendLogs scan; running it for a row that is already maintained would both cost a scan and double count.""" db = _FakeDB(existing_rows=[_existing("key", "k1", "30d")]) - aggregate = _RecordingAggregate(value=5.0) + aggregate = _RecordingAggregate(total=5.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -195,7 +204,7 @@ async def test_existing_row_is_never_reseeded(): @pytest.mark.asyncio async def test_seed_runs_only_for_the_primary_keys_that_are_missing(): db = _FakeDB(existing_rows=[_existing("key", "k1", "30d")]) - aggregate = _RecordingAggregate(value=5.0) + aggregate = _RecordingAggregate(total=5.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -218,7 +227,7 @@ async def test_insert_spend_and_increment_differ_only_when_a_row_is_seeded(): """The conflict arm adds the increment alone so two pods that both seed the same new window cannot add the SpendLogs base twice.""" db = _FakeDB(existing_rows=[]) - aggregate = _RecordingAggregate(value=9.0) + aggregate = _RecordingAggregate(total=9.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -255,7 +264,7 @@ async def test_upsert_sql_adds_for_a_current_window_and_replaces_for_a_newer_one @pytest.mark.asyncio async def test_upsert_never_interpolates_values_into_the_sql(): db = _FakeDB(existing_rows=[]) - aggregate = _RecordingAggregate(value=0.0) + aggregate = _RecordingAggregate(total=0.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -273,7 +282,7 @@ async def test_upserts_are_ordered_by_primary_key_then_window_start(): """Cross-pod lock ordering, plus an older window must be applied before the roll that supersedes it or the roll would be undone.""" db = _FakeDB(existing_rows=[]) - aggregate = _RecordingAggregate(value=0.0) + aggregate = _RecordingAggregate(total=0.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -301,7 +310,7 @@ async def test_upserts_are_ordered_by_primary_key_then_window_start(): @pytest.mark.asyncio async def test_existing_row_lookup_sends_every_primary_key_as_array_params(): db = _FakeDB(existing_rows=[]) - aggregate = _RecordingAggregate(value=0.0) + aggregate = _RecordingAggregate(total=0.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -339,9 +348,7 @@ async def test_unknown_entity_type_contributes_no_seed(): anything else starts from its increment alone.""" db = _FakeDB(existing_rows=[]) - async def no_such_column( - prisma_client, entity_type, entity_id, window_start, batch_started_at - ): + async def no_such_column(prisma_client, entity_type, entity_id, window_start, batch_started_at): return None await commit_window_spend_updates( @@ -396,7 +403,7 @@ async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_o @pytest.mark.asyncio async def test_seed_receives_the_batch_earliest_start_as_its_cutoff(): db = _FakeDB(existing_rows=[]) - aggregate = _RecordingAggregate(value=0.0) + aggregate = _RecordingAggregate(total=0.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -410,7 +417,7 @@ async def test_seed_receives_the_batch_earliest_start_as_its_cutoff(): @pytest.mark.asyncio async def test_seed_passes_no_start_bound_when_the_batch_has_none(): db = _FakeDB(existing_rows=[]) - aggregate = _RecordingAggregate(value=0.0) + aggregate = _RecordingAggregate(total=0.0) await commit_window_spend_updates( prisma_client=_FakePrismaClient(db), @@ -463,15 +470,19 @@ async def test_new_row_still_covers_spend_that_predates_the_batch(): @pytest.mark.asyncio -async def test_seed_skips_logs_from_requests_this_batch_never_saw(): - """A concurrent request on another pod can land its spend log before this - pod seeds the row. Its increment is still queued over there, so the cutoff - has to drop it from the seed even though this batch has no way to know its - id; counting it here and again on that pod's flush is the double count the - old id list could not catch.""" +async def test_seed_keeps_spend_another_pod_persisted_after_this_batch_started(): + """A concurrent request on another pod can land its spend log after this + batch started but before this pod seeds the row. Dropping it on a plain + time cutoff would lose that spend for the rest of the window if that pod + died before flushing its increment, so the seed takes off only this batch's + own spend and keeps everything else.""" db = _FakeDB(existing_rows=[]) spend_logs = _SpendLogsFake( - rows=(("older", 0.5, BEFORE_BATCH), ("other-pod", 0.25, BATCH_STARTED_AT + timedelta(seconds=1))), + rows=( + ("older", 0.5, BEFORE_BATCH), + ("mine", 0.000047, BATCH_STARTED_AT), + ("other-pod", 0.25, BATCH_STARTED_AT + timedelta(seconds=1)), + ), ) await commit_window_spend_updates( @@ -481,7 +492,7 @@ async def test_seed_skips_logs_from_requests_this_batch_never_saw(): ) ((_, params),) = db.batcher.calls - assert params[INSERT_SPEND] == pytest.approx(0.500047) + assert params[INSERT_SPEND] == pytest.approx(0.750047) @pytest.mark.asyncio @@ -506,10 +517,10 @@ async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet(): "entity_type, expected_column", [("key", "api_key = $1"), ("team", "team_id = $1")], ) -async def test_seed_aggregate_sql_stops_at_the_batch_start(entity_type, expected_column): - db = _FakeDB(existing_rows=[{"total": 1.25}]) +async def test_seed_aggregate_sql_splits_the_window_at_the_batch_start(entity_type, expected_column): + db = _FakeDB(existing_rows=[{"total": 1.25, "before_batch": 0.75}]) - total = await spend_logs_total_before_batch( + totals = await spend_logs_seed_totals( prisma_client=_FakePrismaClient(db), entity_type=entity_type, entity_id="e1", @@ -517,11 +528,11 @@ async def test_seed_aggregate_sql_stops_at_the_batch_start(entity_type, expected batch_started_at=BATCH_STARTED_AT, ) - assert total == pytest.approx(1.25) + assert totals == WindowSeedTotals(total=1.25, before_batch=0.75) ((query, params),) = db.query_raw_calls normalized = " ".join(query.split()) assert expected_column in normalized - assert "AND \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')" in normalized + assert "FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC'))" in normalized assert 'FROM "LiteLLM_SpendLogs"' in normalized # startTime is TIMESTAMP(3): the bound is floored to the second so the # batch's own earliest row cannot round under it. @@ -532,12 +543,13 @@ async def test_seed_aggregate_sql_stops_at_the_batch_start(entity_type, expected @pytest.mark.asyncio async def test_seed_aggregate_sums_the_whole_window_without_a_start_bound(): - """A batch with no known start cannot place the cutoff, so the seed counts - everything; at worst that over-counts one batch, which enforcement - tolerates, where under-counting is a budget bypass.""" - db = _FakeDB(existing_rows=[{"total": 1.25}]) + """A batch with no known start cannot place the split, so both halves are + the same sum and the seed counts everything; at worst that over-counts one + batch, which enforcement tolerates, where under-counting is a budget + bypass.""" + db = _FakeDB(existing_rows=[{"total": 1.25, "before_batch": 1.25}]) - total = await spend_logs_total_before_batch( + totals = await spend_logs_seed_totals( prisma_client=_FakePrismaClient(db), entity_type="key", entity_id="e1", @@ -545,7 +557,7 @@ async def test_seed_aggregate_sums_the_whole_window_without_a_start_bound(): batch_started_at=None, ) - assert total == pytest.approx(1.25) + assert totals == WindowSeedTotals(total=1.25, before_batch=1.25) ((query, params),) = db.query_raw_calls assert '"startTime" <' not in query assert params == ("e1", WINDOW_A) @@ -555,7 +567,7 @@ async def test_seed_aggregate_sums_the_whole_window_without_a_start_bound(): async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs_column(): db = _FakeDB(existing_rows=[]) - total = await spend_logs_total_before_batch( + totals = await spend_logs_seed_totals( prisma_client=_FakePrismaClient(db), entity_type="user", entity_id="u1", @@ -563,7 +575,7 @@ async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs batch_started_at=None, ) - assert total is None + assert totals is None assert db.query_raw_calls == [] @@ -571,7 +583,7 @@ async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs async def test_seed_aggregate_treats_an_entity_with_no_rows_as_zero(): db = _FakeDB(existing_rows=[]) - total = await spend_logs_total_before_batch( + totals = await spend_logs_seed_totals( prisma_client=_FakePrismaClient(db), entity_type="key", entity_id="k-unknown", @@ -579,4 +591,4 @@ async def test_seed_aggregate_treats_an_entity_with_no_rows_as_zero(): batch_started_at=None, ) - assert total == 0.0 + assert totals == WindowSeedTotals(total=0.0, before_batch=0.0) From b5c156a7d5b0ec911c34519a0846890ce6489566 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 31 Aug 2026 15:49:49 -0700 Subject: [PATCH 114/120] style(proxy): drop the explicit return None update_database no longer needs Reverting the function to -> None left two bare `return None` statements that RET501 rejects now that None is the only value it can return. Claude-Session: https://claude.ai/code/session_01QvQzYztinxj8ZuD5YxbVdL --- litellm/proxy/db/db_spend_update_writer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index b3fd2c3f22c..202a95ba29b 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -232,7 +232,7 @@ class DBSpendUpdateWriter: team_id, ) if ProxyUpdateSpend.disable_spend_updates() is True: - return None + return if token is not None and isinstance(token, str) and token.startswith("sk-"): hashed_token = hash_token(token=token) else: @@ -318,7 +318,7 @@ class DBSpendUpdateWriter: org_id, end_user_id, ) - return None + return async def _enqueue_tool_usage_transaction( self, From 296bde0d0d70c10ab2f5facdadbeb8e10ab9245c Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 31 Aug 2026 15:50:04 -0700 Subject: [PATCH 115/120] feat(complexity-router): add classification_mode to skip classifier on continuation turns (#38861) --- .../complexity_router/complexity_router.py | 57 +++- .../complexity_router/config.py | 15 + litellm/types/utils.py | 5 + .../router_strategy/test_complexity_router.py | 259 ++++++++++++++++++ .../LogDetailsDrawer/RoutingDecisionCard.tsx | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 9 +- 6 files changed, 338 insertions(+), 8 deletions(-) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 2f4305756e9..1aeedaa97b2 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -479,6 +479,33 @@ def _last_human_ask_index( ) +def _newest_turn_is_human_ask( + messages: Sequence[Mapping[str, object]] | None, + marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS, +) -> bool: + """Whether the request's newest turn carries a real human ask, i.e. this is a new ask rather + than an agent loop's continuation traffic. + + Anchored on `_last_human_ask_index` so every surface's plumbing reads as a continuation: + chat-completions tool turns are role=tool, Messages-surface tool_result turns flatten to empty + human text, and a hybrid turn carrying an ask alongside a tool_result still counts as an ask. + Compared against the newest non-system message rather than the raw tail, because Claude Code + appends a system-role reminder after the human turn; that trailing plumbing is neither an ask + nor loop traffic and must not turn a fresh ask into a continuation. An unreadable request (no + messages) is treated as a continuation: there is no ask to classify, which is the same reading + `_extract_current_ask_and_system_prompt` gives it downstream. + """ + if not messages: + return False + newest_non_system: Final = next( + (index for index in range(len(messages) - 1, -1, -1) if messages[index].get("role") != "system"), + None, + ) + if newest_non_system is None: + return False + return _last_human_ask_index(messages, marker_pairs) == newest_non_system + + def _iter_system_scope_texts( body_system: object, messages: Sequence[Mapping[str, object]], @@ -2247,14 +2274,18 @@ class ComplexityRouter(CustomLogger): @property def _uses_tier_pin(self) -> bool: - return bool(self.config.session_affinity and not self.config.plugins) + """classification_mode 'user_turn' implies the tier pin machinery: the pin write after each + pinnable classification is what gives a continuation a held decision to replay.""" + return bool( + (self.config.session_affinity or self.config.classification_mode == "user_turn") and not self.config.plugins + ) @property def _uses_deployment_pin(self) -> bool: - """session_affinity implies the deployment pin: a session frozen onto one model + """The tier pin implies the deployment pin: a session frozen onto one model group but load-balanced across its deployments would still go cache-cold, which is the exact failure both flags exist to prevent.""" - return bool((self.config.deployment_affinity or self.config.session_affinity) and not self.config.plugins) + return bool(self.config.deployment_affinity and not self.config.plugins) or self._uses_tier_pin def _with_session_deployment_affinity( self, response: PreRoutingHookResponse | None @@ -2282,6 +2313,11 @@ class ComplexityRouter(CustomLogger): pins the model chosen on the session's first turn and reuses it for every later turn, skipping classification entirely. Otherwise delegates to `_classify_and_route`. + When `classification_mode` is 'user_turn', the same pin is replayed only on + continuation turns (an agent loop's tool traffic); a new human ask always falls + through to classification, so the session can still move tiers between asks. + With both knobs on, session_affinity's pin-first behavior wins. + Skipped entirely when `plugins` are configured: reusing a stale pin would bypass the plugin pipeline on every turn after the first, since a pinned model was never re-checked against a policy plugin whose decision can change between turns (e.g. a @@ -2305,7 +2341,13 @@ class ComplexityRouter(CustomLogger): session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs) if use_session_affinity else None cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None - if cache_key is not None: + # In 'user_turn' mode a held pin is replayed only on continuation turns; a new human + # ask falls through and re-classifies. session_affinity restores pin-first for asks too. + pin_replay_allowed: Final = bool(self.config.session_affinity) or not _newest_turn_is_human_ask( + resolved_messages, self._reminder_markers + ) + + if cache_key is not None and pin_replay_allowed: pinned_value: Final = await self.litellm_router_instance.cache.async_get_cache(key=cache_key) pinned_pin: Final = _parse_session_affinity_pin(pinned_value) if pinned_pin is not None: @@ -2354,10 +2396,11 @@ class ComplexityRouter(CustomLogger): kwargs_metadata: Final = request_kwargs.setdefault("metadata", {}) if isinstance(kwargs_metadata, dict): kwargs_metadata[ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] = routed_model + replay_cause: Final[RoutingDecisionCause] = ( + "session_affinity_pin" if self.config.session_affinity else "user_turn_continuation" + ) cause: RoutingDecisionCause = ( - "plan_mode" - if plan_floored - else ("session_affinity_escalation" if escalated else "session_affinity_pin") + "plan_mode" if plan_floored else ("session_affinity_escalation" if escalated else replay_cause) ) verbose_router_logger.info( "ComplexityRouter: routing decision cause=%s, routed_model=%s", cause, routed_model diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 335de11e669..0abe962edf1 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -839,6 +839,21 @@ class ComplexityRouterConfig(BaseModel): description="Minimum cosine similarity for a semantic keyword match", ) + classification_mode: Literal["every_request", "user_turn"] = Field( + default="every_request", + description=( + "When to run the complexity classifier. 'every_request' (the default) classifies every " + "inference request, including the tool-result continuation turns of an agentic loop. " + "'user_turn' classifies only requests whose newest turn is a new human ask and replays " + "the session's held routing decision on continuation turns, which cuts classifier " + "spend and eliminates mid-loop model switches. Continuations with no held decision to " + "replay (no resolvable session_id, expired pin, fresh restart) still classify. Unlike " + "session_affinity, a new human ask always re-classifies, so a session can still move " + "tiers between asks. Suppressed when plugins are configured, for the same reason " + "session_affinity is: a replayed decision would bypass the plugin pipeline." + ), + ) + # Session affinity: pin the first turn's routed model for the rest of the session session_affinity: bool = Field( default=False, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4bf8289d725..14749ecde6a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2842,6 +2842,11 @@ RoutingDecisionCause = Literal[ "housekeeping", "session_affinity_pin", "session_affinity_escalation", + # classification_mode 'user_turn': the request is an agent loop's continuation turn (no new + # human ask), so the session's held routing decision was replayed and the classifier was never + # called. Distinct from "session_affinity_pin", which reports the session_affinity flag pinning + # every turn including new asks; this cause only appears when session_affinity is off. + "user_turn_continuation", "default_fallback", "keyword", "quality_tier", diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index eee4e9aa185..97f60ec8e57 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -4425,6 +4425,265 @@ class _DummyPlugin: return context +class TestClassificationMode: + """Test classification_mode='user_turn': classify only requests whose newest turn is a new + human ask; tool-loop continuation turns replay the session's held routing decision.""" + + REASONING_ASK = { + "role": "user", + "content": "Let's think step by step and reason through this problem carefully.", + } + SIMPLE_ASK = {"role": "user", "content": "Hello!"} + ASSISTANT_ANSWER = {"role": "assistant", "content": "the answer"} + TOOL_CALL_1 = { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "read_file", "arguments": "{}"}}], + } + TOOL_RESULT_1 = {"role": "tool", "tool_call_id": "call_1", "content": "file contents"} + TOOL_CALL_2 = { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_2", "type": "function", "function": {"name": "run_tests", "arguments": "{}"}}], + } + TOOL_RESULT_2 = {"role": "tool", "tool_call_id": "call_2", "content": "3 passed"} + + @pytest.fixture + def user_turn_config(self, basic_config) -> dict: + return {**basic_config, "classification_mode": "user_turn"} + + @staticmethod + def _request_kwargs(session_id: str) -> dict: + return {"metadata": {"session_id": session_id}} + + def _router(self, mock_router_instance, config: dict) -> ComplexityRouter: + mock_router_instance.cache = DualCache() + return ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=config, + ) + + def _tool_loop_turns(self) -> list[list[dict]]: + return [ + [self.REASONING_ASK], + [self.REASONING_ASK, self.TOOL_CALL_1, self.TOOL_RESULT_1], + [self.REASONING_ASK, self.TOOL_CALL_1, self.TOOL_RESULT_1, self.TOOL_CALL_2, self.TOOL_RESULT_2], + ] + + def test_default_mode_is_every_request(self, complexity_router): + assert complexity_router.config.classification_mode == "every_request" + + def test_invalid_classification_mode_rejected(self, mock_router_instance, basic_config): + with pytest.raises(ValidationError): + ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**basic_config, "classification_mode": "sometimes"}, + ) + + @pytest.mark.asyncio + async def test_user_turn_mode_classifies_tool_loop_once(self, mock_router_instance, user_turn_config): + """The mutation check: a 3-request tool loop drives exactly one classification, and both + continuation turns hold the classified model under the user_turn_continuation cause.""" + router = self._router(mock_router_instance, user_turn_config) + with patch.object(router, "_classify_and_route", wraps=router._classify_and_route) as spy: + responses = [ + await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("loop-1"), messages=turn + ) + for turn in self._tool_loop_turns() + ] + assert spy.call_count == 1 + assert [r.model for r in responses] == ["o1-preview", "o1-preview", "o1-preview"] + assert [r.routing_decision["cause"] for r in responses[1:]] == [ + "user_turn_continuation", + "user_turn_continuation", + ] + + @pytest.mark.asyncio + async def test_every_request_default_classifies_every_tool_loop_turn(self, mock_router_instance, basic_config): + """Pins today's default: every request classifies, including tool-loop continuations.""" + router = self._router(mock_router_instance, basic_config) + with patch.object(router, "_classify_and_route", wraps=router._classify_and_route) as spy: + responses = [ + await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("loop-2"), messages=turn + ) + for turn in self._tool_loop_turns() + ] + assert spy.call_count == 3 + assert [r.model for r in responses] == ["o1-preview", "o1-preview", "o1-preview"] + assert all(r.routing_decision["cause"] != "user_turn_continuation" for r in responses) + + @pytest.mark.asyncio + async def test_continuation_without_session_id_still_classifies(self, mock_router_instance, user_turn_config): + """No resolvable session id means no held decision to replay, so every request classifies.""" + router = self._router(mock_router_instance, user_turn_config) + with patch.object(router, "_classify_and_route", wraps=router._classify_and_route) as spy: + responses = [ + await router.async_pre_routing_hook(model="test-model", request_kwargs={}, messages=turn) + for turn in self._tool_loop_turns() + ] + assert spy.call_count == 3 + assert [r.model for r in responses] == ["o1-preview", "o1-preview", "o1-preview"] + assert all(r.routing_decision["cause"] != "user_turn_continuation" for r in responses) + + @pytest.mark.asyncio + async def test_plugins_suppress_user_turn_gate(self, mock_router_instance, basic_config): + """A replayed decision would bypass the plugin pipeline, so plugins force every request + through _classify_and_route, exactly as they do for session_affinity.""" + router = self._router( + mock_router_instance, + {**basic_config, "classification_mode": "user_turn", "plugins": [_DummyPlugin()]}, + ) + with patch.object(router, "_classify_and_route", wraps=router._classify_and_route) as spy: + responses = [ + await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("loop-3"), messages=turn + ) + for turn in self._tool_loop_turns() + ] + assert spy.call_count == 3 + assert [r.model for r in responses] == ["o1-preview", "o1-preview", "o1-preview"] + assert all(r.routing_decision["cause"] != "user_turn_continuation" for r in responses) + + @pytest.mark.asyncio + async def test_new_human_ask_reclassifies_and_repins(self, mock_router_instance, user_turn_config): + """Unlike session_affinity, a new human ask never short-circuits on the pin: the session + re-classifies, moves tier, and the moved decision becomes the next held decision.""" + router = self._router(mock_router_instance, user_turn_config) + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("s-repin"), messages=[self.REASONING_ASK] + ) + second = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._request_kwargs("s-repin"), + messages=[self.REASONING_ASK, self.ASSISTANT_ANSWER, self.SIMPLE_ASK], + ) + third = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._request_kwargs("s-repin"), + messages=[self.REASONING_ASK, self.ASSISTANT_ANSWER, self.SIMPLE_ASK, self.TOOL_CALL_1, self.TOOL_RESULT_1], + ) + assert first.model == "o1-preview" + assert second.model == "gpt-4o-mini" + assert third.model == "gpt-4o-mini" + assert third.routing_decision["cause"] == "user_turn_continuation" + + @pytest.mark.asyncio + async def test_new_ask_with_trailing_system_reminder_reclassifies(self, mock_router_instance, user_turn_config): + """Claude Code appends a system-role reminder after the human turn; that trailing plumbing + must not turn a new ask into a continuation, and a continuation turn carrying the same + trailing reminder stays a continuation.""" + router = self._router(mock_router_instance, user_turn_config) + reminder = {"role": "system", "content": "100 tokens left"} + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("s-reminder"), messages=[self.REASONING_ASK] + ) + second = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._request_kwargs("s-reminder"), + messages=[self.REASONING_ASK, self.ASSISTANT_ANSWER, self.SIMPLE_ASK, reminder], + ) + third = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._request_kwargs("s-reminder"), + messages=[ + self.REASONING_ASK, + self.ASSISTANT_ANSWER, + self.SIMPLE_ASK, + reminder, + self.TOOL_CALL_1, + self.TOOL_RESULT_1, + reminder, + ], + ) + assert first.model == "o1-preview" + assert second.model == "gpt-4o-mini" + assert second.routing_decision["cause"] != "user_turn_continuation" + assert third.model == "gpt-4o-mini" + assert third.routing_decision["cause"] == "user_turn_continuation" + + @pytest.mark.asyncio + async def test_escalation_keyword_turn_is_a_new_ask(self, mock_router_instance, user_turn_config): + """An escalation keyword arrives as human text, so the turn classifies and escalates + instead of replaying the held decision.""" + router = self._router(mock_router_instance, user_turn_config) + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("s-esc"), messages=[self.SIMPLE_ASK] + ) + second = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._request_kwargs("s-esc"), + messages=[self.SIMPLE_ASK, self.ASSISTANT_ANSWER, {"role": "user", "content": "LITELLM ESCALATE"}], + ) + assert first.model == "gpt-4o-mini" + assert second.model == "gpt-4o" + assert second.routing_decision["escalated"] is True + + @pytest.mark.asyncio + async def test_messages_surface_tool_result_shapes(self, mock_router_instance, user_turn_config): + """Messages-surface shapes: a tool_result-only user turn is a continuation, while an ask + riding alongside a tool_result in the same turn is a new ask.""" + router = self._router(mock_router_instance, user_turn_config) + tool_use = {"role": "assistant", "content": [{"type": "tool_use", "id": "x", "name": "t", "input": {}}]} + tool_result = {"type": "tool_result", "tool_use_id": "x", "content": "ok"} + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("s-msgs"), messages=[self.REASONING_ASK] + ) + pure = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._request_kwargs("s-msgs"), + messages=[self.REASONING_ASK, tool_use, {"role": "user", "content": [tool_result]}], + ) + hybrid = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._request_kwargs("s-msgs"), + messages=[ + self.REASONING_ASK, + tool_use, + {"role": "user", "content": [tool_result, {"type": "text", "text": "Hello!"}]}, + ], + ) + assert first.model == "o1-preview" + assert pure.model == "o1-preview" + assert pure.routing_decision["cause"] == "user_turn_continuation" + assert hybrid.model == "gpt-4o-mini" + + @pytest.mark.asyncio + async def test_session_affinity_wins_when_both_knobs_are_on(self, mock_router_instance, user_turn_config): + """With session_affinity also on, the pin short-circuits new asks too and keeps its own + cause, so the session stays on turn 1's model.""" + router = self._router(mock_router_instance, {**user_turn_config, "session_affinity": True}) + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("s-both"), messages=[self.REASONING_ASK] + ) + second = await router.async_pre_routing_hook( + model="test-model", + request_kwargs=self._request_kwargs("s-both"), + messages=[self.REASONING_ASK, self.ASSISTANT_ANSWER, self.SIMPLE_ASK], + ) + assert first.model == "o1-preview" + assert second.model == "o1-preview" + assert second.routing_decision["cause"] == "session_affinity_pin" + + def test_user_turn_mode_enables_tier_and_deployment_pins(self, mock_router_instance, basic_config): + """user_turn implies the tier pin machinery (the pin write is what gives a continuation + a held decision) and the tier pin implies the deployment pin; plugins suppress both.""" + default = self._router(mock_router_instance, basic_config) + enabled = self._router(mock_router_instance, {**basic_config, "classification_mode": "user_turn"}) + suppressed = self._router( + mock_router_instance, + {**basic_config, "classification_mode": "user_turn", "plugins": [_DummyPlugin()]}, + ) + assert default._uses_tier_pin is False + assert enabled._uses_tier_pin is True + assert enabled._uses_deployment_pin is True + assert suppressed._uses_tier_pin is False + assert suppressed._uses_deployment_pin is False + + class TestRoutingPlugins: """Test the `complexity_router_config.plugins` field: narrows the classified tier's candidate pool before a model is picked. Discussion: diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx index 8c77b2db630..d2aa20901f5 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx @@ -89,6 +89,7 @@ const CONSTANT_CAUSE_LABELS: Record = { semantic_keyword_match: "Semantic keyword match", session_affinity_pin: "Pinned to session", session_affinity_escalation: "Escalated from session pin", + user_turn_continuation: "Continuation turn, classifier skipped", quality_tier: "Quality tier mapping", bandit: "Adaptive bandit", default_fallback: "Default model, no route matched", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index cc56483ac85..140130adc4b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -34318,6 +34318,13 @@ export interface components { adaptive_eligible: "all" | "classified_tier"; /** @description Quality vs cost weights for adaptive selection (used when adaptive=True) */ adaptive_weights?: components["schemas"]["AdaptiveRouterWeights"]; + /** + * Classification Mode + * @description When to run the complexity classifier. 'every_request' (the default) classifies every inference request, including the tool-result continuation turns of an agentic loop. 'user_turn' classifies only requests whose newest turn is a new human ask and replays the session's held routing decision on continuation turns, which cuts classifier spend and eliminates mid-loop model switches. Continuations with no held decision to replay (no resolvable session_id, expired pin, fresh restart) still classify. Unlike session_affinity, a new human ask always re-classifies, so a session can still move tiers between asks. Suppressed when plugins are configured, for the same reason session_affinity is: a replayed decision would bypass the plugin pipeline. + * @default every_request + * @enum {string} + */ + classification_mode: "every_request" | "user_turn"; /** * Classification Prompt * @description Replaces the opening instructions of the LLM classifier rubric (the judging-criteria prose) for a custom tier set. The per-tier bullets and the trust-boundary paragraph telling the classifier to ignore tier requests embedded in quoted caller text are always appended after it and cannot be overridden. Requires tier_definitions; a built-in-tier router customizes its prompt via classifier_llm_config.system_prompt or classification_rubric instead. @@ -35590,7 +35597,7 @@ export interface components { * Cause * @enum {string} */ - cause?: "heuristic_scorer" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "session_affinity_pin" | "session_affinity_escalation" | "default_fallback" | "keyword" | "quality_tier" | "bandit"; + cause?: "heuristic_scorer" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit"; /** Classifier Cost */ classifier_cost?: number; /** Classifier Model */ From 859bd01ddab619367671a39ae39d46aa3d349d3c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 15:57:49 -0700 Subject: [PATCH 116/120] fix(e2e): assert the users table's own empty-state copy UsersTable overrides DataTable's default noDataMessage with its own EmptyState, so the row reads "No users found" rather than "No results". Assert that, and pair it with the seeded user being absent so the check cannot pass while the filter silently does nothing. --- tests/e2e/ui/tests/users/searchUsers.spec.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/e2e/ui/tests/users/searchUsers.spec.ts b/tests/e2e/ui/tests/users/searchUsers.spec.ts index ee1a3f18f69..fa8f32764e8 100644 --- a/tests/e2e/ui/tests/users/searchUsers.spec.ts +++ b/tests/e2e/ui/tests/users/searchUsers.spec.ts @@ -46,6 +46,7 @@ test.describe("Internal Users Search", () => { await page.getByTestId("users-filter-sso-id").fill("e2e-sso-id-that-matches-nobody"); await page.getByTestId("filter-drawer-apply").click(); - await expect(page.getByRole("row").filter({ hasText: "No results" })).toBeVisible({ timeout: 30_000 }); + await expect(page.getByText("No users found")).toBeVisible({ timeout: 30_000 }); + await expect(userRows(page).filter({ hasText: "noteam@test.local" })).toHaveCount(0); }); }); From f93d9b6b67e12f88195167cb6f8298092bbd496d Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 31 Aug 2026 16:12:59 -0700 Subject: [PATCH 117/120] feat(complexity_router): escalate oversized prompts to a tier that fits before dispatch (#38844) * feat(complexity_router): escalate oversized prompts to a tier that fits before dispatch The classifier scores complexity and never prompt size, so a long agentic session whose newest ask is trivial classifies SIMPLE onto a small-window tier and the provider rejects it with a context-window 400 that nothing retries. The gate runs after classification on every decision path (classify tail and session-affinity pin), estimates prompt tokens including the out-of-band carriers (top-level system, tools, instructions), and when the decided tier provably cannot hold the prompt moves the request to the lowest configured tier with a model whose declared window fits, restricting the pick to fitting models when the decided tier can keep it. Models with no resolvable window are never escalated away from or onto, escalated decisions are never written as session pins, and the decision records context_escalated plus the original tier in spend logs. Resolves LIT-6503 * fix(complexity_router): judge groups by smallest window, bound skips by bytes, filter adaptive picks Review-round rework, one mechanism per finding. A group is judged by its smallest resolvable deployment window, since the core router picks within a group with no fit check. The counting skip is gated on UTF-8 byte length, which BPE token counts can never exceed, so token-dense scripts cannot slip past it; only a real tokenizer count ever moves a request and a failed count leaves the placement alone. The fit facts now filter every adaptive phase including cold start and the tier fallbacks. Window questions adopt the declared provider and never resolve authenticating providers, and a router instance without get_model_list degrades the gate to a no-op. Tests rebuilt on real Router instances resolving deployment model_info end to end, plus a full-path test through async_get_available_deployment --- .../complexity_router/complexity_router.py | 293 ++++++++++++- .../complexity_router/config.py | 26 ++ litellm/types/utils.py | 4 + .../router_strategy/test_complexity_router.py | 396 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 16 + 5 files changed, 720 insertions(+), 15 deletions(-) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 1aeedaa97b2..329da35eab3 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -733,11 +733,20 @@ def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bo of the three: an agent names the conversation on its first turn, so the cheapest tier would be the pin every session starts with, and the real work that follows would run there for the whole TTL. It describes what that one call is, never what the session's traffic looks like. + + A context-window escalation describes the prompt's size, not the session's complexity, and + size shrinks again the moment the client compacts: pinning the escalated tier would hold the + session on the big-window model long after the oversized context that forced it is gone. The + gate re-fires per request, so leaving these unpinned costs nothing but the classifier call. """ - return decision is None or decision.get("cause") not in ( - "default_model_fallback", - "plan_mode", - "housekeeping", + return decision is None or ( + decision.get("cause") + not in ( + "default_model_fallback", + "plan_mode", + "housekeeping", + ) + and not decision.get("context_escalated") ) @@ -786,6 +795,39 @@ class ClassificationOutcome(NamedTuple): classifier_cost: float | None = None +def _allowed(models: tuple[str, ...], fit_filter: frozenset[str] | None) -> tuple[str, ...]: + return models if fit_filter is None else tuple(model for model in models if model in fit_filter) + + +def _apply_context_placement( + tier: ComplexityTier | str, signals: tuple[str, ...], placement: _ContextWindowPlacement | None +) -> tuple[ComplexityTier | str, tuple[str, ...], ComplexityTier | str | None]: + """(final tier, signals, original tier when the gate escalated, else None).""" + if placement is None: + return tier, signals, None + if _tier_name(placement.tier) == _tier_name(tier): + return placement.tier, signals, None + return placement.tier, (*signals, "context_escalation"), tier + + +def _window_can_hold(window: int | None, needed: int, buffer: float) -> bool: + return window is None or needed <= int(window * buffer) + + +def _group_provably_fits(facts: tuple[int | None, bool], needed: int, buffer: float) -> bool: + window, has_unknown = facts + return window is not None and not has_unknown and needed <= int(window * buffer) + + +class _ContextWindowPlacement(NamedTuple): + """Where the context-window gate placed the request: the placement tier, the subset of its + pool the pick may use, and every configured group not provably misfit (the adaptive filter).""" + + tier: ComplexityTier | str + allowed_models: tuple[str, ...] + holdable_models: frozenset[str] + + class _SessionAffinityPin(NamedTuple): model: str tier: ComplexityTier | None @@ -1222,6 +1264,7 @@ class ComplexityRouter(CustomLogger): classifier_cost: float | None = None, conversation_continuing: bool = True, tier_litellm_params: Mapping[str, object] | None = None, + context_escalation_original_tier: ComplexityTier | str | None = None, ) -> StandardLoggingRoutingDecision: """Assemble the per-request provenance record for this router's decision. @@ -1271,6 +1314,12 @@ class ComplexityRouter(CustomLogger): decision["classifier_model"] = classifier_model if classifier_cost is not None: decision["classifier_cost"] = classifier_cost + if context_escalation_original_tier is not None: + # The pair travels together: the flag says the gate moved the request off its + # decided tier on prompt size, and the original tier names where the decision + # (classifier, keyword rule, or session pin) had placed it before physics did. + decision["context_escalated"] = True + decision["context_escalation_original_tier"] = _tier_name(context_escalation_original_tier) if tier_litellm_params: masked_tier_litellm_params: Final = mask_credentials_in_payload(tier_litellm_params) if isinstance(masked_tier_litellm_params, Mapping): @@ -1671,7 +1720,7 @@ class ComplexityRouter(CustomLogger): return entry.litellm_params if entry is not None else MappingProxyType({}) @staticmethod - def _pick_from_tier_value(model: str | list[str], tier_key: str) -> str: + def _pick_from_tier_value(model: str | Sequence[str], tier_key: str) -> str: if isinstance(model, str): return model if not model: @@ -1687,15 +1736,21 @@ class ComplexityRouter(CustomLogger): raw_messages: list[dict[str, Any]] | None, resolved_messages: list[dict[str, Any]] | None, request_kwargs: dict, + allowed_models: tuple[str, ...] | None = None, ) -> str: if not self.config.plugins: + if allowed_models is not None: + return self._pick_from_tier_value(allowed_models, _tier_name(tier)) return self.get_model_for_tier(tier) from litellm.types.router import RoutingContext tier_key: Final = _tier_name(tier) metadata_key: Final = get_metadata_variable_name_from_kwargs(request_kwargs) - pool: Final = tuple(self._tier_pools().get(tier_key, ())) + full_pool: Final = tuple(self._tier_pools().get(tier_key, ())) + pool: Final = ( + tuple(model for model in full_pool if model in allowed_models) if allowed_models is not None else full_pool + ) if not pool: # Nothing for the plugins to filter. Falling through would raise the # plugin-filtering error below and send the operator hunting for a policy @@ -1789,6 +1844,7 @@ class ComplexityRouter(CustomLogger): request_kwargs: dict[str, Any] | None = None, hard_floor: ComplexityTier | str | None = None, hard_ceiling: ComplexityTier | str | None = None, + fit_filter: frozenset[str] | None = None, ) -> str: """hard_floor excludes every candidate whose tiers all sit below it, turning this pick's soft floors (a distance penalty a high-scoring cheap model can outweigh) into a hard @@ -1801,7 +1857,10 @@ class ComplexityRouter(CustomLogger): tier because that is all it is worth, so a bandit trading cost for quality has nothing to win and must not reach above it. Without it the distance penalty is the only thing holding the tier, and a deployment that lowers tier_distance_penalty silently gets the expensive - model back while the routing decision still reads as the cheapest tier.""" + model back while the routing decision still reads as the cheapest tier. + + fit_filter excludes candidates the context-window gate proved cannot hold the prompt, + in every phase including cold start and the tier fallbacks.""" from litellm.router_strategy.adaptive_router.bandit import ( normalized_cost, thompson_sample, @@ -1812,12 +1871,12 @@ class ComplexityRouter(CustomLogger): if adaptive is None or not isinstance(classified_tier, ComplexityTier): # Custom tier names have no severity index; adaptive is rejected alongside # tier_definitions, so this guard is the contract for any future caller. - return self.get_model_for_tier(classified_tier) + return self._fitting_tier_fallback(classified_tier, fit_filter) request_type: Final = classify_prompt(user_message) classified_idx: Final = TIER_SEVERITY_ORDER.index(classified_tier) pools: Final = self._tier_pools() - classified_candidates: Final = tuple(pools.get(_tier_name(classified_tier), ())) + classified_candidates: Final = _allowed(tuple(pools.get(_tier_name(classified_tier), ())), fit_filter) cold_start_candidates: Final = tuple( model for model in classified_candidates if adaptive._cells[(request_type, model)].total_samples == 0 ) @@ -1847,9 +1906,9 @@ class ComplexityRouter(CustomLogger): if self.config.adaptive_eligible == "classified_tier": candidates = list(classified_candidates) if not candidates: - return self.get_model_for_tier(classified_tier) + return self._fitting_tier_fallback(classified_tier, fit_filter) else: - candidates = list(adaptive.config.available_models) + candidates = list(_allowed(tuple(adaptive.config.available_models), fit_filter)) all_costs: Final = [adaptive.model_to_cost.get(m, 0.0) for m in candidates] quality_weight: Final = self.config.adaptive_weights.quality @@ -1896,7 +1955,7 @@ class ComplexityRouter(CustomLogger): best_score = score best_model = model if best_model is None: - return self.get_model_for_tier(classified_tier) + return self._fitting_tier_fallback(classified_tier, fit_filter) if request_kwargs is not None: metadata = request_kwargs.setdefault("metadata", {}) if isinstance(metadata, dict): @@ -1913,6 +1972,12 @@ class ComplexityRouter(CustomLogger): } return best_model + def _fitting_tier_fallback(self, classified_tier: ComplexityTier | str, fit_filter: frozenset[str] | None) -> str: + fitting: Final = _allowed(tuple(self._tier_pools().get(_tier_name(classified_tier), ())), fit_filter) + if fit_filter is not None and fitting: + return self._pick_from_tier_value(fitting, _tier_name(classified_tier)) + return self.get_model_for_tier(classified_tier) + def _resolve_plan_mode_floor(self) -> ComplexityTier | str | None: """The configured floor as an active tier: the built-in enum member, or the defined name itself for a custom tier set; None when the feature is off.""" @@ -1983,6 +2048,163 @@ class ComplexityRouter(CustomLogger): return None return name if self.config.has_custom_tiers else ComplexityTier(name) + def _deployment_window(self, group: str, deployment: Mapping[str, object]) -> int | None: + from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider + + deployment_model_info: Final = deployment.get("model_info") + declared: Final = ( + deployment_model_info.get("max_input_tokens") if isinstance(deployment_model_info, Mapping) else None + ) + if isinstance(declared, int): + return declared + litellm_params: Final = deployment.get("litellm_params") + params: Final = litellm_params if isinstance(litellm_params, Mapping) else EMPTY_MAPPING + provider_override: Final = params.get("custom_llm_provider") + # get_router_model_info resolves the provider, and get_llm_provider runs the OAuth device + # flow for github_copilot/chatgpt, so a metadata question must never reach it for those. + if declared_authenticating_provider( + str(params.get("model") or ""), provider_override if isinstance(provider_override, str) else None + ): + return None + try: + model_info: Final = self.litellm_router_instance.get_router_model_info( + deployment=cast(dict, deployment), # cast-ok: router deployments are plain dicts + received_model_name=group, + ) + window: Final = model_info.get("max_input_tokens") + except Exception: # noqa: BLE001 # best-effort: an unmappable deployment must not hide the others + return None + return window if isinstance(window, int) else None + + def _group_window_facts(self, group: str) -> tuple[int | None, bool]: + """(smallest declared context window across the group's deployments, whether any deployment + declares none). The core router picks a deployment within the group without a fit check, so + the group is only as safe as its smallest member.""" + list_models: Final = getattr(self.litellm_router_instance, "get_model_list", None) + deployments: Final = list_models(model_name=group) if callable(list_models) else None + if not isinstance(deployments, list) or not deployments: + return (None, True) + windows: Final = tuple( + window for deployment in deployments if (window := self._deployment_window(group, deployment)) is not None + ) + return (min(windows) if windows else None, len(windows) < len(deployments)) + + @staticmethod + def _out_of_band_request_text(request_kwargs: Mapping[str, object]) -> str: + """Prompt content the resolved message list never carries: the Responses API's + `instructions`, the /v1/messages top-level `system` block, and tool definitions. + A coding agent's context is dominated by these.""" + import json + + instructions: Final = request_kwargs.get("instructions") + proxy_request: Final = request_kwargs.get("proxy_server_request") + body: Final = proxy_request.get("body") if isinstance(proxy_request, Mapping) else None + system: Final = body.get("system") if isinstance(body, Mapping) else None + tools: Final = ( + body.get("tools") if isinstance(body, Mapping) and body.get("tools") else request_kwargs.get("tools") + ) + tools_text = "" + if tools: + try: + tools_text = json.dumps(tools, default=str) + except (TypeError, ValueError): + tools_text = str(tools) + return ( + (instructions if isinstance(instructions, str) else "") + + (str(system) if system is not None else "") + + tools_text + ) + + def _request_byte_upper_bound( + self, resolved_messages: Sequence[Mapping[str, object]] | None, request_kwargs: Mapping[str, object] + ) -> int: + """UTF-8 byte length of all prompt content. BPE emits at least one byte per token in every + script, so the token count never exceeds this and 'bytes fit' soundly skips counting.""" + content_bytes: Final = sum(len(str(m.get("content") or "").encode()) for m in resolved_messages or ()) + return content_bytes + len(self._out_of_band_request_text(request_kwargs).encode()) + + async def _counted_request_tokens( + self, resolved_messages: Sequence[Mapping[str, object]], request_kwargs: Mapping[str, object] + ) -> int | None: + """Real-tokenizer count of the resolved messages plus the out-of-band carriers, off the + event loop; None when counting fails, and the gate then leaves the placement alone.""" + import litellm + from litellm.litellm_core_utils.asyncify import asyncify + + out_of_band: Final = self._out_of_band_request_text(request_kwargs) + try: + counted: Final = await asyncify(litellm.token_counter)( + messages=cast(list, resolved_messages) # cast-ok: token_counter only iterates the sequence + ) + return counted + (await asyncify(litellm.token_counter)(text=out_of_band) if out_of_band else 0) + except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request + verbose_router_logger.debug("ComplexityRouter: context-window token count failed. Got - %s", e) + return None + + async def _context_window_placement( + self, + tier: ComplexityTier | str, + resolved_messages: Sequence[Mapping[str, object]] | None, + request_kwargs: Mapping[str, object], + pool_override: tuple[str, ...] | None = None, + ) -> _ContextWindowPlacement | None: + """Correct a decided placement whose models provably cannot hold the prompt, or None + (the placement stands). Only a real tokenizer count ever moves a request, escalation + lands only on groups whose every deployment declares a fitting window, and a group + with no resolvable window is never moved on faith in either direction.""" + if not self.config.enable_context_window_escalation or not resolved_messages: + return None + pools: Final = self._tier_pools() + pool: Final = pool_override if pool_override is not None else tuple(pools.get(_tier_name(tier), ())) + if not pool: + return None + facts: Final = MappingProxyType({group: self._group_window_facts(group) for group in pool}) + known_windows: Final = tuple(window for window, _ in facts.values() if window is not None) + if not known_windows: + return None + buffer: Final = self.config.context_window_escalation_buffer + if self._request_byte_upper_bound(resolved_messages, request_kwargs) <= int(min(known_windows) * buffer): + return None + needed: Final = await self._counted_request_tokens(resolved_messages, request_kwargs) + if needed is None: + return None + return self._placement_for_tokens(tier=tier, pool=pool, pools=pools, facts=facts, needed=needed) + + def _placement_for_tokens( + self, + *, + tier: ComplexityTier | str, + pool: tuple[str, ...], + pools: Mapping[str, list[str]], + facts: Mapping[str, tuple[int | None, bool]], + needed: int, + ) -> _ContextWindowPlacement | None: + buffer: Final = self.config.context_window_escalation_buffer + in_tier: Final = tuple(group for group in pool if _window_can_hold(facts[group][0], needed, buffer)) + if in_tier and len(in_tier) == len(pool): + return None + holdable: Final = frozenset( + group + for tier_pool in pools.values() + for group in tier_pool + if _window_can_hold(self._group_window_facts(group)[0], needed, buffer) + ) + if in_tier: + return _ContextWindowPlacement(tier=tier, allowed_models=in_tier, holdable_models=holdable) + for name in self.config.tier_names()[self._active_tier_severity(tier) + 1 :]: + proven = tuple( + group + for group in pools.get(name, ()) + if _group_provably_fits(self._group_window_facts(group), needed, buffer) + ) + if proven: + return _ContextWindowPlacement( + tier=name if self.config.has_custom_tiers else ComplexityTier(name), + allowed_models=proven, + holdable_models=holdable, + ) + return None + def _apply_plan_mode_floor(self, tier: ComplexityTier | str) -> ComplexityTier | str: """The higher of the decided tier and the plan-mode floor; identity when the floor is unset.""" floor: Final = self._resolve_plan_mode_floor() @@ -2381,6 +2603,26 @@ class ComplexityRouter(CustomLogger): session_model: Final = routed_model if plan_floored and pinned_tier is not None: routed_model = self.get_model_for_tier(self._apply_plan_mode_floor(pinned_tier)) + pin_source_tier: Final = self._tier_for_model(routed_model) + pin_placement: Final = ( + await self._context_window_placement( + pin_source_tier, resolved_messages, request_kwargs, pool_override=(routed_model,) + ) + if pin_source_tier is not None + else None + ) + pin_context_original_tier: Final = ( + pin_source_tier + if pin_placement is not None + and pin_source_tier is not None + and _tier_name(pin_placement.tier) != _tier_name(pin_source_tier) + else None + ) + if pin_placement is not None and pin_context_original_tier is not None: + # The stored pin below keeps the session's own model on purpose. + routed_model = self._pick_from_tier_value( + pin_placement.allowed_models, _tier_name(pin_placement.tier) + ) # Refresh the TTL on every hit so an active session doesn't lose its # pin mid-conversation just because it outlives the original write. await self.litellm_router_instance.cache.async_set_cache( @@ -2405,7 +2647,11 @@ class ComplexityRouter(CustomLogger): verbose_router_logger.info( "ComplexityRouter: routing decision cause=%s, routed_model=%s", cause, routed_model ) - routed_pin_tier: Final = self._tier_for_model(routed_model) if plan_floored else resolved_pin_tier + routed_pin_tier: Final = ( + pin_placement.tier + if pin_placement is not None and pin_context_original_tier is not None + else (self._tier_for_model(routed_model) if plan_floored else resolved_pin_tier) + ) session_tier_litellm_params: Final = self._litellm_params_for_model(routed_pin_tier, routed_model) has_original_messages: Final = messages is not None and len(messages) > 0 return self._with_session_deployment_affinity( @@ -2422,6 +2668,7 @@ class ComplexityRouter(CustomLogger): escalated=escalated, conversation_continuing=conversation_continuing, tier_litellm_params=session_tier_litellm_params, + context_escalation_original_tier=pin_context_original_tier, ), ) ) @@ -2616,6 +2863,8 @@ class ComplexityRouter(CustomLogger): plan_floored: Final = tier != pre_floor_tier if plan_floored: signals = (*signals, "plan_mode_floor") + context_placement: Final = await self._context_window_placement(tier, resolved_messages, request_kwargs) + tier, signals, context_original_tier = _apply_context_placement(tier, signals, context_placement) score_repr: Final = f"{score:.3f}" if score is not None else "n/a" fallback_model: Final = self.config.default_model if not self.config.plugins else None # A sentinel-carrying request skips the failure exit below, whether or not the floor @@ -2662,8 +2911,15 @@ class ComplexityRouter(CustomLogger): # the cheapest tier would then contradict the floor and bound the pick below the tier # the decision reports. housekeeping_ceiling: Final = tier if outcome.cause == "housekeeping" else None + # A context-escalated tier becomes the hard floor: a floor the bandit can slide + # under is not a floor. routed_model = self._soft_floor_pick( - tier, user_message, request_kwargs, hard_floor=plan_floor, hard_ceiling=housekeeping_ceiling + tier, + user_message, + request_kwargs, + hard_floor=tier if context_original_tier is not None else plan_floor, + hard_ceiling=housekeeping_ceiling, + fit_filter=context_placement.holdable_models if context_placement is not None else None, ) adaptive: Final = self._ensure_adaptive_router() if adaptive is not None: @@ -2680,7 +2936,13 @@ class ComplexityRouter(CustomLogger): routed_model, ) else: - routed_model = await self._pick_model_for_tier(tier, messages, resolved_messages, request_kwargs) + routed_model = await self._pick_model_for_tier( + tier, + messages, + resolved_messages, + request_kwargs, + allowed_models=context_placement.allowed_models if context_placement is not None else None, + ) verbose_router_logger.info( "ComplexityRouter: routing decision cause=%s, tier=%s, score=%s, signals=%s, routed_model=%s", outcome.cause, @@ -2733,5 +2995,6 @@ class ComplexityRouter(CustomLogger): classifier_model=classifier_model, classifier_cost=outcome.classifier_cost, tier_litellm_params=tier_litellm_params, + context_escalation_original_tier=context_original_tier, ), ) diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 0abe962edf1..3c5e8aafa18 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -823,6 +823,32 @@ class ComplexityRouterConfig(BaseModel): ), ) + enable_context_window_escalation: bool = Field( + default=True, + description=( + "Escalate a request off a tier whose models provably cannot hold its prompt, before " + "dispatch. The classifier scores complexity and never prompt size, so a long agentic " + "session whose newest ask is trivial lands on a small-window tier and the provider " + "rejects it with a context-window 400 that nothing retries. When every model of the " + "decided tier has a declared window smaller than the estimated prompt, the request " + "moves to the lowest configured tier with a model whose declared window fits; when " + "only some of the tier's models fit, the pick is restricted to those and the tier " + "keeps the request. Models with no resolvable window are never escalated away from " + "and never escalated onto. Set false to dispatch on complexity alone, as before." + ), + ) + context_window_escalation_buffer: float = Field( + default=0.95, + gt=0, + le=1, + description=( + "Fraction of a model's declared context window the estimated prompt must fit within. " + "The token count is an estimate, so fitting against the full window would dispatch " + "prompts that the provider's own tokenizer then rejects; 0.95 leaves room for that " + "drift plus the response tokens." + ), + ) + # Semantic (embedding) matching for keyword_tier_rules instead of literal text matching semantic_keyword_matching: bool = Field( default=False, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 14749ecde6a..a1b3523442b 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2886,6 +2886,8 @@ class StandardLoggingRoutingDecision(TypedDict, total=False): classifier_model: str classifier_cost: float escalated: bool + context_escalated: bool # writable-ok: Pydantic warns on ReadOnly TypedDict fields + context_escalation_original_tier: str # writable-ok: Pydantic warns on ReadOnly TypedDict fields tier_boundaries: StandardLoggingRoutingDecisionTierBoundaries reasoning_override_min_score: float # writable-ok: Pydantic warns on ReadOnly TypedDict fields conversation_continuing: bool @@ -2912,6 +2914,8 @@ DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset( "classifier_model", "classifier_cost", "escalated", + "context_escalated", + "context_escalation_original_tier", "tier_boundaries", "reasoning_override_min_score", "conversation_continuing", diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 97f60ec8e57..3f7844cffba 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -10021,3 +10021,399 @@ class TestHeuristicFirst: ) outcome = await router.aclassify(NO_SIGNAL_PROMPT) assert outcome.cause == "default_model_fallback" + + +def _windowed_router(*deployments: tuple) -> Router: + """Real Router; each deployment is (group, provider_model, declared window or None). + None means no declared override on a model the cost map does not know: unresolvable.""" + return Router( + model_list=[ + { + "model_name": group, + "litellm_params": {"model": provider_model, "mock_response": "ok"}, + **({"model_info": {"max_input_tokens": window}} if window is not None else {}), + } + for group, provider_model, window in deployments + ] + ) + + +_SMALL = ("small-model", "openai/gpt-3.5-turbo", 16385) +_BIG = ("big-model", "openai/gpt-4o-mini", 200000) + +# A long agentic session whose newest ask is trivial: low-density filler the heuristic scores +# SIMPLE, sized well past a 16,385-token window so the fit check must move it. +_CONTEXT_FILLER = "The meeting notes were saved to the shared folder for later review this week. " * 2000 +_OVERSIZED_TURNS = [ + {"role": "user", "content": "Here is everything discussed so far. " + _CONTEXT_FILLER}, + {"role": "assistant", "content": "Noted, I have read all of it."}, + {"role": "user", "content": "ok continue"}, +] +# ~40k CJK chars: chars/4 says ~10k tokens, the real tokenizer says several times that. A +# character-based shortcut would skip counting and dispatch this to a 16k window. +_CJK_TURNS = [ + {"role": "user", "content": "会议记录已经保存到共享文件夹里，供大家本周晚些时候查阅和讨论使用。" * 1300}, + {"role": "user", "content": "ok continue"}, +] + + +def _tier_config(**overrides) -> Dict: + return {"tiers": {"SIMPLE": "small-model", "COMPLEX": "big-model"}, **overrides} + + +class TestContextWindowEscalation: + """A tier decided on complexity alone must still hold the prompt, or the provider 400s. + + The classifier never weighs prompt size (token count is a 0.10-weight scoring dimension, + below every tier boundary), so a long session ending in a trivial ask lands on the + smallest tier and dies upstream with no retry. The gate checks fit pre-dispatch, against + windows resolved through the real Router deployment chain. + """ + + @pytest.mark.asyncio + async def test_an_oversized_simple_prompt_escalates_to_the_lowest_tier_that_fits(self): + """The LIT-6503 regression: SIMPLE verdict, 17k-token prompt, 16,385-token tier model. + + Unfixed, this dispatched to the small model and the provider rejected it with a + context-window 400 that neither the retry layer nor tier-keyed fallbacks catch. + """ + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=_windowed_router(_SMALL, _BIG), + complexity_router_config=_tier_config(), + ) + + result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) + + assert result is not None + assert result.model == "big-model" + assert result.routing_decision["context_escalated"] is True + assert result.routing_decision["context_escalation_original_tier"] == "SIMPLE" + assert result.routing_decision["tier"] == "COMPLEX" + assert "context_escalation" in result.routing_decision["signals"] + + @pytest.mark.asyncio + async def test_a_prompt_that_fits_routes_exactly_as_before(self): + """The gate must be invisible for normal traffic: same model, no escalation facts.""" + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=_windowed_router(_SMALL, _BIG), + complexity_router_config=_tier_config(), + ) + + result = await router.async_pre_routing_hook( + model="test-router", request_kwargs={}, messages=[{"role": "user", "content": "ok continue"}] + ) + + assert result is not None + assert result.model == "small-model" + assert "context_escalated" not in result.routing_decision + assert "context_escalation_original_tier" not in result.routing_decision + + @pytest.mark.asyncio + async def test_the_pick_prefers_a_fitting_group_inside_the_decided_tier(self): + """A tier holding both a small and a large group keeps the request and picks the one + that fits, which is cheaper than escalating and preserves the classifier's decision.""" + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=_windowed_router(_SMALL, ("mid-model", "openai/gpt-4o-mini", 200000), _BIG), + complexity_router_config={"tiers": {"SIMPLE": ["small-model", "mid-model"], "COMPLEX": "big-model"}}, + ) + + result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) + + assert result is not None + assert result.model == "mid-model" + assert result.routing_decision["tier"] == "SIMPLE" + assert "context_escalated" not in result.routing_decision + + @pytest.mark.asyncio + async def test_a_group_is_only_as_safe_as_its_smallest_deployment(self): + """One group name can front deployments with different windows, and the core router + picks among them with no fit check, so retaining the group on its largest member + turns the pick into a coin flip against a 400. The gate judges the group by its + smallest resolvable window and escalates past it.""" + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=Router( + model_list=[ + { + "model_name": "mixed-pool", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "mock_response": "ok"}, + "model_info": {"max_input_tokens": 16385}, + }, + { + "model_name": "mixed-pool", + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"max_input_tokens": 200000}, + }, + { + "model_name": "big-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"max_input_tokens": 200000}, + }, + ] + ), + complexity_router_config={"tiers": {"SIMPLE": "mixed-pool", "COMPLEX": "big-model"}}, + ) + + result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) + + assert result is not None + assert result.model == "big-model" + assert result.routing_decision["context_escalated"] is True + + @pytest.mark.asyncio + async def test_token_dense_text_cannot_slip_past_the_counting_shortcut(self): + """CJK text runs several tokens per four characters, so a chars/4 shortcut would skip + the real count and dispatch an oversized prompt. The skip is gated on the UTF-8 byte + length, which the token count can never exceed.""" + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=_windowed_router(_SMALL, _BIG), + complexity_router_config=_tier_config(), + ) + + result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_CJK_TURNS) + + assert result is not None + assert result.model == "big-model" + assert result.routing_decision["context_escalated"] is True + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "deployments,tiers,expected_model", + [ + ( + (("small-model", "openai/unmapped-model-under-test", None), _BIG), + {"SIMPLE": "small-model", "COMPLEX": "big-model"}, + "small-model", + ), + ( + (_SMALL, ("mid-model", "openai/another-unmapped-model", None), _BIG), + {"SIMPLE": "small-model", "MEDIUM": "mid-model", "COMPLEX": "big-model"}, + "big-model", + ), + ((_SMALL,), {"SIMPLE": "small-model"}, "small-model"), + ], + ids=["unknown-window-stays", "unproven-target-skipped", "nothing-fits-stays"], + ) + async def test_unknown_windows_are_never_acted_on(self, deployments, tiers, expected_model): + """No faith in either direction: a model with no resolvable window is never escalated + away from (its misfit is unprovable) and never escalated onto (its fit is unprovable); + when nothing provably fits, the classified tier stands and the client owns overflow.""" + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=_windowed_router(*deployments), + complexity_router_config={"tiers": tiers}, + ) + + result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) + + assert result is not None + assert result.model == expected_model + + @pytest.mark.asyncio + async def test_the_disabled_gate_dispatches_on_complexity_alone(self): + """The escape hatch: enable_context_window_escalation false restores today's behavior.""" + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=_windowed_router(_SMALL, _BIG), + complexity_router_config=_tier_config(enable_context_window_escalation=False), + ) + + result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) + + assert result is not None + assert result.model == "small-model" + assert "context_escalated" not in result.routing_decision + + @pytest.mark.asyncio + async def test_out_of_band_system_and_tools_count_against_the_window(self): + """The Claude Code shape that live-testing caught: a tiny ask riding a top-level + `system` block and tool definitions that together dwarf the message list. None of + that reaches resolved messages on /v1/messages, so a gate reading only messages + dispatches a provably oversized request and the provider 400s anyway.""" + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=_windowed_router(_SMALL, _BIG), + complexity_router_config=_tier_config(), + ) + + result = await router.async_pre_routing_hook( + model="test-router", + request_kwargs={ + "proxy_server_request": { + "body": { + "system": _CONTEXT_FILLER, + "tools": [{"name": f"tool_{i}", "description": _CONTEXT_FILLER[:500]} for i in range(20)], + } + } + }, + messages=[{"role": "user", "content": "reply with exactly: rig check ok"}], + ) + + assert result is not None + assert result.model == "big-model" + assert result.routing_decision["context_escalated"] is True + + @pytest.mark.asyncio + async def test_an_escalated_first_turn_never_becomes_the_session_pin(self): + """Escalation describes the prompt's size, not the session: once the client compacts, + the next turn fits again, so pinning the big-window tier would hold the whole session + on it for the TTL. The escalated turn routes big, and the next fitting turn classifies + fresh instead of inheriting a pin.""" + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=_windowed_router(_SMALL, _BIG), + complexity_router_config=_tier_config(session_affinity=True), + ) + session_kwargs = lambda: {"metadata": {"session_id": "s-1", "user_api_key_hash": "k-1"}} # noqa: E731 + + first = await router.async_pre_routing_hook( + model="test-router", request_kwargs=session_kwargs(), messages=_OVERSIZED_TURNS + ) + second = await router.async_pre_routing_hook( + model="test-router", request_kwargs=session_kwargs(), messages=[{"role": "user", "content": "ok continue"}] + ) + + assert first is not None and first.model == "big-model" + assert second is not None and second.model == "small-model" + assert second.routing_decision["cause"] != "session_affinity_pin" + + @pytest.mark.asyncio + async def test_a_pinned_session_escalates_per_request_and_keeps_its_pin(self): + """The pin fast path skips classification, not physics: an oversized turn on a session + pinned to the small tier is served by the fitting tier, while the stored pin keeps the + session's own model so the first turn that fits again routes exactly as pinned.""" + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=_windowed_router(_SMALL, _BIG), + complexity_router_config=_tier_config(session_affinity=True), + ) + session_kwargs = lambda: {"metadata": {"session_id": "s-2", "user_api_key_hash": "k-2"}} # noqa: E731 + + pinned = await router.async_pre_routing_hook( + model="test-router", request_kwargs=session_kwargs(), messages=[{"role": "user", "content": "ok continue"}] + ) + oversized = await router.async_pre_routing_hook( + model="test-router", request_kwargs=session_kwargs(), messages=_OVERSIZED_TURNS + ) + back_to_small = await router.async_pre_routing_hook( + model="test-router", request_kwargs=session_kwargs(), messages=[{"role": "user", "content": "ok continue"}] + ) + + assert pinned is not None and pinned.model == "small-model" + assert oversized is not None and oversized.model == "big-model" + assert oversized.routing_decision["cause"] == "session_affinity_pin" + assert oversized.routing_decision["context_escalated"] is True + assert oversized.routing_decision["context_escalation_original_tier"] == "SIMPLE" + assert back_to_small is not None and back_to_small.model == "small-model" + assert back_to_small.routing_decision["cause"] == "session_affinity_pin" + + @pytest.mark.asyncio + async def test_the_adaptive_cold_start_never_samples_a_model_that_cannot_hold_the_prompt(self): + """The bandit's exploration is still bounded by physics: with the whole classified tier + unobserved, cold start samples only among models whose window holds the prompt.""" + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=Router( + model_list=[ + { + "model_name": "small-model", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "mock_response": "ok"}, + "model_info": {"max_input_tokens": 16385}, + }, + { + "model_name": "mid-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"max_input_tokens": 200000}, + }, + ] + ), + complexity_router_config={"adaptive": True, "tiers": {"SIMPLE": ["small-model", "mid-model"]}}, + ) + + result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) + + assert result is not None + assert result.model == "mid-model" + + @pytest.mark.asyncio + async def test_the_gate_never_resolves_an_authenticating_provider(self, monkeypatch, tmp_path): + """Resolving github_copilot runs its OAuth device flow, so a window question must adopt + the declaration instead of resolving: the copilot group reads as unknown-window and the + request stays put, with zero copilot resolutions recorded.""" + import json + import time + + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) + (tmp_path / "api-key.json").write_text(json.dumps({"token": "tid=test", "expires_at": int(time.time()) + 3600})) + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=Router( + model_list=[ + {"model_name": "cop-pool", "litellm_params": {"model": "github_copilot/gpt-4o"}}, + { + "model_name": "big-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"max_input_tokens": 200000}, + }, + ] + ), + complexity_router_config={"tiers": {"SIMPLE": "cop-pool", "COMPLEX": "big-model"}}, + ) + real_get_llm_provider = litellm.get_llm_provider + copilot_resolutions: List = [] + + def _guarded(*args, **kwargs): + target = str(kwargs.get("model") or (args[0] if args else "")) + str(kwargs.get("custom_llm_provider") or "") + if "github_copilot" in target: + copilot_resolutions.append(target) + raise RuntimeError("the gate must not resolve an authenticating provider") + return real_get_llm_provider(*args, **kwargs) + + monkeypatch.setattr(litellm, "get_llm_provider", _guarded) + + result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) + + assert result is not None + assert result.model == "cop-pool" + assert copilot_resolutions == [] + + @pytest.mark.asyncio + async def test_the_full_routing_path_serves_the_escalated_deployment(self): + """End to end through Router.async_get_available_deployment: the auto-router alias with + an oversized prompt resolves to the big tier's deployment, and a small prompt to the + small tier's, with no mocking anywhere in the resolution chain.""" + router = Router( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"SIMPLE": "small-model", "COMPLEX": "big-model"}}, + }, + }, + { + "model_name": "small-model", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "mock_response": "ok"}, + "model_info": {"max_input_tokens": 16385}, + }, + { + "model_name": "big-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"max_input_tokens": 200000}, + }, + ] + ) + + oversized = await router.async_get_available_deployment( + model="smart-router", request_kwargs={}, messages=_OVERSIZED_TURNS + ) + small = await router.async_get_available_deployment( + model="smart-router", request_kwargs={}, messages=[{"role": "user", "content": "ok continue"}] + ) + + assert oversized["model_name"] == "big-model" + assert small["model_name"] == "small-model" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 140130adc4b..b79e2c12447 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -34385,6 +34385,12 @@ export interface components { * @description Keywords indicating code-related content */ code_keywords?: string[] | null; + /** + * Context Window Escalation Buffer + * @description Fraction of a model's declared context window the estimated prompt must fit within. The token count is an estimate, so fitting against the full window would dispatch prompts that the provider's own tokenizer then rejects; 0.95 leaves room for that drift plus the response tokens. + * @default 0.95 + */ + context_window_escalation_buffer: number; /** * Custom Technical Keywords * @description Domain-specific technical keywords appended to the effective base list (technical_keywords if set, otherwise DEFAULT_TECHNICAL_KEYWORDS). Order is preserved; duplicates are removed case-insensitively against the base list and within this list. @@ -34413,6 +34419,12 @@ export interface components { * @description Embedding model (LiteLLM model name) used when semantic_keyword_matching is enabled */ embedding_model?: string | null; + /** + * Enable Context Window Escalation + * @description Escalate a request off a tier whose models provably cannot hold its prompt, before dispatch. The classifier scores complexity and never prompt size, so a long agentic session whose newest ask is trivial lands on a small-window tier and the provider rejects it with a context-window 400 that nothing retries. When every model of the decided tier has a declared window smaller than the estimated prompt, the request moves to the lowest configured tier with a model whose declared window fits; when only some of the tier's models fit, the pick is restricted to those and the tier keeps the request. Models with no resolvable window are never escalated away from and never escalated onto. Set false to dispatch on complexity alone, as before. + * @default true + */ + enable_context_window_escalation: boolean; /** * Escalation Keywords * @description Case-sensitive phrases a user can include to force a bump to the next-higher complexity tier when they aren't satisfied with results (they can force a stronger model, but not choose which one). Defaults to ['LITELLM ESCALATE'] when unset; set to an empty list to disable. @@ -35602,6 +35614,10 @@ export interface components { classifier_cost?: number; /** Classifier Model */ classifier_model?: string; + /** Context Escalated */ + context_escalated?: boolean; + /** Context Escalation Original Tier */ + context_escalation_original_tier?: string; /** Conversation Continuing */ conversation_continuing?: boolean; /** Escalated */ From 38294188789cecfd6dc44f042f337393b4f0c7e8 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 31 Aug 2026 16:37:38 -0700 Subject: [PATCH 118/120] feat(shadow_eval): target teams and users so JWT-auth traffic can be evaluated (#39015) Shadow eval jobs previously targeted only virtual keys, so deployments on pure JWT auth (which present no key at all) could never sample their traffic. Jobs now carry a typed (target_type, target_id) pair covering keys, teams, and users; sampling matches the identity every request resolves to at auth time, so team and user jobs cover JWT traffic with no client changes. Resolves LIT-6578 --- .../migration.sql | 21 + .../litellm_proxy_extras/schema.prisma | 7 +- litellm/integrations/shadow_eval_logger.py | 42 +- .../auto_router_endpoints.py | 381 ++++++++++++----- litellm/proxy/schema.prisma | 7 +- .../auto_router_endpoints.py | 120 ++++-- schema.prisma | 7 +- .../integrations/test_shadow_eval_logger.py | 115 ++++- .../test_auto_router_endpoints.py | 395 +++++++++++++++--- .../_components/ShadowEvalSection.test.tsx | 144 +++++-- .../_components/ShadowEvalSection.tsx | 128 ++++-- .../_components/useShadowEval.ts | 2 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 203 +++++---- 13 files changed, 1181 insertions(+), 391 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260831000000_shadow_eval_typed_targets/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831000000_shadow_eval_typed_targets/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831000000_shadow_eval_typed_targets/migration.sql new file mode 100644 index 00000000000..b7dbe931dd2 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831000000_shadow_eval_typed_targets/migration.sql @@ -0,0 +1,21 @@ +DO $$ +BEGIN + IF EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_name = 'LiteLLM_ShadowEvalJob' AND column_name = 'api_key_id' + ) THEN + ALTER TABLE "LiteLLM_ShadowEvalJob" RENAME COLUMN "api_key_id" TO "target_id"; + END IF; +END $$; + +ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN IF NOT EXISTS "target_type" TEXT NOT NULL DEFAULT 'key'; + +DROP INDEX IF EXISTS "LiteLLM_ShadowEvalJob_one_active_per_key_direction"; + +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJob_one_active_per_target_direction" + ON "LiteLLM_ShadowEvalJob"("target_type", "target_id", "direction") WHERE "stopped_at" IS NULL; + +DROP INDEX IF EXISTS "LiteLLM_ShadowEvalJob_api_key_id_idx"; + +CREATE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJob_target_type_target_id_idx" + ON "LiteLLM_ShadowEvalJob"("target_type", "target_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 60223265211..01a607b68a9 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1529,14 +1529,15 @@ model LiteLLM_AutoRouterSession { model LiteLLM_ShadowEvalJob { id String @id @default(cuid()) group_id String // legs of one job share this; the API's job id - api_key_id String // hashed virtual key whose traffic this leg shadows + target_type String @default("key") // key | team | user + target_id String // hashed virtual key, team_id, or user_id whose traffic this leg shadows router_name String // the auto-router under evaluation, in either direction direction String @default("forward") // forward | reverse baseline_model String? // reverse only: the fixed model the router is judged against judge_model String shadow_percentage Float max_turns Int // sample-count ceiling: the whole budget on pre-max_budget jobs, the error-loop valve otherwise - max_budget Float? // per-key USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets + max_budget Float? // per-target USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets created_at DateTime @default(now()) created_by String? ends_at DateTime @@ -1544,7 +1545,7 @@ model LiteLLM_ShadowEvalJob { stopped_by String? // operator who stopped it early; null when it ended on its own @@index([group_id]) - @@index([api_key_id]) + @@index([target_type, target_id]) @@index([created_at]) } diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index cf8aa38d86e..18bda0a9d55 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -592,7 +592,12 @@ _JOBS_CACHE_KEY: Final = "shadow_eval:active_jobs" class ShadowEvalLogger(CustomLogger): - """Fires blind pairwise shadow evaluations for keys with an active shadow-eval job.""" + """Fires blind pairwise shadow evaluations for targets with an active shadow-eval job. + + A job targets a virtual key, a team, or a user; a request qualifies for a job when + any of its resolved identities (key hash, team id, user id) matches the job's + target, so team and user jobs cover JWT-authenticated traffic, which carries no + key hash at all.""" def __init__( self, @@ -617,10 +622,10 @@ class ShadowEvalLogger(CustomLogger): # generation; the refill absorbs written rows and resets. self._job_starts: dict[str, int] = {} # mutable-ok: per-generation counter - async def _active_jobs(self) -> Mapping[str, tuple[ActiveShadowEvalJob, ...]]: - """Active jobs by api_key_id, cache-first. A key holds at most one job per - direction, so the value is a collection. A DB fault returns empty without - caching, so sampling pauses for that request and the next one retries.""" + async def _active_jobs(self) -> Mapping[tuple[str, str], tuple[ActiveShadowEvalJob, ...]]: + """Active jobs by (target_type, target_id), cache-first. A target holds at most + one job per direction, so the value is a collection. A DB fault returns empty + without caching, so sampling pauses for that request and the next one retries.""" cached: Final = await self._jobs_cache.async_get_cache(_JOBS_CACHE_KEY) if cached is not None: return cached # pyright: ignore[reportReturnType] # cache stores exactly this mapping shape @@ -652,10 +657,10 @@ class ShadowEvalLogger(CustomLogger): ) for row in grouped or [] } - by_key: Final = tuple( + by_target: Final = tuple( sorted( ( - (str(record.api_key_id), job) + ((str(record.target_type), str(record.target_id)), job) for record in records or [] if (job := _as_active_job(record, *attempt_stats.get(str(record.id), (0, 0.0)))) is not None ), @@ -663,7 +668,7 @@ class ShadowEvalLogger(CustomLogger): ) ) jobs: Final = MappingProxyType( - {key: tuple(job for _, job in group) for key, group in groupby(by_key, key=itemgetter(0))} + {target: tuple(job for _, job in group) for target, group in groupby(by_target, key=itemgetter(0))} ) await self._jobs_cache.async_set_cache(_JOBS_CACHE_KEY, jobs) self._job_starts = {} # rebind-ok: new generation, counts absorbed into the fill @@ -720,8 +725,18 @@ class ShadowEvalLogger(CustomLogger): if should_redact_message_logging(dict(kwargs)): # mutable-ok: predicate takes a plain dict return metadata: Final = payload.get("metadata") or _EMPTY_METADATA - api_key_hash: Final = metadata.get("user_api_key_hash") - if not api_key_hash: + # Each identity the request resolved to is a candidate target; JWT-auth + # requests carry no key hash but do carry a team and user. + targets: Final = tuple( + (target_type, str(value)) + for target_type, value in ( + ("key", metadata.get("user_api_key_hash")), + ("team", metadata.get("user_api_key_team_id")), + ("user", metadata.get("user_api_key_user_id")), + ) + if value + ) + if not targets: return request_id: Final = payload.get("id") or "" if not request_id: @@ -731,8 +746,11 @@ class ShadowEvalLogger(CustomLogger): return # only surfaces this table can normalize are comparable; unknown types fail closed if ops.wire_params and _request_mutating_guardrail_ran(request_metadata): return # the wire-body snapshot predates the rewrite; replaying it would resurrect stripped content + active_jobs: Final = await self._active_jobs() eligible: Final = self._sampled_jobs( - (await self._active_jobs()).get(str(api_key_hash), ()), request_metadata, request_id + tuple(job for target in targets for job in active_jobs.get(target, ())), + request_metadata, + request_id, ) if not eligible: return @@ -1056,7 +1074,7 @@ class ShadowEvalLogger(CustomLogger): ) -_EMPTY_JOBS: Final[Mapping[str, tuple[ActiveShadowEvalJob, ...]]] = MappingProxyType({}) +_EMPTY_JOBS: Final[Mapping[tuple[str, str], tuple[ActiveShadowEvalJob, ...]]] = MappingProxyType({}) def _default_prisma_provider() -> "PrismaClient | None": diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index b5533d548e5..44b0cdcca2e 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -7,7 +7,7 @@ POST /auto_router/validate_complexity_router_config - Dry-run the complexity-rou from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone -from itertools import groupby +from itertools import chain, groupby from operator import attrgetter from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Protocol @@ -58,10 +58,11 @@ from litellm.types.management_endpoints.auto_router_endpoints import ( ComplexityRouterConfigValidationResponse, RequestComplexityRouterConfig, ShadowEvalDirection, - ShadowEvalJobKeyResponse, ShadowEvalJobResponse, + ShadowEvalJobTargetResponse, ShadowEvalResult, ShadowEvalSlice, + ShadowEvalTargetType, StartShadowEvalRequest, ) @@ -104,10 +105,43 @@ class _VerificationTokenTable(Protocol): async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_VerificationTokenRow]: ... +class _TeamRow(Protocol): + @property + def team_id(self) -> str: ... + + @property + def team_alias(self) -> str | None: ... + + +class _TeamRowsTable(Protocol): + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_TeamRow]: ... + + +class _UserRow(Protocol): + @property + def user_id(self) -> str: ... + + @property + def user_email(self) -> str | None: ... + + +class _UserRowsTable(Protocol): + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_UserRow]: ... + + class _ShadowEvalJobRow(Protocol): @property def id(self) -> str: ... + @property + def group_id(self) -> str: ... + + @property + def target_type(self) -> str: ... + + @property + def target_id(self) -> str: ... + class _ShadowEvalJobTable(Protocol): async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_ShadowEvalJobRow]: ... @@ -138,6 +172,14 @@ def _verification_tokens(prisma_client: "PrismaClient") -> _VerificationTokenTab return prisma_client.db.litellm_verificationtoken +def _team_rows(prisma_client: "PrismaClient") -> _TeamRowsTable: + return prisma_client.db.litellm_teamtable + + +def _user_rows(prisma_client: "PrismaClient") -> _UserRowsTable: + return prisma_client.db.litellm_usertable + + def _shadow_eval_jobs(prisma_client: "PrismaClient") -> _ShadowEvalJobTable: return prisma_client.db.litellm_shadowevaljob @@ -836,7 +878,7 @@ def _validate_judge_is_not_a_candidate( def _is_unique_violation(error: Exception) -> bool: - """Whether a Prisma create failed on a unique index. One active job per key and + """Whether a Prisma create failed on a unique index. One active job per target and direction lives in a partial unique index (raw SQL in the migration; schema.prisma cannot express partial indexes), so the read-then-create check above it is advisory: two concurrent starts pass the read, and the loser must surface as the same 409 @@ -885,7 +927,7 @@ _ATTEMPT_AGG_BY_LEG_SQL: Final = "SELECT job_id AS grp," + _ATTEMPT_AGG_SELECT # direction, and mid-deploy rows from old pods price as judge-only until the deploy ends). _SWEEP_FINISHED_JOBS_SQL: Final = """ UPDATE "LiteLLM_ShadowEvalJob" j SET stopped_at = (NOW() AT TIME ZONE 'utc') -WHERE j.api_key_id = ANY($1::text[]) AND j.stopped_at IS NULL +WHERE j.target_type = $2 AND j.target_id = ANY($1::text[]) AND j.stopped_at IS NULL AND ( j.ends_at <= (NOW() AT TIME ZONE 'utc') OR (SELECT COUNT(*) FROM "LiteLLM_ShadowEvalAttempt" a WHERE a.job_id = j.id) >= j.max_turns @@ -966,10 +1008,10 @@ WHERE group_id IN ( ) """ -_LIST_LEGS_BY_KEY_SQL: Final = """ +_LIST_LEGS_BY_TARGET_SQL: Final = """ SELECT * FROM "LiteLLM_ShadowEvalJob" WHERE group_id IN ( - SELECT group_id FROM "LiteLLM_ShadowEvalJob" WHERE api_key_id = $2 + SELECT group_id FROM "LiteLLM_ShadowEvalJob" WHERE target_type = $2 AND target_id = $3 GROUP BY group_id ORDER BY MAX(created_at) DESC LIMIT $1::int ) """ @@ -1007,15 +1049,16 @@ def _slices(rows: Sequence[_AttemptAggRow]) -> tuple[ShadowEvalSlice, ...]: class _LegRow(BaseModel): """One LiteLLM_ShadowEvalJob row, validated off the untyped prisma record. A row is - one key's leg of a job; the legs of a job share group_id and identical config, written - together by one create_many. The API's job id is the group id, so leg ids never leave - the server (attempts reference them internally).""" + one target's leg of a job; the legs of a job share group_id and identical config, + written together by one create_many. The API's job id is the group id, so leg ids + never leave the server (attempts reference them internally).""" model_config = ConfigDict(from_attributes=True) id: str group_id: str - api_key_id: str + target_type: ShadowEvalTargetType + target_id: str router_name: str direction: ShadowEvalDirection baseline_model: str | None = None @@ -1068,16 +1111,17 @@ def _group_response( first: Final = legs[0] return ShadowEvalJobResponse( job_id=group_id, - keys=tuple( - ShadowEvalJobKeyResponse( - api_key_id=leg.api_key_id, + targets=tuple( + ShadowEvalJobTargetResponse( + target_type=leg.target_type, + target_id=leg.target_id, max_turns=leg.max_turns, max_budget=leg.max_budget, stopped_at=leg.stopped_at, attempt_count=stats.attempt_count if (stats := attempt_counts.get(leg.id)) else 0, spend=round(stats.spend, 6) if stats else 0.0, ) - for leg in sorted(legs, key=lambda leg: leg.api_key_id) + for leg in sorted(legs, key=lambda leg: (leg.target_type, leg.target_id)) ), router_name=first.router_name, direction=first.direction, @@ -1090,34 +1134,85 @@ def _group_response( ) -_NO_KEY_LABELS: Final[tuple[str | None, str | None]] = (None, None) +_NO_TARGET_LABELS: Final[tuple[str | None, str | None]] = (None, None) -async def _with_key_labels( +def _target_labels( + key_rows: Sequence[_VerificationTokenRow], + team_rows: Sequence[_TeamRow], + user_rows: Sequence[_UserRow], +) -> Mapping[tuple[str, str], tuple[str | None, str | None]]: + """Display labels by (target_type, target_id): a key's (alias, masked name), a + team's (alias, None), a user's (email, None).""" + return MappingProxyType( + { # mutable-ok: MappingProxyType needs a dict to wrap + key: value + for key, value in chain( + ((("key", row.token), (row.key_alias, row.key_name)) for row in key_rows), + ((("team", row.team_id), (row.team_alias, None)) for row in team_rows), + ((("user", row.user_id), (row.user_email, None)) for row in user_rows), + ) + } + ) + + +def _target_ids_of(responses: Sequence[ShadowEvalJobResponse], target_type: ShadowEvalTargetType) -> tuple[str, ...]: + return tuple( + sorted( + frozenset( + target.target_id + for response in responses + for target in response.targets + if target.target_type == target_type + ) + ) + ) + + +async def _with_target_labels( prisma_client: "PrismaClient", responses: Sequence[ShadowEvalJobResponse] ) -> tuple[ShadowEvalJobResponse, ...]: - """Resolve every scoped key's hash to its alias and masked name in one batched read, - so the UI can say whose traffic a job shadows. Deleted keys resolve to None.""" + """Resolve every scoped target's id to a display label in one batched read per kind, + so the UI can say whose traffic a job shadows: a key's alias and masked name, a + team's alias, a user's email. Deleted targets resolve to None.""" if not responses: return () - tokens: Final = sorted(frozenset(key.api_key_id for response in responses for key in response.keys)) - key_rows: Final = await _verification_tokens(prisma_client).find_many( - where={"token": {"in": tokens}} # mutable-ok: Prisma filter + tokens: Final = _target_ids_of(responses, "key") + team_ids: Final = _target_ids_of(responses, "team") + user_ids: Final = _target_ids_of(responses, "user") + key_rows: Final = ( + await _verification_tokens(prisma_client).find_many( + where={"token": {"in": list(tokens)}} # mutable-ok: Prisma filter + ) + if tokens + else () ) - labels: Final[Mapping[str, tuple[str | None, str | None]]] = { - row.token: (row.key_alias, row.key_name) for row in key_rows or () - } + team_rows: Final = ( + await _team_rows(prisma_client).find_many( + where={"team_id": {"in": list(team_ids)}} # mutable-ok: Prisma filter + ) + if team_ids + else () + ) + user_rows: Final = ( + await _user_rows(prisma_client).find_many( + where={"user_id": {"in": list(user_ids)}} # mutable-ok: Prisma filter + ) + if user_ids + else () + ) + labels: Final = _target_labels(key_rows or (), team_rows or (), user_rows or ()) return tuple( response.model_copy( update={ # mutable-ok: pydantic update payload - "keys": tuple( - key.model_copy( + "targets": tuple( + target.model_copy( update={ # mutable-ok: pydantic update payload - "key_alias": labels.get(key.api_key_id, _NO_KEY_LABELS)[0], - "key_name": labels.get(key.api_key_id, _NO_KEY_LABELS)[1], + "target_alias": labels.get((target.target_type, target.target_id), _NO_TARGET_LABELS)[0], + "key_name": labels.get((target.target_type, target.target_id), _NO_TARGET_LABELS)[1], } ) - for key in response.keys + for target in response.targets ) } ) @@ -1125,29 +1220,37 @@ async def _with_key_labels( ) -async def _shadow_eval_results(prisma_client: "PrismaClient", legs: Sequence[_LegRow]) -> ShadowEvalResult | None: - """All three stratifications of one job's verdicts. Tier answers "where does the router - do well"; the model stratification groups by whichever model served the real arm, so it - answers "which of the models these keys use today would the router beat" forward, and - "for the turns the router sent to X, did X beat the baseline" in reverse; key answers - "which key's traffic does the router suit". Reads are bounded by the job's own attempts - (<= the sum of its keys' max_turns) via the job_id index.""" +async def _shadow_eval_results( + prisma_client: "PrismaClient", legs: Sequence[_LegRow] +) -> tuple[ShadowEvalResult | None, Mapping[tuple[str, str], ShadowEvalSlice]]: + """One job's stratified verdicts, plus each target's own slice keyed by the + (target_type, target_id) pair so a key, team, and user sharing an id can never + collapse into one entry. Tier answers "where does the router do well"; the model + stratification groups by whichever model served the real arm, so it answers "which + of the models these targets use today would the router beat" forward, and "for the + turns the router sent to X, did X beat the baseline" in reverse; the per-target + slices answer "which target's traffic does the router suit". Reads are bounded by + the job's own attempts (<= the sum of its targets' max_turns) via the job_id index.""" leg_ids: Final = [leg.id for leg in legs] # mutable-ok: query param by_tier: Final = _ATTEMPT_AGG_ROWS.validate_python( await _query_raw(prisma_client, _ATTEMPT_AGG_BY_TIER_SQL, leg_ids) or () ) if not by_tier: - return None + return None, MappingProxyType({}) by_model: Final = _ATTEMPT_AGG_ROWS.validate_python( await _query_raw(prisma_client, _ATTEMPT_AGG_BY_MODEL_SQL, leg_ids) or () ) - key_by_leg: Final = MappingProxyType({leg.id: leg.api_key_id for leg in legs}) + target_by_leg: Final = MappingProxyType({leg.id: (leg.target_type, leg.target_id) for leg in legs}) by_leg: Final = _ATTEMPT_AGG_ROWS.validate_python( await _query_raw(prisma_client, _ATTEMPT_AGG_BY_LEG_SQL, leg_ids) or () ) - by_key: Final = tuple( - row.model_copy(update={"grp": key_by_leg[row.grp]}) # mutable-ok: pydantic update payload - for row in by_leg + verdicts_by_target: Final[Mapping[tuple[str, str], ShadowEvalSlice]] = MappingProxyType( + { # mutable-ok: MappingProxyType needs a dict to wrap + target_by_leg[slice.group]: slice.model_copy( + update={"group": target_by_leg[slice.group][1]} # mutable-ok: pydantic update payload + ) + for slice in _slices(by_leg) + } ) total_turns: Final = sum(r.turn_count for r in by_tier) funnel_rows: Final = await _query_raw(prisma_client, _FUNNEL_TOTALS_SQL, leg_ids) @@ -1155,10 +1258,9 @@ async def _shadow_eval_results(prisma_client: "PrismaClient", legs: Sequence[_Le # Coverage only when EVERY leg has a funnel row: a partial seed (one leg's insert # failed) must read as unknown, not as job-level counts missing a leg's traffic. funnel: Final = counted if counted is not None and counted.legs_with_rows == len(leg_ids) else None - return ShadowEvalResult( + result: Final = ShadowEvalResult( by_tier=_slices(by_tier), by_current_model=_slices(by_model), - by_key=_slices(by_key), overall_shadow_win_rate_pct=_pct_of(sum(r.shadow_wins for r in by_tier), total_turns), overall_tie_rate_pct=_pct_of(sum(r.ties for r in by_tier), total_turns), sampled_real_spend=sum(r.real_spend for r in by_tier), @@ -1168,6 +1270,7 @@ async def _shadow_eval_results(prisma_client: "PrismaClient", legs: Sequence[_Le shed_count=funnel.shed if funnel is not None else None, withheld_count=funnel.withheld if funnel is not None else None, ) + return result, verdicts_by_target @router.post( @@ -1182,22 +1285,29 @@ async def start_shadow_eval( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ) -> ShadowEvalJobResponse: """ - Start a shadow eval: duplicate a sampled slice of one or more keys' live traffic against - a second arm, judge the two responses blind, and stratify win rates by tier, by the model - that served the real arm, and by key. + Start a shadow eval: duplicate a sampled slice of one or more targets' live traffic + against a second arm, judge the two responses blind, and stratify win rates by tier, + by the model that served the real arm, and by target. - A forward job answers whether the keys should adopt router_name: it samples the requests - the router did not serve and duplicates them through it. A reverse job answers whether a - key already on the router still gains from it: it samples the requests the router did - serve and duplicates them against baseline_model. A key can hold one active job per - direction, so both questions can run at once. + A target is a virtual key, a team, or a user. Team and user targets match on the + identity every request resolves to at auth time, so they cover JWT-authenticated + traffic, which presents no virtual key; a user target samples that user's traffic + across all their teams, whether it arrives on a JWT or a key they own. - Shadow responses are never served to users. Each key samples until its recorded eval - spend, the shadow and judge calls' own cost, reaches max_budget dollars, the job's - window ends, or the job is stopped, so one key running out of budget does not end - sampling for the others; sampling changes propagate to pods within about 10 seconds. - Shadow and judge calls bill to the shadowed key but are excluded from request counts - and auto-router adoption metrics. + A forward job answers whether the targets should adopt router_name: it samples the + requests the router did not serve and duplicates them through it. A reverse job + answers whether a target already on the router still gains from it: it samples the + requests the router did serve and duplicates them against baseline_model. A target + can hold one active job per direction, so both questions can run at once, and a + request matching several jobs' targets (say its key and its team) is sampled by + each, separately budgeted. + + Shadow responses are never served to users. Each target samples until its recorded + eval spend, the shadow and judge calls' own cost, reaches max_budget dollars, the + job's window ends, or the job is stopped, so one target running out of budget does + not end sampling for the others; sampling changes propagate to pods within about 10 + seconds. Shadow and judge calls bill to the sampled request's own identity but are + excluded from request counts and auto-router adoption metrics. """ from litellm.proxy.proxy_server import llm_router, prisma_client @@ -1206,35 +1316,88 @@ async def start_shadow_eval( raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) if llm_router is None or not _is_configured_pre_routing_strategy(llm_router, data.router_name): raise HTTPException(status_code=400, detail=f"'{data.router_name}' is not a configured auto-router") - token_rows: Final = await _verification_tokens(prisma_client).find_many( - where={"token": {"in": list(data.api_key_ids)}} # mutable-ok: Prisma filter - ) - unknown: Final = tuple(sorted(frozenset(data.api_key_ids) - frozenset(row.token for row in token_rows or ()))) - if unknown: - raise HTTPException( - status_code=400, - detail=( - f"api_key_ids not on this proxy: {', '.join(unknown)}; pass each key's token hash, " - "the value the key list and key info endpoints report" - ), + token_rows: Final = ( + await _verification_tokens(prisma_client).find_many( + where={"token": {"in": list(data.api_key_ids)}} # mutable-ok: Prisma filter ) + if data.api_key_ids + else () + ) + team_rows: Final = ( + await _team_rows(prisma_client).find_many( + where={"team_id": {"in": list(data.team_ids)}} # mutable-ok: Prisma filter + ) + if data.team_ids + else () + ) + user_rows: Final = ( + await _user_rows(prisma_client).find_many( + where={"user_id": {"in": list(data.user_ids)}} # mutable-ok: Prisma filter + ) + if data.user_ids + else () + ) + unknown_keys: Final = sorted(frozenset(data.api_key_ids) - frozenset(row.token for row in token_rows or ())) + unknown_teams: Final = sorted(frozenset(data.team_ids) - frozenset(row.team_id for row in team_rows or ())) + unknown_users: Final = sorted(frozenset(data.user_ids) - frozenset(row.user_id for row in user_rows or ())) + unknown_parts: Final = tuple( + part + for part in ( + ( + f"api_key_ids not on this proxy: {', '.join(unknown_keys)}; pass each key's token hash, " + "the value the key list and key info endpoints report" + ) + if unknown_keys + else None, + f"team_ids not on this proxy: {', '.join(unknown_teams)}" if unknown_teams else None, + f"user_ids not on this proxy: {', '.join(unknown_users)}" if unknown_users else None, + ) + if part is not None + ) + if unknown_parts: + raise HTTPException(status_code=400, detail=". ".join(unknown_parts)) # Every model check below runs once per team the job samples for, since that is the # identity the shadow and judge calls carry and therefore what the router selects on. - team_ids: Final = tuple(dict.fromkeys(row.team_id for row in token_rows or ())) + # A user target's traffic can span teams, so it validates unscoped (None); each + # sampled attempt still resolves the judge under its own request's team at eval time. + team_ids: Final = tuple( + dict.fromkeys( + ( + *(row.team_id for row in token_rows or ()), + *data.team_ids, + *((None,) if data.user_ids else ()), + ) + ) + ) _validate_plain_model(llm_router, data.judge_model, "judge_model", team_ids) if data.baseline_model is not None: _validate_plain_model(llm_router, data.baseline_model, "baseline_model", team_ids) _validate_judge_is_not_a_candidate(llm_router, data, team_ids) + requested_targets: Final[tuple[tuple[ShadowEvalTargetType, str], ...]] = ( + *(("key", key) for key in data.api_key_ids), + *(("team", team) for team in data.team_ids), + *(("user", user) for user in data.user_ids), + ) + requested_by_type: Final[tuple[tuple[ShadowEvalTargetType, tuple[str, ...]], ...]] = tuple( + (target_type, ids) + for target_type, ids in (("key", data.api_key_ids), ("team", data.team_ids), ("user", data.user_ids)) + if ids + ) # A job whose window passed or whose budget ran out stopped sampling on its own, - # but its legs still hold their slots in the per-key, per-direction partial unique index - # until stamped; free them so a new eval can start. Sweeping both directions is deliberate. - requested: Final = list(data.api_key_ids) # mutable-ok: query param - await prisma_client.db.execute_raw(_SWEEP_FINISHED_JOBS_SQL, requested) + # but its legs still hold their slots in the per-target, per-direction partial unique + # index until stamped; free them so a new eval can start. Sweeping both directions is + # deliberate. Sweep and claim filter on exact (target_type, id) pairs so a team id + # that happens to equal a key hash never matches the other kind's slot. + for target_type, ids in requested_by_type: + await prisma_client.db.execute_raw(_SWEEP_FINISHED_JOBS_SQL, list(ids), target_type) # mutable-ok: query param claimed: Final = await _shadow_eval_jobs(prisma_client).find_many( where={ # mutable-ok: Prisma filter - "api_key_id": {"in": requested}, # mutable-ok: Prisma filter + "OR": [ # mutable-ok: Prisma filter + {"target_type": target_type, "target_id": {"in": list(ids)}} # mutable-ok: Prisma filter + for target_type, ids in requested_by_type + ], "direction": data.direction, "stopped_at": None, }, @@ -1244,7 +1407,7 @@ async def start_shadow_eval( status_code=409, detail=( f"Already in an active {data.direction} shadow eval job: " - + ", ".join(sorted(f"{row.api_key_id} (job {row.group_id})" for row in claimed)) + + ", ".join(sorted(f"{row.target_type} {row.target_id} (job {row.group_id})" for row in claimed)) + ". Stop it first." ), ) @@ -1268,10 +1431,16 @@ async def start_shadow_eval( # Leg ids are minted here rather than by the DB default so the funnel seed below # writes from the same values with no read-back, which a lagging read replica # (DATABASE_URL_READ_REPLICA) could otherwise return empty. - leg_ids: Final = tuple(str(uuid4()) for _ in data.api_key_ids) + leg_ids: Final = tuple(str(uuid4()) for _ in requested_targets) await _shadow_eval_jobs(prisma_client).create_many( data=[ # mutable-ok: Prisma payload - {**shared_config, "id": leg_id, "api_key_id": key} for leg_id, key in zip(leg_ids, data.api_key_ids) + { # mutable-ok: Prisma payload + **shared_config, + "id": leg_id, + "target_type": target_type, + "target_id": target_id, + } # mutable-ok: Prisma payload + for leg_id, (target_type, target_id) in zip(leg_ids, requested_targets) ] ) except Exception as e: @@ -1280,7 +1449,8 @@ async def start_shadow_eval( raise HTTPException( status_code=409, detail=( - f"A requested key was claimed by another {data.direction} shadow eval job concurrently. Stop it first." + f"A requested target was claimed by another {data.direction} shadow eval job concurrently. " + "Stop it first." ), ) from e # Seed a zero funnel row per leg NOW: a fully covered job never skips a request, so @@ -1293,18 +1463,19 @@ async def start_shadow_eval( ) except Exception as seed_err: # noqa: BLE001 # coverage is advisory; the job must still start verbose_proxy_logger.error("shadow_eval: funnel seed failed for job %s: %s", group_id, seed_err) - labels: Final = MappingProxyType({row.token: row for row in token_rows}) + labels: Final = _target_labels(token_rows or (), team_rows or (), user_rows or ()) return ShadowEvalJobResponse( job_id=group_id, - keys=tuple( - ShadowEvalJobKeyResponse( - api_key_id=api_key_id, + targets=tuple( + ShadowEvalJobTargetResponse( + target_type=target_type, + target_id=target_id, max_turns=SHADOW_EVAL_TURN_VALVE, max_budget=data.max_budget, - key_alias=labels[api_key_id].key_alias, - key_name=labels[api_key_id].key_name, + target_alias=labels.get((target_type, target_id), _NO_TARGET_LABELS)[0], + key_name=labels.get((target_type, target_id), _NO_TARGET_LABELS)[1], ) - for api_key_id in sorted(data.api_key_ids) + for target_type, target_id in sorted(requested_targets) ), router_name=data.router_name, direction=data.direction, @@ -1324,22 +1495,29 @@ async def start_shadow_eval( ) async def list_shadow_eval_jobs( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], - api_key_id: Annotated[ - str | None, Query(description="Filter to jobs that shadow this key, alone or alongside others") + target_type: Annotated[ + ShadowEvalTargetType | None, Query(description="Kind of target to filter on; requires target_id") + ] = None, + target_id: Annotated[ + str | None, Query(description="Filter to jobs that shadow this target, alone or alongside others") ] = None, limit: Annotated[int, Query(ge=1, le=200, description="Newest jobs to return")] = 50, ) -> tuple[ShadowEvalJobResponse, ...]: - """List shadow eval jobs, newest first, each key with its attempt count so status is - accurate. Judged counts, spend, and results ride the detail endpoint only.""" + """List shadow eval jobs, newest first, each target with its attempt count so status + is accurate. Judged counts, spend, and results ride the detail endpoint only.""" from litellm.proxy.proxy_server import prisma_client _require_admin_viewer(user_api_key_dict, "view shadow evals") if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) + filter_type: Final = target_type if isinstance(target_type, str) else None + filter_id: Final = target_id if isinstance(target_id, str) else None + if (filter_type is None) != (filter_id is None): + raise HTTPException(status_code=400, detail="target_type and target_id filter together; pass both or neither") legs: Final = _LEG_ROWS.validate_python( ( - await _query_raw(prisma_client, _LIST_LEGS_BY_KEY_SQL, limit, api_key_id) - if api_key_id + await _query_raw(prisma_client, _LIST_LEGS_BY_TARGET_SQL, limit, filter_type, filter_id) + if filter_type and filter_id else await _query_raw(prisma_client, _LIST_LEGS_SQL, limit) ) or () @@ -1354,7 +1532,7 @@ async def list_shadow_eval_jobs( by_group, key=lambda group_id: max(leg.created_at for leg in by_group[group_id]), reverse=True ) counts: Final = await _leg_attempt_counts(prisma_client, legs) - return await _with_key_labels( + return await _with_target_labels( prisma_client, tuple(_group_response(group_id, by_group[group_id], counts) for group_id in newest_first) ) @@ -1391,16 +1569,25 @@ async def get_shadow_eval_job( where={"job_id": {"in": leg_ids}, "outcome": "error"}, # mutable-ok: Prisma filter order={"created_at": "desc"}, # mutable-ok: Prisma order ) - labeled: Final = await _with_key_labels( + labeled: Final = await _with_target_labels( prisma_client, (_group_response(job_id, legs, await _leg_attempt_counts(prisma_client, legs)),) ) + results, verdicts_by_target = await _shadow_eval_results(prisma_client, legs) return labeled[0].model_copy( update={ # mutable-ok: pydantic update payload "judged_count": totals[0].judged_count if totals else 0, "error_count": totals[0].error_count if totals else 0, "judge_spend": round(totals[0].judge_spend, 6) if totals else 0.0, "last_error": latest_error.error if latest_error else None, - "results": await _shadow_eval_results(prisma_client, legs), + "results": results, + "targets": tuple( + target.model_copy( + update={ # mutable-ok: pydantic update payload + "verdicts": verdicts_by_target.get((target.target_type, target.target_id)) + } + ) + for target in labeled[0].targets + ), } ) @@ -1415,8 +1602,8 @@ async def stop_shadow_eval_job( job_id: str, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ) -> ShadowEvalJobResponse: - """Stop an active shadow eval job, every key it scopes at once. Attempts are kept; - sampling halts within ~10s. Keys that already stopped on their own budget keep the + """Stop an active shadow eval job, every target it scopes at once. Attempts are kept; + sampling halts within ~10s. Targets that already stopped on their own budget keep the stopped_at they earned. The statement is the whole state machine: it claims the job only while a leg still samples inside the window with no stop recorded, so a racing operator, a same-instant budget spend, and a repeat stop all read the same 400 with @@ -1443,5 +1630,5 @@ async def stop_shadow_eval_job( current: Final = _group_response(job_id, legs, counts) if claimed == 0: raise HTTPException(status_code=400, detail=f"Job {job_id} is already {current.status}") - labeled: Final = await _with_key_labels(prisma_client, (current,)) + labeled: Final = await _with_target_labels(prisma_client, (current,)) return labeled[0] diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 60223265211..01a607b68a9 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1529,14 +1529,15 @@ model LiteLLM_AutoRouterSession { model LiteLLM_ShadowEvalJob { id String @id @default(cuid()) group_id String // legs of one job share this; the API's job id - api_key_id String // hashed virtual key whose traffic this leg shadows + target_type String @default("key") // key | team | user + target_id String // hashed virtual key, team_id, or user_id whose traffic this leg shadows router_name String // the auto-router under evaluation, in either direction direction String @default("forward") // forward | reverse baseline_model String? // reverse only: the fixed model the router is judged against judge_model String shadow_percentage Float max_turns Int // sample-count ceiling: the whole budget on pre-max_budget jobs, the error-loop valve otherwise - max_budget Float? // per-key USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets + max_budget Float? // per-target USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets created_at DateTime @default(now()) created_by String? ends_at DateTime @@ -1544,7 +1545,7 @@ model LiteLLM_ShadowEvalJob { stopped_by String? // operator who stopped it early; null when it ended on its own @@index([group_id]) - @@index([api_key_id]) + @@index([target_type, target_id]) @@index([created_at]) } diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index bde3f5f9e7e..dfc45ccf9bf 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -245,6 +245,8 @@ ShadowEvalStatus: TypeAlias = Literal["running", "completed", "stopped"] ShadowEvalDirection: TypeAlias = Literal["forward", "reverse"] +ShadowEvalTargetType: TypeAlias = Literal["key", "team", "user"] + DEFAULT_SHADOW_EVAL_JUDGE_MODEL: Final[str] = "anthropic/claude-sonnet-5" # Sample-count ceiling written on every new job: a zero-cost error loop (a shadow arm that @@ -253,16 +255,37 @@ SHADOW_EVAL_TURN_VALVE: Final[int] = 10_000 class StartShadowEvalRequest(BaseModel): - """Start duplicating one or more keys' traffic for blind comparison against an auto-router.""" + """Start duplicating one or more targets' traffic for blind comparison against an auto-router. + + A target is a virtual key, a team, or a user; each becomes its own leg with its own + budget and stop state. Team and user targets match on the identity every request + carries after auth (user_api_key_team_id / user_api_key_user_id), so they cover + JWT-authenticated traffic, which presents no virtual key at all.""" api_key_ids: tuple[str, ...] = Field( - min_length=1, + default=(), max_length=100, description=( - "The hashed virtual keys whose traffic will be shadowed. Shadow evaluation runs ONLY on these " - "keys' traffic; requests made with any other key are not sampled. Each key carries its own " - "max_budget spend budget, so one key exhausting its budget leaves the others sampling. At most 100 " - "keys per job, which also bounds every read the job's endpoints make." + "Hashed virtual keys whose traffic will be shadowed. Combined with team_ids and user_ids the job " + "needs at least one target and at most 100, which also bounds every read the job's endpoints make. " + "Each target carries its own max_budget spend budget, so one exhausting its budget leaves the " + "others sampling." + ), + ) + team_ids: tuple[str, ...] = Field( + default=(), + max_length=100, + description=( + "Teams whose traffic will be shadowed, matched on the team every authenticated request resolves " + "to, so a team's JWT-auth and virtual-key traffic are both sampled" + ), + ) + user_ids: tuple[str, ...] = Field( + default=(), + max_length=100, + description=( + "Users whose traffic will be shadowed, matched on the user every authenticated request resolves " + "to across all their teams: JWT requests carrying their subject claim and virtual keys they own" ), ) router_name: str = Field(description="The auto-router under evaluation, in either direction") @@ -285,7 +308,7 @@ class StartShadowEvalRequest(BaseModel): shadow_percentage: float = Field( ge=0.1, le=100.0, - description="Percentage of the key's requests to duplicate through the router", + description="Percentage of each target's requests to duplicate through the router", ) judge_model: str = Field( default=DEFAULT_SHADOW_EVAL_JUDGE_MODEL, @@ -306,9 +329,9 @@ class StartShadowEvalRequest(BaseModel): ge=0.01, le=10_000, description=( - "Per-key USD budget for the eval's own overhead, the shadow-arm and judge calls, priced with " - "the same figures the spend pipeline bills. EACH scoped key samples until its recorded eval " - "spend reaches this, so a job over N keys spends at most about N times max_budget; in-flight " + "Per-target USD budget for the eval's own overhead, the shadow-arm and judge calls, priced with " + "the same figures the spend pipeline bills. EACH scoped target samples until its recorded eval " + "spend reaches this, so a job over N targets spends at most about N times max_budget; in-flight " "samples can overshoot the cap by one sampling cache window" ), ) @@ -319,7 +342,7 @@ class StartShadowEvalRequest(BaseModel): """Pydantic ignores unknown fields, so a caller still sending max_turns would silently run on the default dollar budget instead of the bound they asked for.""" if isinstance(values, Mapping) and "max_turns" in values: - raise ValueError("max_turns was replaced by max_budget, the per-key USD cap on the eval's own spend") + raise ValueError("max_turns was replaced by max_budget, the per-target USD cap on the eval's own spend") return values @field_validator("shadow_percentage") @@ -327,12 +350,21 @@ class StartShadowEvalRequest(BaseModel): def _round_percentage(cls, value: float) -> float: return round(value, 2) - @field_validator("api_key_ids") + @field_validator("api_key_ids", "team_ids", "user_ids") @classmethod - def _dedupe_keys(cls, value: tuple[str, ...]) -> tuple[str, ...]: - """A key named twice would collide with itself on the one-active-per-(key, direction) index.""" + def _dedupe_targets(cls, value: tuple[str, ...]) -> tuple[str, ...]: + """A target named twice would collide with itself on the one-active-per-(target, direction) index.""" return tuple(dict.fromkeys(value)) + @model_validator(mode="after") + def _at_least_one_target_at_most_hundred(self) -> "StartShadowEvalRequest": + total: Final = len(self.api_key_ids) + len(self.team_ids) + len(self.user_ids) + if total < 1: + raise ValueError("at least one target is required: pass api_key_ids, team_ids, or user_ids") + if total > 100: + raise ValueError("at most 100 targets per job across api_key_ids, team_ids, and user_ids") + return self + @model_validator(mode="after") def _baseline_model_matches_direction(self) -> "StartShadowEvalRequest": if self.direction == "reverse" and self.baseline_model is None: @@ -343,8 +375,9 @@ class StartShadowEvalRequest(BaseModel): class ShadowEvalSlice(BaseModel): - """Judge outcomes for one slice of a job's verdicts (a router tier, or one of the - models that served the real arm).""" + """Judge outcomes for one slice of a job's verdicts: a router tier, one of the + models that served the real arm, or one scoped target (embedded on that target's + own entry, so slices never need re-joining to a target by id).""" group: str turn_count: int @@ -395,12 +428,6 @@ class ShadowEvalResult(BaseModel): "and in reverse the models the router itself picked" ) ) - by_key: tuple[ShadowEvalSlice, ...] = Field( - description=( - "One slice per scoped key that has judged verdicts, grouped on the raw key hash. Keys the job " - "scopes but has not judged a turn for yet are absent rather than reported as zero" - ), - ) overall_shadow_win_rate_pct: float overall_tie_rate_pct: float sampled_real_spend: float = Field( @@ -436,27 +463,28 @@ class ShadowEvalResult(BaseModel): ) -class ShadowEvalJobKeyResponse(BaseModel): - """One key a job shadows, with its own budget and stop state.""" +class ShadowEvalJobTargetResponse(BaseModel): + """One target a job shadows (a key, team, or user), with its own budget and stop state.""" - api_key_id: str = Field(description="The hashed virtual key whose traffic this entry scopes") + target_type: ShadowEvalTargetType = Field(description="What kind of entity this entry scopes") + target_id: str = Field(description="The hashed virtual key, team id, or user id whose traffic this entry scopes") max_turns: int = Field( description=( - "This key's sample-count ceiling: the whole budget for jobs created before max_budget " + "This target's sample-count ceiling: the whole budget for jobs created before max_budget " "existed, and the error-loop safety valve otherwise" ) ) max_budget: float | None = Field( default=None, description=( - "This key's own USD budget for the eval's shadow and judge spend, independent of its " + "This target's own USD budget for the eval's shadow and judge spend, independent of its " "siblings'; None on jobs created before spend budgets existed, which max_turns alone bounds" ), ) stopped_at: datetime | None = Field( default=None, description=( - "When this key's slot was stamped free, whether its own budget ran out, the window closed, " + "When this target's slot was stamped free, whether its own budget ran out, the window closed, " "or an operator stopped the job; status is derived, so a spent budget reads completed even " "while this is still unset" ), @@ -464,45 +492,53 @@ class ShadowEvalJobKeyResponse(BaseModel): attempt_count: int | None = Field( default=None, description=( - "This key's sampled attempts so far, judged and errored alike, the same count the sampler " + "This target's sampled attempts so far, judged and errored alike, the same count the sampler " "budgets against max_turns; populated on list and detail responses. Frozen at stopped_at " - "once the key is stamped, so in-flight attempts landing after a stop never reclassify it" + "once the target is stamped, so in-flight attempts landing after a stop never reclassify it" ), ) spend: float | None = Field( default=None, description=( - "This key's recorded shadow plus judge spend in USD, the same figure the sampler budgets " + "This target's recorded shadow plus judge spend in USD, the same figure the sampler budgets " "against max_budget; populated on list and detail responses and frozen at stopped_at " "exactly like attempt_count" ), ) + verdicts: "ShadowEvalSlice | None" = Field( + default=None, + description="This target's own judged-verdict slice; detail endpoint only, None until a turn is judged", + ) + @property def budget_spent(self) -> bool: over_spend: Final = self.max_budget is not None and self.spend is not None and self.spend >= self.max_budget return over_spend or (self.attempt_count is not None and self.attempt_count >= self.max_turns) - key_alias: str | None = Field( + target_alias: str | None = Field( default=None, - description="Alias of the shadowed key, resolved from the key row at read time; None when unset or deleted", + description=( + "Display label resolved from the target's own row at read time: the key's alias, the team's " + "alias, or the user's email; None when unset or deleted" + ), ) key_name: str | None = Field( default=None, - description="Masked display name (sk-...) of the shadowed key, resolved at read time like key_alias", + description="Masked display name (sk-...) for key targets, resolved at read time; None for teams and users", ) class ShadowEvalJobResponse(BaseModel): - """A shadow-eval job over one or more keys, each with its own budget and stop state; - status is derived from stopped_by, the keys' stop and budget state, and ends_at, + """A shadow-eval job over one or more targets, each with its own budget and stop state; + status is derived from stopped_by, the targets' stop and budget state, and ends_at, never stored, so no writer anywhere can produce an inconsistent one. Aggregate fields are populated by the detail endpoint only and stay None on list responses.""" job_id: str - keys: tuple[ShadowEvalJobKeyResponse, ...] = Field( + targets: tuple[ShadowEvalJobTargetResponse, ...] = Field( min_length=1, - description="The keys whose traffic this job evaluates, and only those keys', each with its own budget", + description="The targets whose traffic this job evaluates, and only theirs, each with its own budget", ) router_name: str direction: ShadowEvalDirection = "forward" @@ -531,8 +567,8 @@ class ShadowEvalJobResponse(BaseModel): def status(self) -> ShadowEvalStatus: """Three recorded facts, no history-guessing: a stop is stopped_by (the migration backfills it for every job that displayed stopped when the column arrived, so the - pre-column population is closed), completion is the window passing or every key - spending its budget, and anything else is running. The all-keys-stamped fallback + pre-column population is closed), completion is the window passing or every target + spending its budget, and anything else is running. The all-targets-stamped fallback covers only stops written by pre-column pods during a rolling deploy.""" if self.stopped_by is not None: return "stopped" @@ -540,8 +576,8 @@ class ShadowEvalJobResponse(BaseModel): self.ends_at if self.ends_at.tzinfo else self.ends_at.replace(tzinfo=timezone.utc) ): return "completed" - if all(key.budget_spent for key in self.keys): + if all(target.budget_spent for target in self.targets): return "completed" - if all(key.stopped_at is not None for key in self.keys): + if all(target.stopped_at is not None for target in self.targets): return "stopped" return "running" diff --git a/schema.prisma b/schema.prisma index 60223265211..01a607b68a9 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1529,14 +1529,15 @@ model LiteLLM_AutoRouterSession { model LiteLLM_ShadowEvalJob { id String @id @default(cuid()) group_id String // legs of one job share this; the API's job id - api_key_id String // hashed virtual key whose traffic this leg shadows + target_type String @default("key") // key | team | user + target_id String // hashed virtual key, team_id, or user_id whose traffic this leg shadows router_name String // the auto-router under evaluation, in either direction direction String @default("forward") // forward | reverse baseline_model String? // reverse only: the fixed model the router is judged against judge_model String shadow_percentage Float max_turns Int // sample-count ceiling: the whole budget on pre-max_budget jobs, the error-loop valve otherwise - max_budget Float? // per-key USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets + max_budget Float? // per-target USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets created_at DateTime @default(now()) created_by String? ends_at DateTime @@ -1544,7 +1545,7 @@ model LiteLLM_ShadowEvalJob { stopped_by String? // operator who stopped it early; null when it ended on its own @@index([group_id]) - @@index([api_key_id]) + @@index([target_type, target_id]) @@index([created_at]) } diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index f9c287fc7b7..1af3dd3f613 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -58,11 +58,12 @@ def _prisma(jobs=(), attempt_counts=(), attempt_costs=()) -> MagicMock: return prisma -def _job_record(job: ActiveShadowEvalJob, api_key_id="key-hash") -> MagicMock: +def _job_record(job: ActiveShadowEvalJob, target_type="key", target_id="key-hash") -> MagicMock: record = MagicMock() for field, value in dict( id=job.id, - api_key_id=api_key_id, + target_type=target_type, + target_id=target_id, router_name=job.router_name, direction=job.direction, baseline_model=job.baseline_model, @@ -123,7 +124,7 @@ def _spend_counter(store=None): return counter, read, write -def _logger(router=None, prisma=None, jobs=(), counter_store=None) -> ShadowEvalLogger: +def _logger(router=None, prisma=None, jobs=(), counter_store=None, jobs_by_target=None) -> ShadowEvalLogger: cache = InMemoryCache(max_size_in_memory=4, default_ttl=60) counter, read, write = _spend_counter(counter_store) funnel_events = [] @@ -137,8 +138,9 @@ def _logger(router=None, prisma=None, jobs=(), counter_store=None) -> ShadowEval ) logger._test_counter = counter logger._test_funnel = funnel_events - if jobs: - cache.set_cache("shadow_eval:active_jobs", {"key-hash": tuple(jobs)}) + seeded = jobs_by_target if jobs_by_target is not None else ({("key", "key-hash"): tuple(jobs)} if jobs else None) + if seeded is not None: + cache.set_cache("shadow_eval:active_jobs", seeded) return logger @@ -837,6 +839,86 @@ class TestSuccessHookSkipChain: prisma.db.litellm_shadowevalattempt.create.assert_not_called() +JWT_IDENTITY = {"user_api_key_hash": None, "user_api_key_team_id": "team-eng", "user_api_key_user_id": "dev-alice"} + + +@pytest.mark.asyncio +class TestTargetMatching: + """A request qualifies for a job through ANY of its resolved identities: key hash, + team id, or user id. Team and user jobs must therefore sample JWT-authenticated + traffic, which carries no key hash at all.""" + + @pytest.mark.parametrize( + "target,sampled", + [ + (("team", "team-eng"), True), + (("user", "dev-alice"), True), + (("key", "some-key"), False), + ], + ids=["team-job-samples-jwt-traffic", "user-job-samples-jwt-traffic", "key-jobs-never-match-keyless-traffic"], + ) + async def test_jwt_shaped_traffic_matches_team_and_user_jobs_but_no_key_job(self, target, sampled): + prisma = _prisma() + router = _router() + logger = _logger(router=router, prisma=prisma, jobs_by_target={target: (_job(),)}) + hook_kwargs = _success_kwargs() + hook_kwargs["standard_logging_object"]["metadata"] = dict(JWT_IDENTITY) + + await logger.async_log_success_event(hook_kwargs, RESPONSE, None, None) + await _drain(logger) + + if sampled: + prisma.db.litellm_shadowevalattempt.create.assert_awaited_once() + assert prisma.db.litellm_shadowevalattempt.create.call_args.kwargs["data"]["job_id"] == "job-1" + else: + router.acompletion.assert_not_called() + prisma.db.litellm_shadowevalattempt.create.assert_not_called() + + async def test_an_event_with_no_identity_early_returns_without_a_cache_read(self): + prisma = _prisma() + router = _router() + cache = MagicMock(spec=InMemoryCache) + cache.async_get_cache = AsyncMock() + logger = ShadowEvalLogger( + router_provider=lambda: router, + prisma_provider=lambda: prisma, + jobs_cache=cache, + ) + hook_kwargs = _success_kwargs() + hook_kwargs["standard_logging_object"]["metadata"] = {} + + await logger.async_log_success_event(hook_kwargs, RESPONSE, None, None) + + cache.async_get_cache.assert_not_awaited() + router.acompletion.assert_not_called() + prisma.db.litellm_shadowevalattempt.create.assert_not_called() + + async def test_an_event_matching_a_key_job_and_a_team_job_fires_both(self): + """A request's key and its team can each hold a job; the two are separately + budgeted experiments, so both fire and each counts its own start.""" + prisma = _prisma() + logger = _logger( + router=_router(), + prisma=prisma, + jobs_by_target={ + ("key", "key-hash"): (_job(id="key-job"),), + ("team", "team-eng"): (_job(id="team-job"),), + }, + ) + hook_kwargs = _success_kwargs() + hook_kwargs["standard_logging_object"]["metadata"] = { + "user_api_key_hash": "key-hash", + "user_api_key_team_id": "team-eng", + } + + await logger.async_log_success_event(hook_kwargs, RESPONSE, None, None) + await _drain(logger) + + rows = [call.kwargs["data"] for call in prisma.db.litellm_shadowevalattempt.create.call_args_list] + assert sorted(row["job_id"] for row in rows) == ["key-job", "team-job"] + assert logger._job_starts == {"key-job": 1, "team-job": 1} + + @pytest.mark.asyncio class TestActiveJobsCache: async def test_cache_miss_reads_db_once_then_serves_from_cache(self): @@ -851,8 +933,8 @@ class TestActiveJobsCache: first = await logger._active_jobs() second = await logger._active_jobs() - assert [job.id for job in first["key-hash"]] == ["job-1"] - assert second["key-hash"][0].attempts == 7 + assert [job.id for job in first[("key", "key-hash")]] == ["job-1"] + assert second[("key", "key-hash")][0].attempts == 7 assert prisma.db.litellm_shadowevaljob.find_many.await_count == 1 where = prisma.db.litellm_shadowevaljob.find_many.call_args.kwargs["where"] assert where["stopped_at"] is None @@ -899,8 +981,8 @@ class TestActiveJobsCache: jobs = await logger._active_jobs() assert logger._job_starts == {} - assert jobs["key-hash"][0].attempts == 7 - assert jobs["key-hash"][0].spend == 0.05 + assert jobs[("key", "key-hash")][0].attempts == 7 + assert jobs[("key", "key-hash")][0].spend == 0.05 @pytest.mark.asyncio @@ -1249,13 +1331,14 @@ class TestActiveJobsFailClosed: jobs_cache=InMemoryCache(max_size_in_memory=4, default_ttl=60), ) - assert [job.id for job in (await logger._active_jobs())["key-hash"]] == ["job-ok"] + assert [job.id for job in (await logger._active_jobs())[("key", "key-hash")]] == ["job-ok"] - async def test_both_of_a_key_s_jobs_survive_the_lookup(self): + async def test_every_targets_jobs_survive_the_lookup_keyed_by_type_and_id(self): records = [ _job_record(_job(id="job-forward")), _job_record(_reverse_job(id="job-reverse")), - _job_record(_job(id="job-other"), api_key_id="other-key"), + _job_record(_job(id="job-other"), target_id="other-key"), + _job_record(_job(id="job-team"), target_type="team", target_id="team-eng"), ] prisma = _prisma(jobs=records, attempt_counts=[("job-reverse", 3)]) logger = ShadowEvalLogger( @@ -1266,9 +1349,11 @@ class TestActiveJobsFailClosed: jobs = await logger._active_jobs() - assert sorted(job.id for job in jobs["key-hash"]) == ["job-forward", "job-reverse"] - assert [job.id for job in jobs["other-key"]] == ["job-other"] - assert {job.id: job.attempts for job in jobs["key-hash"]}["job-reverse"] == 3 + assert sorted(job.id for job in jobs[("key", "key-hash")]) == ["job-forward", "job-reverse"] + assert [job.id for job in jobs[("key", "other-key")]] == ["job-other"] + assert [job.id for job in jobs[("team", "team-eng")]] == ["job-team"] + assert ("team-eng",) not in jobs and "team-eng" not in jobs + assert {job.id: job.attempts for job in jobs[("key", "key-hash")]}["job-reverse"] == 3 def _failing_router(): diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 726e09f3162..a74aa553449 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -878,7 +878,8 @@ def _leg_record(**overrides: object) -> MagicMock: defaults = { "id": "leg-1", "group_id": "job-1", - "api_key_id": "key-hash", + "target_type": "key", + "target_id": "key-hash", "router_name": "my-router", "direction": "forward", "baseline_model": None, @@ -912,8 +913,28 @@ def _key_record( return record +def _team_record(team_id: str, team_alias: str | None) -> MagicMock: + record = MagicMock(spec=["team_id", "team_alias"]) + record.team_id = team_id + record.team_alias = team_alias + return record + + +def _user_record(user_id: str, user_email: str | None) -> MagicMock: + record = MagicMock(spec=["user_id", "user_email"]) + record.user_id = user_id + record.user_email = user_email + return record + + def _shadow_prisma( - legs=(), agg_rows=None, by_leg_rows=None, known_keys=("key-hash", "key-hash-2"), key_teams=None + legs=(), + agg_rows=None, + by_leg_rows=None, + known_keys=("key-hash", "key-hash-2"), + key_teams=None, + known_teams=None, + known_users=None, ) -> MagicMock: """The job-table fake honours the filters it is handed, so a read that forgets stopped_at sees rows the partial index would have released, one that forgets @@ -921,6 +942,8 @@ def _shadow_prisma( group read that matched on a leg id would come back empty.""" prisma = MagicMock() teams: Final = key_teams or {} + team_aliases: Final = known_teams or {} + user_emails: Final = known_users or {} async def find_tokens(*, where): """Honours the token filter, like the job-table fake below: the endpoint derives the @@ -931,6 +954,17 @@ def _shadow_prisma( prisma.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=find_tokens) + async def find_teams(*, where): + requested = where["team_id"]["in"] + return [_team_record(t, alias) for t, alias in team_aliases.items() if t in requested] + + async def find_users(*, where): + requested = where["user_id"]["in"] + return [_user_record(u, email) for u, email in user_emails.items() if u in requested] + + prisma.db.litellm_teamtable.find_many = AsyncMock(side_effect=find_teams) + prisma.db.litellm_usertable.find_many = AsyncMock(side_effect=find_users) + async def execute_raw(sql: str, *params: object): if "SET stopped_by" in sql: group = [row for row in stored if row.group_id == params[0]] @@ -959,9 +993,19 @@ def _shadow_prisma( async def find_many_legs(where=None, **_: object): current = list(stored) w = dict(where or {}) - if "api_key_id" in w: - wanted = w["api_key_id"]["in"] if isinstance(w["api_key_id"], dict) else [w["api_key_id"]] - current = [row for row in current if row.api_key_id in wanted] + if "OR" in w: + pairs = [ + ( + branch["target_type"], + branch["target_id"]["in"] if isinstance(branch["target_id"], dict) else [branch["target_id"]], + ) + for branch in w["OR"] + ] + current = [ + row + for row in current + if any(row.target_type == target_type and row.target_id in ids for target_type, ids in pairs) + ] if "direction" in w: current = [row for row in current if row.direction == w["direction"]] if "stopped_at" in w: @@ -983,7 +1027,8 @@ def _shadow_prisma( fields = ( "id", "group_id", - "api_key_id", + "target_type", + "target_id", "router_name", "direction", "baseline_model", @@ -1009,7 +1054,11 @@ def _shadow_prisma( if "AS attempt_count" in sql: return prisma.attempt_rows if "GROUP BY group_id" in sql: - scoped = [row for row in stored if "api_key_id = $2" not in sql or row.api_key_id == params[1]] + scoped = [ + row + for row in stored + if "target_type = $2" not in sql or (row.target_type == params[1] and row.target_id == params[2]) + ] keep = set(newest_groups(scoped, params[0])) return [leg_dict(row) for row in stored if row.group_id in keep] if "FILTER (WHERE outcome != 'error')::int AS judged_count" in sql: @@ -1058,7 +1107,7 @@ async def test_start_shadow_eval_writes_one_leg_per_key_in_one_statement(monkeyp response = await start_shadow_eval(_start_request(api_key_ids=("key-hash", "key-hash-2")), ADMIN) - sweep_sql, sweep_keys = prisma.db.execute_raw.call_args.args + sweep_sql, sweep_ids, sweep_type = prisma.db.execute_raw.call_args.args assert "stopped_at IS NULL" in sweep_sql assert "j.ends_at <= (NOW() AT TIME ZONE 'utc')" in sweep_sql assert "SET stopped_at = (NOW() AT TIME ZONE 'utc')" in sweep_sql @@ -1066,12 +1115,13 @@ async def test_start_shadow_eval_writes_one_leg_per_key_in_one_statement(monkeyp assert "j.max_budget IS NOT NULL" in sweep_sql assert ">= j.max_budget" in sweep_sql assert "SUM(a.judge_cost + a.shadow_cost + a.shadow_classifier_cost)" in sweep_sql - assert "j.api_key_id = ANY($1::text[])" in sweep_sql - assert sweep_keys == ["key-hash", "key-hash-2"] + assert "j.target_type = $2 AND j.target_id = ANY($1::text[])" in sweep_sql + assert sweep_ids == ["key-hash", "key-hash-2"] + assert sweep_type == "key" prisma.db.litellm_shadowevaljob.create_many.assert_awaited_once() rows = prisma.db.litellm_shadowevaljob.create_many.call_args.kwargs["data"] - assert [row["api_key_id"] for row in rows] == ["key-hash", "key-hash-2"] - assert len({frozenset((k, v) for k, v in row.items() if k not in ("api_key_id", "id")) for row in rows}) == 1 + assert [(row["target_type"], row["target_id"]) for row in rows] == [("key", "key-hash"), ("key", "key-hash-2")] + assert len({frozenset((k, v) for k, v in row.items() if k not in ("target_id", "id")) for row in rows}) == 1 assert len({row["id"] for row in rows}) == len(rows) assert len({row["group_id"] for row in rows}) == 1 assert all(row["max_turns"] == SHADOW_EVAL_TURN_VALVE and row["created_by"] == "admin" for row in rows) @@ -1080,11 +1130,12 @@ async def test_start_shadow_eval_writes_one_leg_per_key_in_one_statement(monkeyp assert response.job_id == rows[0]["group_id"] assert response.status == "running" assert response.judged_count is None - assert [(key.api_key_id, key.max_budget, key.key_alias) for key in response.keys] == [ + assert [(target.target_id, target.max_budget, target.target_alias) for target in response.targets] == [ ("key-hash", 5.0, "prod-alpha"), ("key-hash-2", 5.0, "prod-alpha"), ] - assert all(key.max_turns == SHADOW_EVAL_TURN_VALVE for key in response.keys) + assert all(target.target_type == "key" for target in response.targets) + assert all(target.max_turns == SHADOW_EVAL_TURN_VALVE for target in response.targets) @pytest.mark.asyncio @@ -1232,7 +1283,7 @@ async def test_start_shadow_eval_rejections( import litellm.proxy.proxy_server as proxy_server _configure_anthropic_sdk_judge(monkeypatch) - prisma = _shadow_prisma(legs=[_leg_record(id=f"leg-{key}", group_id="job-7", api_key_id=key) for key in claimed]) + prisma = _shadow_prisma(legs=[_leg_record(id=f"leg-{key}", group_id="job-7", target_id=key) for key in claimed]) monkeypatch.setattr(proxy_server, "prisma_client", prisma) monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) @@ -1315,14 +1366,14 @@ async def test_start_shadow_eval_names_the_busy_key_and_its_job(monkeypatch: pyt import litellm.proxy.proxy_server as proxy_server _configure_anthropic_sdk_judge(monkeypatch) - prisma = _shadow_prisma(legs=[_leg_record(id="leg-b", group_id="job-7", api_key_id="key-hash-2")]) + prisma = _shadow_prisma(legs=[_leg_record(id="leg-b", group_id="job-7", target_id="key-hash-2")]) monkeypatch.setattr(proxy_server, "prisma_client", prisma) monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) with pytest.raises(HTTPException) as exc: await start_shadow_eval(_start_request(api_key_ids=("key-hash", "key-hash-2")), ADMIN) assert exc.value.status_code == 409 - assert "key-hash-2 (job job-7)" in exc.value.detail + assert "key key-hash-2 (job job-7)" in exc.value.detail @pytest.mark.asyncio @@ -1441,6 +1492,180 @@ def test_start_shadow_eval_request_dedupes_and_bounds_the_key_set(): _start_request(api_key_ids=tuple(f"k{i}" for i in range(101))) +def test_start_request_bounds_the_combined_target_count_across_types(): + """The 1..100 bound counts keys, teams, and users together, so a caller cannot dodge + it by spreading targets over the three fields, and a request naming no target of any + type samples nothing and is rejected.""" + with pytest.raises(ValidationError, match="at least one target"): + _start_request(api_key_ids=(), team_ids=(), user_ids=()) + with pytest.raises(ValidationError, match="at most 100 targets"): + _start_request(api_key_ids=tuple(f"k{i}" for i in range(60)), team_ids=tuple(f"t{i}" for i in range(41))) + mixed = _start_request(api_key_ids=tuple(f"k{i}" for i in range(60)), team_ids=tuple(f"t{i}" for i in range(40))) + assert len(mixed.api_key_ids) + len(mixed.team_ids) == 100 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "overrides,prisma_kwargs,expected_target", + [ + ( + {"api_key_ids": (), "team_ids": ("team-eng",)}, + {"known_teams": {"team-eng": "Engineering"}}, + ("team", "team-eng", "Engineering"), + ), + ( + {"api_key_ids": (), "user_ids": ("dev-alice",)}, + {"known_users": {"dev-alice": "alice@example.com"}}, + ("user", "dev-alice", "alice@example.com"), + ), + ], + ids=["team-target-labeled-by-team-alias", "user-target-labeled-by-user-email"], +) +async def test_start_shadow_eval_creates_typed_legs_for_team_and_user_targets( + monkeypatch: pytest.MonkeyPatch, overrides, prisma_kwargs, expected_target +): + import litellm.proxy.proxy_server as proxy_server + + _configure_anthropic_sdk_judge(monkeypatch) + prisma = _shadow_prisma(**prisma_kwargs) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + response = await start_shadow_eval(_start_request(**overrides), ADMIN) + + target_type, target_id, target_alias = expected_target + rows = prisma.db.litellm_shadowevaljob.create_many.call_args.kwargs["data"] + assert [(row["target_type"], row["target_id"]) for row in rows] == [(target_type, target_id)] + assert response.status == "running" + target = response.targets[0] + assert (target.target_type, target.target_id, target.target_alias, target.key_name) == ( + target_type, + target_id, + target_alias, + None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "overrides,prisma_kwargs,expected_detail", + [ + ( + {"api_key_ids": (), "team_ids": ("team-eng", "team-ghost")}, + {"known_teams": {"team-eng": "Engineering"}}, + "team_ids not on this proxy: team-ghost", + ), + ( + {"api_key_ids": (), "user_ids": ("dev-alice", "dev-ghost")}, + {"known_users": {"dev-alice": "alice@example.com"}}, + "user_ids not on this proxy: dev-ghost", + ), + ], + ids=["unknown-team", "unknown-user"], +) +async def test_start_shadow_eval_rejects_teams_and_users_this_proxy_does_not_know( + monkeypatch: pytest.MonkeyPatch, overrides, prisma_kwargs, expected_detail +): + import litellm.proxy.proxy_server as proxy_server + + _configure_anthropic_sdk_judge(monkeypatch) + prisma = _shadow_prisma(**prisma_kwargs) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + with pytest.raises(HTTPException) as exc: + await start_shadow_eval(_start_request(**overrides), ADMIN) + assert exc.value.status_code == 400 + assert expected_detail in exc.value.detail + prisma.db.litellm_shadowevaljob.create_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_start_shadow_eval_mixed_targets_create_both_legs_and_sweep_once_per_type( + monkeypatch: pytest.MonkeyPatch, +): + import litellm.proxy.proxy_server as proxy_server + + _configure_anthropic_sdk_judge(monkeypatch) + prisma = _shadow_prisma(known_teams={"team-eng": "Engineering"}) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + response = await start_shadow_eval(_start_request(team_ids=("team-eng",)), ADMIN) + + rows = prisma.db.litellm_shadowevaljob.create_many.call_args.kwargs["data"] + assert [(row["target_type"], row["target_id"]) for row in rows] == [("key", "key-hash"), ("team", "team-eng")] + assert len({row["group_id"] for row in rows}) == 1 + sweeps = [ + call.args + for call in prisma.db.execute_raw.await_args_list + if "SET stopped_at = (NOW() AT TIME ZONE 'utc')" in call.args[0] + ] + assert [(ids, target_type) for _, ids, target_type in sweeps] == [(["key-hash"], "key"), (["team-eng"], "team")] + assert [(t.target_type, t.target_id, t.target_alias) for t in response.targets] == [ + ("key", "key-hash", "prod-alpha"), + ("team", "team-eng", "Engineering"), + ] + + +@pytest.mark.asyncio +async def test_start_shadow_eval_names_the_busy_team_target(monkeypatch: pytest.MonkeyPatch): + import litellm.proxy.proxy_server as proxy_server + + _configure_anthropic_sdk_judge(monkeypatch) + prisma = _shadow_prisma( + legs=[_leg_record(id="leg-t", group_id="job-7", target_type="team", target_id="team-eng")], + known_teams={"team-eng": "Engineering"}, + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + with pytest.raises(HTTPException) as exc: + await start_shadow_eval(_start_request(api_key_ids=(), team_ids=("team-eng",)), ADMIN) + assert exc.value.status_code == 409 + assert "team team-eng (job job-7)" in exc.value.detail + prisma.db.litellm_shadowevaljob.create_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_start_shadow_eval_claim_matches_exact_target_pairs_not_bare_ids(monkeypatch: pytest.MonkeyPatch): + """A key whose hash happens to spell a team's id must not hold the team's slot: the + claim matches (target_type, target_id) pairs, never ids across kinds.""" + import litellm.proxy.proxy_server as proxy_server + + _configure_anthropic_sdk_judge(monkeypatch) + prisma = _shadow_prisma( + legs=[_leg_record(id="leg-k", group_id="job-7", target_type="key", target_id="team-eng")], + known_teams={"team-eng": "Engineering"}, + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + response = await start_shadow_eval(_start_request(api_key_ids=(), team_ids=("team-eng",)), ADMIN) + + assert response.status == "running" + prisma.db.litellm_shadowevaljob.create_many.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_list_shadow_eval_jobs_rejects_a_lone_filter_half(monkeypatch: pytest.MonkeyPatch): + """target_type and target_id only mean anything together: a bare id could name a key + or a team, and a bare type filters nothing.""" + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma(legs=[_leg_record()]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + with pytest.raises(HTTPException) as id_only: + await list_shadow_eval_jobs(VIEWER, target_type=None, target_id="key-hash", limit=50) + assert id_only.value.status_code == 400 + + with pytest.raises(HTTPException) as type_only: + await list_shadow_eval_jobs(VIEWER, target_type="key", target_id=None, limit=50) + assert type_only.value.status_code == 400 + prisma.db.query_raw.assert_not_called() + + @pytest.mark.asyncio async def test_start_shadow_eval_concurrent_unique_violation_is_a_409(monkeypatch: pytest.MonkeyPatch): import litellm.proxy.proxy_server as proxy_server @@ -1530,7 +1755,7 @@ async def test_get_shadow_eval_job_pools_counts_and_slices_results_per_key(monke }, ] prisma = _shadow_prisma( - legs=[_leg_record(), _leg_record(id="leg-2", api_key_id="key-hash-2", max_turns=50)], + legs=[_leg_record(), _leg_record(id="leg-2", target_id="key-hash-2", max_turns=50)], agg_rows=tier_rows, by_leg_rows=leg_rows, ) @@ -1549,8 +1774,10 @@ async def test_get_shadow_eval_job_pools_counts_and_slices_results_per_key(monke assert response.results.by_tier[0].shadow_win_rate_pct == 50.0 assert response.results.overall_shadow_win_rate_pct == 40.0 assert response.results.overall_tie_rate_pct == 20.0 - assert [(s.group, s.turn_count) for s in response.results.by_key] == [("key-hash", 6), ("key-hash-2", 4)] - assert response.results.by_key[0].shadow_win_rate_pct == 66.7 + verdicts_by_target = {(t.target_type, t.target_id): t.verdicts for t in response.targets} + assert verdicts_by_target[("key", "key-hash")].turn_count == 6 + assert verdicts_by_target[("key", "key-hash")].shadow_win_rate_pct == 66.7 + assert verdicts_by_target[("key", "key-hash-2")].turn_count == 4 agg_sql = next(call.args[0] for call in prisma.db.query_raw.await_args_list if "real_spend" in call.args[0]) assert agg_sql.count("FILTER (WHERE real_cost IS NOT NULL AND NOT real_cache_hit)") == 2 assert response.results.by_tier[0].real_spend == 0.08 @@ -1561,7 +1788,10 @@ async def test_get_shadow_eval_job_pools_counts_and_slices_results_per_key(monke assert response.results.not_sampled_count is None assert response.results.unjudgeable_count is None assert response.results.shed_count is None - assert [(key.api_key_id, key.max_turns) for key in response.keys] == [("key-hash", 200), ("key-hash-2", 50)] + assert [(target.target_id, target.max_turns) for target in response.targets] == [ + ("key-hash", 200), + ("key-hash-2", 50), + ] totals_args = [call.args for call in prisma.db.query_raw.await_args_list if "judged_count" in call.args[0]] assert totals_args == [(totals_args[0][0], ["leg-1", "leg-2"])] error_where = prisma.db.litellm_shadowevalattempt.find_first.call_args.kwargs["where"] @@ -1595,7 +1825,7 @@ async def test_list_shadow_eval_jobs_collapses_legs_into_jobs_newest_first(monke _leg_record(created_at=datetime(2026, 8, 13, tzinfo=timezone.utc)), _leg_record( id="leg-2", - api_key_id="key-hash-2", + target_id="key-hash-2", stopped_at=stamp, created_at=datetime(2026, 8, 13, tzinfo=timezone.utc), ), @@ -1615,14 +1845,14 @@ async def test_list_shadow_eval_jobs_collapses_legs_into_jobs_newest_first(monke ) monkeypatch.setattr(proxy_server, "prisma_client", prisma) - jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) + jobs = await list_shadow_eval_jobs(VIEWER, target_type=None, target_id=None, limit=50) assert [(job.job_id, job.status) for job in jobs] == [ ("job-1", "running"), ("job-2", "stopped"), ("job-3", "completed"), ] - assert [key.api_key_id for key in jobs[0].keys] == ["key-hash", "key-hash-2"] + assert [target.target_id for target in jobs[0].targets] == ["key-hash", "key-hash-2"] assert all(job.judged_count is None and job.results is None for job in jobs) legs_sql, legs_limit = prisma.db.query_raw.await_args_list[0].args assert "GROUP BY group_id ORDER BY MAX(created_at) DESC LIMIT $1::int" in legs_sql @@ -1648,17 +1878,20 @@ async def test_list_shadow_eval_jobs_filters_to_jobs_containing_the_key(monkeypa prisma = _shadow_prisma( legs=[ _leg_record(), - _leg_record(id="leg-2", api_key_id="key-hash-2"), - _leg_record(id="leg-3", group_id="job-2", api_key_id="key-hash-2"), + _leg_record(id="leg-2", target_id="key-hash-2"), + _leg_record(id="leg-3", group_id="job-2", target_id="key-hash-2"), _leg_record(id="leg-4", group_id="job-3"), ] ) monkeypatch.setattr(proxy_server, "prisma_client", prisma) - jobs = await list_shadow_eval_jobs(VIEWER, api_key_id="key-hash-2", limit=50) + jobs = await list_shadow_eval_jobs(VIEWER, target_type="key", target_id="key-hash-2", limit=50) assert [job.job_id for job in jobs] == ["job-1", "job-2"] - assert [key.api_key_id for key in jobs[0].keys] == ["key-hash", "key-hash-2"] + assert [target.target_id for target in jobs[0].targets] == ["key-hash", "key-hash-2"] + legs_sql, *legs_params = prisma.db.query_raw.await_args_list[0].args + assert "WHERE target_type = $2 AND target_id = $3" in legs_sql + assert legs_params == [50, "key", "key-hash-2"] @pytest.mark.parametrize( @@ -1682,7 +1915,7 @@ async def test_job_status_runs_until_every_key_stops_and_completed_outranks_stop legs=[ _leg_record( id=f"leg-{index}", - api_key_id=f"key-{index}", + target_id=f"key-{index}", stopped_at=stamp if stopped else None, ends_at=datetime.now(timezone.utc) + timedelta(days=days_left), ) @@ -1691,7 +1924,7 @@ async def test_job_status_runs_until_every_key_stops_and_completed_outranks_stop ) monkeypatch.setattr(proxy_server, "prisma_client", prisma) - jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) + jobs = await list_shadow_eval_jobs(VIEWER, target_type=None, target_id=None, limit=50) assert [job.status for job in jobs] == [expected] @@ -1707,9 +1940,9 @@ async def test_list_reads_completed_once_every_key_spends_its_budget(monkeypatch prisma = _shadow_prisma( legs=[ _leg_record(max_turns=5), - _leg_record(id="leg-2", api_key_id="key-hash-2", max_turns=5), - _leg_record(id="leg-3", group_id="job-2", api_key_id="key-hash", max_turns=5), - _leg_record(id="leg-4", group_id="job-2", api_key_id="key-hash-2", max_turns=5), + _leg_record(id="leg-2", target_id="key-hash-2", max_turns=5), + _leg_record(id="leg-3", group_id="job-2", target_id="key-hash", max_turns=5), + _leg_record(id="leg-4", group_id="job-2", target_id="key-hash-2", max_turns=5), ] ) prisma.attempt_rows = [ @@ -1720,13 +1953,13 @@ async def test_list_reads_completed_once_every_key_spends_its_budget(monkeypatch ] monkeypatch.setattr(proxy_server, "prisma_client", prisma) - jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) + jobs = await list_shadow_eval_jobs(VIEWER, target_type=None, target_id=None, limit=50) by_id = {job.job_id: job for job in jobs} assert by_id["job-1"].status == "completed" - assert all(key.stopped_at is None for key in by_id["job-1"].keys) + assert all(target.stopped_at is None for target in by_id["job-1"].targets) assert by_id["job-2"].status == "running" - assert {key.api_key_id: key.attempt_count for key in by_id["job-2"].keys} == {"key-hash": 5, "key-hash-2": 3} + assert {t.target_id: t.attempt_count for t in by_id["job-2"].targets} == {"key-hash": 5, "key-hash-2": 3} @pytest.mark.asyncio @@ -1740,7 +1973,7 @@ async def test_recorded_operator_stop_outranks_budget_arithmetic(monkeypatch: py prisma.attempt_rows = [{"job_id": "leg-1", "attempt_count": 6, "spend": 0.0}] monkeypatch.setattr(proxy_server, "prisma_client", prisma) - jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) + jobs = await list_shadow_eval_jobs(VIEWER, target_type=None, target_id=None, limit=50) assert jobs[0].status == "stopped" assert jobs[0].stopped_by == "admin" @@ -1760,7 +1993,7 @@ async def test_backfilled_legacy_stop_never_reads_as_completion(monkeypatch: pyt prisma.attempt_rows = [{"job_id": "leg-1", "attempt_count": 6, "spend": 0.0}] monkeypatch.setattr(proxy_server, "prisma_client", prisma) - jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) + jobs = await list_shadow_eval_jobs(VIEWER, target_type=None, target_id=None, limit=50) assert jobs[0].status == "stopped" @@ -1808,6 +2041,56 @@ def test_max_budget_migration_is_additive_and_leaves_legacy_rows_null(): @pytest.mark.asyncio +@pytest.mark.asyncio +async def test_verdicts_keep_same_id_targets_of_different_kinds_distinct(monkeypatch): + """A team and a user can legitimately share an id; their slices must not merge.""" + from litellm.proxy import proxy_server + + leg_rows = [ + { + "grp": "leg-1", + "turn_count": 6, + "real_wins": 2, + "shadow_wins": 4, + "ties": 0, + "avg_confidence": 0.8, + "real_spend": 0.02, + "shadow_spend": 0.01, + "cache_hit_turns": 0, + }, + { + "grp": "leg-2", + "turn_count": 4, + "real_wins": 3, + "shadow_wins": 0, + "ties": 1, + "avg_confidence": 0.6, + "real_spend": 0.05, + "shadow_spend": 0.04, + "cache_hit_turns": 1, + }, + ] + prisma = _shadow_prisma( + legs=[ + _leg_record(target_type="team", target_id="dev-alice"), + _leg_record(id="leg-2", target_type="user", target_id="dev-alice"), + ], + agg_rows=leg_rows[:1], + by_leg_rows=leg_rows, + known_teams={"dev-alice": "alias"}, + known_users={"dev-alice": "alice@example.com"}, + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + response = await get_shadow_eval_job("job-1", VIEWER) + + verdicts_by_target = {(t.target_type, t.target_id): t.verdicts for t in response.targets} + assert verdicts_by_target[("team", "dev-alice")].turn_count == 6 + assert verdicts_by_target[("team", "dev-alice")].shadow_win_rate_pct == 66.7 + assert verdicts_by_target[("user", "dev-alice")].turn_count == 4 + assert verdicts_by_target[("user", "dev-alice")].shadow_win_rate_pct == 0.0 + + async def test_stop_rejects_a_job_that_already_spent_its_budget(monkeypatch: pytest.MonkeyPatch): import litellm.proxy.proxy_server as proxy_server @@ -1832,9 +2115,9 @@ async def test_list_reads_completed_once_every_key_spends_its_dollar_budget(monk prisma = _shadow_prisma( legs=[ _leg_record(max_turns=SHADOW_EVAL_TURN_VALVE, max_budget=1.0), - _leg_record(id="leg-2", api_key_id="key-hash-2", max_turns=SHADOW_EVAL_TURN_VALVE, max_budget=1.0), + _leg_record(id="leg-2", target_id="key-hash-2", max_turns=SHADOW_EVAL_TURN_VALVE, max_budget=1.0), _leg_record( - id="leg-3", group_id="job-2", api_key_id="key-hash", max_turns=SHADOW_EVAL_TURN_VALVE, max_budget=1.0 + id="leg-3", group_id="job-2", target_id="key-hash", max_turns=SHADOW_EVAL_TURN_VALVE, max_budget=1.0 ), ] ) @@ -1845,13 +2128,13 @@ async def test_list_reads_completed_once_every_key_spends_its_dollar_budget(monk ] monkeypatch.setattr(proxy_server, "prisma_client", prisma) - jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) + jobs = await list_shadow_eval_jobs(VIEWER, target_type=None, target_id=None, limit=50) by_id = {job.job_id: job for job in jobs} assert by_id["job-1"].status == "completed" assert by_id["job-2"].status == "running" - assert {key.api_key_id: key.spend for key in by_id["job-1"].keys} == {"key-hash": 1.0, "key-hash-2": 1.25} - assert all(key.max_budget == 1.0 for key in by_id["job-1"].keys) + assert {t.target_id: t.spend for t in by_id["job-1"].targets} == {"key-hash": 1.0, "key-hash-2": 1.25} + assert all(target.max_budget == 1.0 for target in by_id["job-1"].targets) @pytest.mark.asyncio @@ -1880,11 +2163,11 @@ async def test_legacy_jobs_without_a_dollar_budget_stay_turn_gated(monkeypatch: prisma.attempt_rows = [{"job_id": "leg-1", "attempt_count": 40, "spend": 250.0}] monkeypatch.setattr(proxy_server, "prisma_client", prisma) - jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) + jobs = await list_shadow_eval_jobs(VIEWER, target_type=None, target_id=None, limit=50) assert jobs[0].status == "running" - assert jobs[0].keys[0].max_budget is None - assert jobs[0].keys[0].spend == 250.0 + assert jobs[0].targets[0].max_budget is None + assert jobs[0].targets[0].spend == 250.0 @pytest.mark.asyncio @@ -1892,14 +2175,14 @@ async def test_shadow_eval_responses_name_every_shadowed_key(monkeypatch: pytest import litellm.proxy.proxy_server as proxy_server prisma = _shadow_prisma( - legs=[_leg_record(), _leg_record(id="leg-2", api_key_id="deleted-key-hash")], + legs=[_leg_record(), _leg_record(id="leg-2", target_id="deleted-key-hash")], known_keys=("key-hash", "key-hash-2"), ) monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) monkeypatch.setattr(proxy_server, "prisma_client", prisma) - jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) - assert [(key.key_alias, key.key_name) for key in jobs[0].keys] == [ + jobs = await list_shadow_eval_jobs(VIEWER, target_type=None, target_id=None, limit=50) + assert [(target.target_alias, target.key_name) for target in jobs[0].targets] == [ (None, None), ("prod-alpha", "sk-...lpha"), ] @@ -1907,7 +2190,7 @@ async def test_shadow_eval_responses_name_every_shadowed_key(monkeypatch: pytest assert batched_where == {"token": {"in": ["deleted-key-hash", "key-hash"]}} detail = await get_shadow_eval_job("job-1", VIEWER) - assert [key.key_alias for key in detail.keys] == [None, "prod-alpha"] + assert [target.target_alias for target in detail.targets] == [None, "prod-alpha"] @pytest.mark.asyncio @@ -1919,7 +2202,7 @@ async def test_stop_shadow_eval_stops_every_unstopped_leg_and_rejects_non_runnin import litellm.proxy.proxy_server as proxy_server earned = datetime.now(timezone.utc) - timedelta(hours=1) - prisma = _shadow_prisma(legs=[_leg_record(), _leg_record(id="leg-2", api_key_id="key-hash-2", stopped_at=earned)]) + prisma = _shadow_prisma(legs=[_leg_record(), _leg_record(id="leg-2", target_id="key-hash-2", stopped_at=earned)]) monkeypatch.setattr(proxy_server, "prisma_client", prisma) stopped = await stop_shadow_eval_job("job-1", ADMIN) @@ -1938,9 +2221,9 @@ async def test_stop_shadow_eval_stops_every_unstopped_leg_and_rejects_non_runnin assert datetime.fromisoformat(stop_stamp).tzinfo is None assert prisma.db.execute_raw.await_count == 1 prisma.db.litellm_shadowevaljob.update_many.assert_not_called() - by_key = {key.api_key_id: key.stopped_at for key in stopped.keys} - assert by_key["key-hash-2"] == earned - assert by_key["key-hash"] is not None and by_key["key-hash"] != earned + by_target = {target.target_id: target.stopped_at for target in stopped.targets} + assert by_target["key-hash-2"] == earned + assert by_target["key-hash"] is not None and by_target["key-hash"] != earned done_leg = _leg_record(ends_at=datetime.now(timezone.utc) - timedelta(days=1)) prisma_done = _shadow_prisma(legs=[done_leg]) @@ -2331,7 +2614,7 @@ async def test_get_shadow_eval_job_sums_funnel_rows_across_legs(monkeypatch: pyt }, ] prisma = _shadow_prisma( - legs=[_leg_record(), _leg_record(id="leg-2", api_key_id="key-hash-2")], + legs=[_leg_record(), _leg_record(id="leg-2", target_id="key-hash-2")], agg_rows=tier_rows, ) prisma.funnel_rows = [{"legs_with_rows": 2, "not_sampled": 30, "unjudgeable": 5, "shed": 2, "withheld": 3}] @@ -2366,7 +2649,7 @@ async def test_partially_seeded_funnel_reads_as_unknown_coverage(monkeypatch: py }, ] prisma = _shadow_prisma( - legs=[_leg_record(), _leg_record(id="leg-2", api_key_id="key-hash-2")], + legs=[_leg_record(), _leg_record(id="leg-2", target_id="key-hash-2")], agg_rows=tier_rows, ) prisma.funnel_rows = [{"legs_with_rows": 1, "not_sampled": 30, "unjudgeable": 5, "shed": 2, "withheld": 0}] diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx index ef6e224761d..8b397f20552 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx @@ -39,6 +39,35 @@ vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ })), })); +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + useInfiniteTeams: vi.fn(() => ({ + data: { pages: [{ teams: [{ team_id: "team-eng", team_alias: "engineering" }], page: 1, total_pages: 1 }] }, + isLoading: false, + fetchNextPage: vi.fn(), + hasNextPage: false, + isFetchingNextPage: false, + })), +})); + +vi.mock("@/app/(dashboard)/hooks/users/useUsers", () => ({ + useInfiniteUsers: vi.fn(() => ({ + data: { + pages: [ + { + users: [{ user_id: "dev-alice", user_alias: null, user_email: "alice@example.com" }], + page: 1, + total_pages: 1, + }, + ], + }, + isPending: false, + isError: false, + fetchNextPage: vi.fn(), + hasNextPage: false, + isFetchingNextPage: false, + })), +})); + vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ useAutoRouters: vi.fn(() => ({ data: [ @@ -60,7 +89,7 @@ vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({ })), })); -import ShadowEvalSection, { shadowedKeyLabel } from "./ShadowEvalSection"; +import ShadowEvalSection, { shadowedTargetLabel } from "./ShadowEvalSection"; import { useShadowEvalJob, useShadowEvalJobs, @@ -77,14 +106,15 @@ const job = (overrides: Partial = {}): ShadowEvalJob => ({ baseline_model: null, judge_model: "anthropic/claude-sonnet-5", shadow_percentage: 10, - keys: [ + targets: [ { - api_key_id: "hashed-key-abc", + target_type: "key", + target_id: "hashed-key-abc", max_turns: 10000, max_budget: 10, spend: 3.21, stopped_at: null, - key_alias: "prod-alpha", + target_alias: "prod-alpha", key_name: "sk-...alpha", }, ], @@ -129,7 +159,6 @@ const job = (overrides: Partial = {}): ShadowEvalJob => ({ cache_hit_turns: 2, }, ], - by_key: [], overall_shadow_win_rate_pct: 48.0, overall_tie_rate_pct: 22.0, sampled_real_spend: 0.6, @@ -144,17 +173,18 @@ const job = (overrides: Partial = {}): ShadowEvalJob => ({ ...overrides, }); -const keyEntry = ( - api_key_id: string, - overrides: Partial = {}, -): ShadowEvalJob["keys"][number] => ({ - api_key_id, +const targetEntry = ( + target_id: string, + overrides: Partial = {}, +): ShadowEvalJob["targets"][number] => ({ + target_type: "key", + target_id, max_turns: 10000, max_budget: 10, spend: 0, stopped_at: null, attempt_count: null, - key_alias: null, + target_alias: null, key_name: null, ...overrides, }); @@ -235,8 +265,8 @@ describe("ShadowEvalSection", () => { it("gives every active job its own card with a stop button, with the form still offered", () => { mockHooks({ jobs: [ - job({ job_id: "job-a", status: "running", keys: [keyEntry("key-a")] }), - job({ job_id: "job-b", status: "running", keys: [keyEntry("key-b")] }), + job({ job_id: "job-a", status: "running", targets: [targetEntry("key-a")] }), + job({ job_id: "job-b", status: "running", targets: [targetEntry("key-b")] }), ], }); render(); @@ -347,7 +377,7 @@ describe("ShadowEvalSection", () => { }); it("shows spend without a budget cap for a job from before spend budgets existed", () => { - const j = job({ keys: [keyEntry("hashed-key-abc", { max_budget: null, spend: 3.21 })] }); + const j = job({ targets: [targetEntry("hashed-key-abc", { max_budget: null, spend: 3.21 })] }); mockHooks({ jobs: [j], detailsById: { "job-1": j } }); render(); expect(screen.getByText(/\$3\.21 eval spend/)).toBeInTheDocument(); @@ -417,6 +447,38 @@ describe("ShadowEvalSection", () => { const expectedBody = { api_key_ids: ["hash-alpha", "hash-beta"], + team_ids: [], + user_ids: [], + router_name: "gpt-auto", + direction: "forward", + shadow_percentage: 10, + duration_days: 7, + max_budget: 10, + judge_model: "anthropic/claude-sonnet-5", + }; + expect(start.mutate).toHaveBeenCalledWith(expectedBody); + }); + + it("submits a team-only job with team_ids and no keys", async () => { + const user = userEvent.setup(); + const { start } = mockHooks({}); + render(); + + expect(screen.getByText("Start shadow eval")).toBeDisabled(); + + await user.click(screen.getByPlaceholderText("Search teams by alias")); + const teamList = await screen.findByTestId("paginated-multi-select-list"); + await user.click(within(teamList).getByText("engineering")); + await user.click(screen.getByPlaceholderText("Select an auto-router")); + await user.click(await screen.findByText("gpt-auto")); + await user.click(screen.getByPlaceholderText("Select a judge model")); + await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + await user.click(screen.getByText("Start shadow eval")); + + const expectedBody = { + api_key_ids: [], + team_ids: ["team-eng"], + user_ids: [], router_name: "gpt-auto", direction: "forward", shadow_percentage: 10, @@ -453,6 +515,8 @@ describe("ShadowEvalSection", () => { const expectedBody = { api_key_ids: ["hash-alpha"], + team_ids: [], + user_ids: [], router_name: "gpt-auto", direction: "reverse", baseline_model: "prod-claude", @@ -485,9 +549,13 @@ describe("ShadowEvalSection", () => { }); it("labels the shadowed key by alias, then masked name, then truncated hash", () => { - expect(shadowedKeyLabel(job().keys[0])).toBe("prod-alpha"); - expect(shadowedKeyLabel(keyEntry("hashed-key-abc", { key_name: "sk-...alpha" }))).toBe("sk-...alpha"); - expect(shadowedKeyLabel(keyEntry("hashed-key-abc"))).toBe("hashed-key…"); + expect(shadowedTargetLabel(job().targets[0])).toBe("prod-alpha"); + expect(shadowedTargetLabel(targetEntry("hashed-key-abc", { key_name: "sk-...alpha" }))).toBe("sk-...alpha"); + expect(shadowedTargetLabel(targetEntry("hashed-key-abc"))).toBe("hashed-key…"); + expect(shadowedTargetLabel(targetEntry("team-eng", { target_type: "team" }))).toBe("team-eng"); + expect(shadowedTargetLabel(targetEntry("team-eng", { target_type: "team", target_alias: "engineering" }))).toBe( + "engineering", + ); }); it("breaks results down per key, so one key exhausting its own budget is visible while a sibling runs on", () => { @@ -495,15 +563,12 @@ describe("ShadowEvalSection", () => { jobs: [ job({ judged_count: 205, - keys: [ - keyEntry("hash-spent", { max_budget: 2, spend: 1.5, stopped_at: "2026-08-08T00:00:00Z" }), - keyEntry("hash-hungry", { max_budget: 5, spend: 0.2 }), - ], - results: { - by_tier: [], - by_current_model: [], - by_key: [ - { + targets: [ + targetEntry("hash-spent", { + max_budget: 2, + spend: 1.5, + stopped_at: "2026-08-08T00:00:00Z", + verdicts: { group: "hash-spent", turn_count: 200, real_win_rate_pct: 20.0, @@ -514,7 +579,12 @@ describe("ShadowEvalSection", () => { shadow_spend: 0.5, cache_hit_turns: 0, }, - ], + }), + targetEntry("hash-hungry", { max_budget: 5, spend: 0.2 }), + ], + results: { + by_tier: [], + by_current_model: [], overall_shadow_win_rate_pct: 60.0, overall_tie_rate_pct: 20.0, sampled_real_spend: 0.9, @@ -539,7 +609,7 @@ describe("ShadowEvalSection", () => { expect(screen.getByText(/205 turns judged/)).toBeInTheDocument(); expect(screen.getByText(/Shadowing 10% of/)).toBeInTheDocument(); - expect(screen.getByText("2 keys")).toBeInTheDocument(); + expect(screen.getByText("2 targets")).toBeInTheDocument(); }); it("reads a key that spent its budget as completed even before the sweep stamps it", () => { @@ -547,9 +617,9 @@ describe("ShadowEvalSection", () => { mockHooks({ jobs: [ job({ - keys: [ - keyEntry("hash-spent", { max_budget: 2, spend: 2, attempt_count: 40 }), - keyEntry("hash-hungry", legacyTurnBudgetLeg), + targets: [ + targetEntry("hash-spent", { max_budget: 2, spend: 2, attempt_count: 40 }), + targetEntry("hash-hungry", legacyTurnBudgetLeg), ], }), ], @@ -571,9 +641,9 @@ describe("ShadowEvalSection", () => { job({ judged_count: 0, results: null, - keys: [ - keyEntry("hash-spent", { max_budget: 0.5, spend: 0.5, attempt_count: 2 }), - keyEntry("hash-hungry", { max_budget: 5, spend: 0.01, attempt_count: 1 }), + targets: [ + targetEntry("hash-spent", { max_budget: 0.5, spend: 0.5, attempt_count: 2 }), + targetEntry("hash-hungry", { max_budget: 5, spend: 0.01, attempt_count: 1 }), ], }), ], @@ -594,9 +664,9 @@ describe("ShadowEvalSection", () => { jobs: [ job({ status: "completed", - keys: [ - keyEntry("hash-spent", { max_turns: 200, stopped_at: "2026-08-08T00:00:00Z" }), - keyEntry("hash-hungry", { max_turns: 500 }), + targets: [ + targetEntry("hash-spent", { max_turns: 200, stopped_at: "2026-08-08T00:00:00Z" }), + targetEntry("hash-hungry", { max_turns: 500 }), ], }), ], diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx index 39dee28390a..ec145bc617a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx @@ -3,10 +3,13 @@ import React, { useMemo, useState } from "react"; import { useInfiniteKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; import { useAutoRouters, usePlainModelGroups } from "@/app/(dashboard)/hooks/models/useModels"; import { PaginatedMultiSelect } from "@/components/shared/PaginatedMultiSelect"; +import TeamMultiSelect from "@/components/common_components/team_multi_select"; +import { userOptionLabel } from "@/components/common_components/UserDropdown"; import { SearchSelect, type SearchSelectOption } from "@/components/shared/SearchSelect"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; @@ -27,7 +30,7 @@ import { useStartShadowEval, useStopShadowEval, type ShadowEvalJob, - type ShadowEvalJobKey, + type ShadowEvalJobTarget, type ShadowEvalSlice, } from "./useShadowEval"; @@ -66,29 +69,31 @@ const routerMatchedOrBeatPct = ( ? 100 - results.overall_shadow_win_rate_pct : results.overall_shadow_win_rate_pct + results.overall_tie_rate_pct; -export const shadowedKeyLabel = (key: ShadowEvalJobKey): string => - key.key_alias || key.key_name || `${key.api_key_id.slice(0, 10)}…`; +export const shadowedTargetLabel = (target: ShadowEvalJobTarget): string => + target.target_alias || + target.key_name || + (target.target_type === "key" ? `${target.target_id.slice(0, 10)}…` : target.target_id); -const shadowedKeysLabel = (job: ShadowEvalJob): string => - job.keys.length === 1 ? shadowedKeyLabel(job.keys[0]) : `${job.keys.length} keys`; +const shadowedTargetsLabel = (job: ShadowEvalJob): string => + job.targets.length === 1 ? shadowedTargetLabel(job.targets[0]) : `${job.targets.length} targets`; const totalBudget = (job: ShadowEvalJob): number | null => - job.keys.reduce( - (sum, key) => (sum === null || key.max_budget == null ? null : sum + key.max_budget), + job.targets.reduce( + (sum, target) => (sum === null || target.max_budget == null ? null : sum + target.max_budget), 0, ); -const totalSpend = (job: ShadowEvalJob): number => job.keys.reduce((sum, key) => sum + (key.spend ?? 0), 0); +const totalSpend = (job: ShadowEvalJob): number => job.targets.reduce((sum, target) => sum + (target.spend ?? 0), 0); -const keySpent = (key: ShadowEvalJobKey): boolean => { - const spendBudgetReached = key.max_budget != null && key.spend != null && key.spend >= key.max_budget; - const turnValveReached = key.attempt_count != null && key.attempt_count >= key.max_turns; +const targetSpent = (target: ShadowEvalJobTarget): boolean => { + const spendBudgetReached = target.max_budget != null && target.spend != null && target.spend >= target.max_budget; + const turnValveReached = target.attempt_count != null && target.attempt_count >= target.max_turns; return spendBudgetReached || turnValveReached; }; -const keyStatus = (job: ShadowEvalJob, key: ShadowEvalJobKey): string => { - if (job.status === "completed" || (key.stopped_at == null && keySpent(key))) return "completed"; - return key.stopped_at != null ? "stopped" : "running"; +const targetStatus = (job: ShadowEvalJob, target: ShadowEvalJobTarget): string => { + if (job.status === "completed" || (target.stopped_at == null && targetSpent(target))) return "completed"; + return target.stopped_at != null ? "stopped" : "running"; }; const jobHeadline = (job: ShadowEvalJob): React.ReactNode => @@ -96,12 +101,12 @@ const jobHeadline = (job: ShadowEvalJob): React.ReactNode => <> Comparing {job.router_name} to{" "} {job.baseline_model} on {job.shadow_percentage}% of{" "} - {shadowedKeysLabel(job)} traffic + {shadowedTargetsLabel(job)} traffic ) : ( <> - Shadowing {job.shadow_percentage}% of {shadowedKeysLabel(job)} traffic - via {job.router_name} + Shadowing {job.shadow_percentage}% of {shadowedTargetsLabel(job)}{" "} + traffic via {job.router_name} ); @@ -255,13 +260,12 @@ const VerdictBar: React.FC<{ direction: ShadowEvalDirection; results: NonNullabl ); }; -const KeyTable: React.FC<{ job: ShadowEvalJob }> = ({ job }) => { - const slices = new Map((job.results?.by_key ?? []).map((slice) => [slice.group, slice])); +const TargetTable: React.FC<{ job: ShadowEvalJob }> = ({ job }) => { return ( - Key + Target Status {["Budget used", "Router wins", `${otherArmLabel(job.direction)} wins`].map((label) => ( @@ -271,18 +275,23 @@ const KeyTable: React.FC<{ job: ShadowEvalJob }> = ({ job }) => { - {job.keys.map((key) => { - const slice = slices.get(key.api_key_id); + {job.targets.map((target) => { + const slice = target.verdicts; return ( - - {shadowedKeyLabel(key)} + + + {shadowedTargetLabel(target)} + {target.target_type !== "key" && ( + {target.target_type} + )} + - + - {key.max_budget != null - ? `${usd(key.spend ?? 0)} / ${usd(key.max_budget)}` - : `${(key.attempt_count ?? slice?.turn_count ?? 0).toLocaleString()} / ${key.max_turns.toLocaleString()} turns`} + {target.max_budget != null + ? `${usd(target.spend ?? 0)} / ${usd(target.max_budget)}` + : `${(target.attempt_count ?? slice?.turn_count ?? 0).toLocaleString()} / ${target.max_turns.toLocaleString()} turns`} {slice ? ( <> @@ -318,9 +327,9 @@ const ResultsBody: React.FC<{ job: ShadowEvalJob; resultsError?: boolean }> = ({ const hasVerdicts = results != null && (results.by_tier.length > 0 || results.by_current_model.length > 0); return ( <> - {job.keys.length > 1 && ( + {job.targets.length > 1 && (
- +
)} {/* results == null re-stated for TS narrowing; hasVerdicts alone cannot narrow it */} @@ -454,9 +463,9 @@ const DIRECTION_OPTIONS: readonly { value: ShadowEvalDirection; label: string }[ const START_FORM_DESCRIPTION: Record = { forward: - "Duplicates a sampled slice of the selected keys' traffic through the auto-router and has an LLM judge compare both answers blind. Each key gets its own spend budget. The router's answers are never served to users; judge calls bill to the shadowed key.", + "Duplicates a sampled slice of the selected targets' traffic (keys, teams, or users) through the auto-router and has an LLM judge compare both answers blind. Each target gets its own spend budget. The router's answers are never served to users; judge calls bill to the sampled traffic's own identity.", reverse: - "Duplicates a sampled slice of the traffic the auto-router already serves against a fixed baseline model and has an LLM judge compare both answers blind. Each key gets its own spend budget. The baseline's answers are never served to users; judge calls bill to the shadowed key.", + "Duplicates a sampled slice of the traffic the auto-router already serves against a fixed baseline model and has an LLM judge compare both answers blind. Each target gets its own spend budget. The baseline's answers are never served to users; judge calls bill to the sampled traffic's own identity.", }; const DURATION_OPTIONS = [ @@ -515,9 +524,46 @@ const KeySelect: React.FC<{ value: string[]; onChange: (tokens: string[]) => voi ); }; +const UserSelect: React.FC<{ value: string[]; onChange: (ids: string[]) => void }> = ({ value, onChange }) => { + const [search, setSearch] = useState(""); + const { data, isPending, isError, fetchNextPage, hasNextPage, isFetchingNextPage } = useInfiniteUsers( + 50, + search || undefined, + ); + const options = useMemo( + () => + Array.from( + new Map( + (data?.pages ?? []) + .flatMap((page) => page.users) + .map((user) => [user.user_id, { label: userOptionLabel(user), value: user.user_id }] as const), + ).values(), + ), + [data], + ); + return ( + void fetchNextPage()} + hasNextPage={hasNextPage} + isFetchingNextPage={isFetchingNextPage} + isLoading={isPending} + placeholder="Search users by email" + emptyText="No matching users" + errorText={isError ? "Users could not be loaded. Refresh the page to retry." : undefined} + /> + ); +}; + const StartForm: React.FC = () => { const { accessToken } = useAuthorized(); const [apiKeyIds, setApiKeyIds] = useState([]); + const [teamIds, setTeamIds] = useState([]); + const [userIds, setUserIds] = useState([]); const [routerName, setRouterName] = useState(""); const [direction, setDirection] = useState("forward"); const [baselineModel, setBaselineModel] = useState(""); @@ -542,12 +588,15 @@ const StartForm: React.FC = () => { const parsedMaxBudget = Number.parseFloat(maxBudget); const maxBudgetValid = parsedMaxBudget >= 0.01 && parsedMaxBudget <= 10000; const baselinePicked = direction === "forward" || baselineModel !== ""; - const filled = apiKeyIds.length > 0 && [routerName, judgeModel].every((field) => field !== "") && baselinePicked; + const targetsPicked = apiKeyIds.length + teamIds.length + userIds.length > 0; + const filled = targetsPicked && [routerName, judgeModel].every((field) => field !== "") && baselinePicked; const boundsValid = percentageValid && maxBudgetValid; const valid = Boolean(accessToken) && filled && boundsValid; const handleStart = () => { const startBody = { api_key_ids: apiKeyIds, + team_ids: teamIds, + user_ids: userIds, router_name: routerName, direction, ...(direction === "reverse" ? { baseline_model: baselineModel } : {}), @@ -587,6 +636,12 @@ const StartForm: React.FC = () => { + + + + + + { value={maxBudget} onChange={(e) => setMaxBudget(e.target.value)} /> - max shadow + judge spend, per key + max shadow + judge spend, per target {maxBudget.trim() !== "" && !maxBudgetValid && (

Enter a value from 0.01 to 10000

@@ -774,8 +829,9 @@ const ShadowEvalSection: React.FC = () => {

Shadow eval

- Blind-judge the auto-router on your real traffic: against the models a key uses today before switching, or - against a fixed baseline after it has switched. + Blind-judge the auto-router on the real traffic of a key, team, or user (teams and users cover + JWT-authenticated traffic): against the models they use today before switching, or against a fixed baseline + after they have switched.

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useShadowEval.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useShadowEval.ts index eef98320e67..107df0f594a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useShadowEval.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useShadowEval.ts @@ -7,7 +7,7 @@ import { $api, fetchClient } from "@/lib/http/api"; import type { components } from "@/lib/http/schema"; export type ShadowEvalJob = components["schemas"]["ShadowEvalJobResponse"]; -export type ShadowEvalJobKey = components["schemas"]["ShadowEvalJobKeyResponse"]; +export type ShadowEvalJobTarget = components["schemas"]["ShadowEvalJobTargetResponse"]; export type ShadowEvalSlice = components["schemas"]["ShadowEvalSlice"]; export type StartShadowEvalRequest = components["schemas"]["StartShadowEvalRequest"]; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b79e2c12447..c20545f6fb2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -1234,8 +1234,8 @@ export interface paths { }; /** * List Shadow Eval Jobs - * @description List shadow eval jobs, newest first, each key with its attempt count so status is - * accurate. Judged counts, spend, and results ride the detail endpoint only. + * @description List shadow eval jobs, newest first, each target with its attempt count so status + * is accurate. Judged counts, spend, and results ride the detail endpoint only. */ get: operations["list_shadow_eval_jobs_auto_router_shadow_eval_get"]; put?: never; @@ -1257,22 +1257,29 @@ export interface paths { put?: never; /** * Start Shadow Eval - * @description Start a shadow eval: duplicate a sampled slice of one or more keys' live traffic against - * a second arm, judge the two responses blind, and stratify win rates by tier, by the model - * that served the real arm, and by key. + * @description Start a shadow eval: duplicate a sampled slice of one or more targets' live traffic + * against a second arm, judge the two responses blind, and stratify win rates by tier, + * by the model that served the real arm, and by target. * - * A forward job answers whether the keys should adopt router_name: it samples the requests - * the router did not serve and duplicates them through it. A reverse job answers whether a - * key already on the router still gains from it: it samples the requests the router did - * serve and duplicates them against baseline_model. A key can hold one active job per - * direction, so both questions can run at once. + * A target is a virtual key, a team, or a user. Team and user targets match on the + * identity every request resolves to at auth time, so they cover JWT-authenticated + * traffic, which presents no virtual key; a user target samples that user's traffic + * across all their teams, whether it arrives on a JWT or a key they own. * - * Shadow responses are never served to users. Each key samples until its recorded eval - * spend, the shadow and judge calls' own cost, reaches max_budget dollars, the job's - * window ends, or the job is stopped, so one key running out of budget does not end - * sampling for the others; sampling changes propagate to pods within about 10 seconds. - * Shadow and judge calls bill to the shadowed key but are excluded from request counts - * and auto-router adoption metrics. + * A forward job answers whether the targets should adopt router_name: it samples the + * requests the router did not serve and duplicates them through it. A reverse job + * answers whether a target already on the router still gains from it: it samples the + * requests the router did serve and duplicates them against baseline_model. A target + * can hold one active job per direction, so both questions can run at once, and a + * request matching several jobs' targets (say its key and its team) is sampled by + * each, separately budgeted. + * + * Shadow responses are never served to users. Each target samples until its recorded + * eval spend, the shadow and judge calls' own cost, reaches max_budget dollars, the + * job's window ends, or the job is stopped, so one target running out of budget does + * not end sampling for the others; sampling changes propagate to pods within about 10 + * seconds. Shadow and judge calls bill to the sampled request's own identity but are + * excluded from request counts and auto-router adoption metrics. */ post: operations["start_shadow_eval_auto_router_shadow_eval_start_post"]; delete?: never; @@ -1312,8 +1319,8 @@ export interface paths { put?: never; /** * Stop Shadow Eval Job - * @description Stop an active shadow eval job, every key it scopes at once. Attempts are kept; - * sampling halts within ~10s. Keys that already stopped on their own budget keep the + * @description Stop an active shadow eval job, every target it scopes at once. Attempts are kept; + * sampling halts within ~10s. Targets that already stopped on their own budget keep the * stopped_at they earned. The statement is the whole state machine: it claims the job * only while a leg still samples inside the window with no stop recorded, so a racing * operator, a same-instant budget spend, and a repeat stop all read the same 400 with @@ -35262,56 +35269,10 @@ export interface components { /** Timeout */ timeout?: number | null; }; - /** - * ShadowEvalJobKeyResponse - * @description One key a job shadows, with its own budget and stop state. - */ - ShadowEvalJobKeyResponse: { - /** - * Api Key Id - * @description The hashed virtual key whose traffic this entry scopes - */ - api_key_id: string; - /** - * Attempt Count - * @description This key's sampled attempts so far, judged and errored alike, the same count the sampler budgets against max_turns; populated on list and detail responses. Frozen at stopped_at once the key is stamped, so in-flight attempts landing after a stop never reclassify it - */ - attempt_count?: number | null; - /** - * Key Alias - * @description Alias of the shadowed key, resolved from the key row at read time; None when unset or deleted - */ - key_alias?: string | null; - /** - * Key Name - * @description Masked display name (sk-...) of the shadowed key, resolved at read time like key_alias - */ - key_name?: string | null; - /** - * Max Budget - * @description This key's own USD budget for the eval's shadow and judge spend, independent of its siblings'; None on jobs created before spend budgets existed, which max_turns alone bounds - */ - max_budget?: number | null; - /** - * Max Turns - * @description This key's sample-count ceiling: the whole budget for jobs created before max_budget existed, and the error-loop safety valve otherwise - */ - max_turns: number; - /** - * Spend - * @description This key's recorded shadow plus judge spend in USD, the same figure the sampler budgets against max_budget; populated on list and detail responses and frozen at stopped_at exactly like attempt_count - */ - spend?: number | null; - /** - * Stopped At - * @description When this key's slot was stamped free, whether its own budget ran out, the window closed, or an operator stopped the job; status is derived, so a spent budget reads completed even while this is still unset - */ - stopped_at?: string | null; - }; /** * ShadowEvalJobResponse - * @description A shadow-eval job over one or more keys, each with its own budget and stop state; - * status is derived from stopped_by, the keys' stop and budget state, and ends_at, + * @description A shadow-eval job over one or more targets, each with its own budget and stop state; + * status is derived from stopped_by, the targets' stop and budget state, and ends_at, * never stored, so no writer anywhere can produce an inconsistent one. Aggregate * fields are populated by the detail endpoint only and stay None on list responses. */ @@ -35353,11 +35314,6 @@ export interface components { * @description Verdicts recorded; detail endpoint only */ judged_count?: number | null; - /** - * Keys - * @description The keys whose traffic this job evaluates, and only those keys', each with its own budget - */ - keys: components["schemas"]["ShadowEvalJobKeyResponse"][]; /** * Last Error * @description Most recent attempt error; detail endpoint only @@ -35373,8 +35329,8 @@ export interface components { * Status * @description Three recorded facts, no history-guessing: a stop is stopped_by (the migration * backfills it for every job that displayed stopped when the column arrived, so the - * pre-column population is closed), completion is the window passing or every key - * spending its budget, and anything else is running. The all-keys-stamped fallback + * pre-column population is closed), completion is the window passing or every target + * spending its budget, and anything else is running. The all-targets-stamped fallback * covers only stops written by pre-column pods during a rolling deploy. * @enum {string} */ @@ -35384,6 +35340,65 @@ export interface components { * @description The operator who stopped the job early, recorded by the stop endpoint; 'unknown' backfilled by migration for jobs that displayed stopped when the column arrived; None when the job ended on its own. Its presence is what makes a job read stopped rather than completed */ stopped_by?: string | null; + /** + * Targets + * @description The targets whose traffic this job evaluates, and only theirs, each with its own budget + */ + targets: components["schemas"]["ShadowEvalJobTargetResponse"][]; + }; + /** + * ShadowEvalJobTargetResponse + * @description One target a job shadows (a key, team, or user), with its own budget and stop state. + */ + ShadowEvalJobTargetResponse: { + /** + * Attempt Count + * @description This target's sampled attempts so far, judged and errored alike, the same count the sampler budgets against max_turns; populated on list and detail responses. Frozen at stopped_at once the target is stamped, so in-flight attempts landing after a stop never reclassify it + */ + attempt_count?: number | null; + /** + * Key Name + * @description Masked display name (sk-...) for key targets, resolved at read time; None for teams and users + */ + key_name?: string | null; + /** + * Max Budget + * @description This target's own USD budget for the eval's shadow and judge spend, independent of its siblings'; None on jobs created before spend budgets existed, which max_turns alone bounds + */ + max_budget?: number | null; + /** + * Max Turns + * @description This target's sample-count ceiling: the whole budget for jobs created before max_budget existed, and the error-loop safety valve otherwise + */ + max_turns: number; + /** + * Spend + * @description This target's recorded shadow plus judge spend in USD, the same figure the sampler budgets against max_budget; populated on list and detail responses and frozen at stopped_at exactly like attempt_count + */ + spend?: number | null; + /** + * Stopped At + * @description When this target's slot was stamped free, whether its own budget ran out, the window closed, or an operator stopped the job; status is derived, so a spent budget reads completed even while this is still unset + */ + stopped_at?: string | null; + /** + * Target Alias + * @description Display label resolved from the target's own row at read time: the key's alias, the team's alias, or the user's email; None when unset or deleted + */ + target_alias?: string | null; + /** + * Target Id + * @description The hashed virtual key, team id, or user id whose traffic this entry scopes + */ + target_id: string; + /** + * Target Type + * @description What kind of entity this entry scopes + * @enum {string} + */ + target_type: "key" | "team" | "user"; + /** @description This target's own judged-verdict slice; detail endpoint only, None until a turn is judged */ + verdicts?: components["schemas"]["ShadowEvalSlice"] | null; }; /** * ShadowEvalResult @@ -35395,11 +35410,6 @@ export interface components { * @description Sliced by the model that served the real arm: the keys' incumbent models in forward mode, and in reverse the models the router itself picked */ by_current_model: components["schemas"]["ShadowEvalSlice"][]; - /** - * By Key - * @description One slice per scoped key that has judged verdicts, grouped on the raw key hash. Keys the job scopes but has not judged a turn for yet are absent rather than reported as zero - */ - by_key: components["schemas"]["ShadowEvalSlice"][]; /** By Tier */ by_tier: components["schemas"]["ShadowEvalSlice"][]; /** @@ -35441,8 +35451,9 @@ export interface components { }; /** * ShadowEvalSlice - * @description Judge outcomes for one slice of a job's verdicts (a router tier, or one of the - * models that served the real arm). + * @description Judge outcomes for one slice of a job's verdicts: a router tier, one of the + * models that served the real arm, or one scoped target (embedded on that target's + * own entry, so slices never need re-joining to a target by id). */ ShadowEvalSlice: { /** Avg Judge Confidence */ @@ -35672,12 +35683,18 @@ export interface components { }; /** * StartShadowEvalRequest - * @description Start duplicating one or more keys' traffic for blind comparison against an auto-router. + * @description Start duplicating one or more targets' traffic for blind comparison against an auto-router. + * + * A target is a virtual key, a team, or a user; each becomes its own leg with its own + * budget and stop state. Team and user targets match on the identity every request + * carries after auth (user_api_key_team_id / user_api_key_user_id), so they cover + * JWT-authenticated traffic, which presents no virtual key at all. */ StartShadowEvalRequest: { /** * Api Key Ids - * @description The hashed virtual keys whose traffic will be shadowed. Shadow evaluation runs ONLY on these keys' traffic; requests made with any other key are not sampled. Each key carries its own max_budget spend budget, so one key exhausting its budget leaves the others sampling. At most 100 keys per job, which also bounds every read the job's endpoints make. + * @description Hashed virtual keys whose traffic will be shadowed. Combined with team_ids and user_ids the job needs at least one target and at most 100, which also bounds every read the job's endpoints make. Each target carries its own max_budget spend budget, so one exhausting its budget leaves the others sampling. + * @default [] */ api_key_ids: string[]; /** @@ -35706,7 +35723,7 @@ export interface components { judge_model: string; /** * Max Budget - * @description Per-key USD budget for the eval's own overhead, the shadow-arm and judge calls, priced with the same figures the spend pipeline bills. EACH scoped key samples until its recorded eval spend reaches this, so a job over N keys spends at most about N times max_budget; in-flight samples can overshoot the cap by one sampling cache window + * @description Per-target USD budget for the eval's own overhead, the shadow-arm and judge calls, priced with the same figures the spend pipeline bills. EACH scoped target samples until its recorded eval spend reaches this, so a job over N targets spends at most about N times max_budget; in-flight samples can overshoot the cap by one sampling cache window * @default 10 */ max_budget: number; @@ -35717,9 +35734,21 @@ export interface components { router_name: string; /** * Shadow Percentage - * @description Percentage of the key's requests to duplicate through the router + * @description Percentage of each target's requests to duplicate through the router */ shadow_percentage: number; + /** + * Team Ids + * @description Teams whose traffic will be shadowed, matched on the team every authenticated request resolves to, so a team's JWT-auth and virtual-key traffic are both sampled + * @default [] + */ + team_ids: string[]; + /** + * User Ids + * @description Users whose traffic will be shadowed, matched on the user every authenticated request resolves to across all their teams: JWT requests carrying their subject claim and virtual keys they own + * @default [] + */ + user_ids: string[]; }; /** * SuccessfulKeyUpdate @@ -40681,8 +40710,10 @@ export interface operations { list_shadow_eval_jobs_auto_router_shadow_eval_get: { parameters: { query?: { - /** @description Filter to jobs that shadow this key, alone or alongside others */ - api_key_id?: string | null; + /** @description Kind of target to filter on; requires target_id */ + target_type?: ("key" | "team" | "user") | null; + /** @description Filter to jobs that shadow this target, alone or alongside others */ + target_id?: string | null; /** @description Newest jobs to return */ limit?: number; }; From a99f62d1bc70cdc8c6a6ad0b8b358ed320ad919e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:46:12 -0700 Subject: [PATCH 119/120] fix(anthropic_messages): park deferred billing before the end-of-stream sentinel At end of drain the pump enqueued the sentinel first and picked the billing mode from client_detached afterward, so a client that consumed the sentinel and tore the relay down before the pump resumed (possible whenever the sentinel enqueue hit a full queue) had its fully delivered response billed through the teardown path, skipping the proxy's post-response hook. Bill or park before the sentinel goes out, and let an unconsumed sentinel fall back to dispatching the parked billing. --- .../messages/streaming_iterator.py | 23 +++- .../messages/test_streaming_iterator.py | 120 ++++++++++++++++++ 2 files changed, 137 insertions(+), 6 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 55a64dedee4..45c7825344b 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -598,7 +598,12 @@ class BaseAnthropicMessagesStreamingIterator: """Drain the whole upstream into ``queue`` (backpressured) and bill once. Runs detached so a client disconnect can't interrupt the upstream read; - see ``async_sse_wrapper`` for the full rationale. Returns after billing. + see ``async_sse_wrapper`` for the full rationale. On a completed drain + the success billing (or deferred park) happens before the end-of-stream + sentinel is enqueued: the relay can only tear down after consuming the + sentinel, so its teardown can never outrun the park and get mistaken + for a client disconnect, and a sentinel the client never consumes falls + back to dispatching the parked billing here. """ from litellm._logging import verbose_proxy_logger @@ -631,11 +636,17 @@ class BaseAnthropicMessagesStreamingIterator: await self._handle_pump_upstream_error(queue, client_detached, collected_chunks, exc) return - if not client_detached.is_set(): - if not saw_terminal_event: - await self._enqueue_for_client(queue, client_detached, _incomplete_stream_error_sse_event()) - await self._enqueue_for_client(queue, client_detached, None) - await self._bill_collected_chunks(collected_chunks, stream_teardown=client_detached.is_set()) + if client_detached.is_set(): + await self._bill_collected_chunks(collected_chunks, stream_teardown=True) + return + if not saw_terminal_event and not await self._enqueue_for_client( + queue, client_detached, _incomplete_stream_error_sse_event() + ): + await self._bill_collected_chunks(collected_chunks, stream_teardown=True) + return + await self._bill_collected_chunks(collected_chunks, stream_teardown=False) + if not await self._enqueue_for_client(queue, client_detached, None): + self._dispatch_pending_deferred_logging() async def _handle_pump_upstream_error( self, diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index e8099a4217b..11a048edc1f 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -1156,3 +1156,123 @@ async def test_normal_end_without_deferred_dispatch_enqueues_immediately(monkeyp assert len(worker.enqueued) == 1 assert getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None) is None worker.close_enqueued() + + +def _backpressured_wrapper(iterator, upstream_exhausted: asyncio.Event): + async def _stream(): + try: + for event in COMPLETE_STREAM_EVENTS: + yield event + finally: + upstream_exhausted.set() + + return iterator.async_sse_wrapper(_stream()) + + +async def _drain_with_pauses_until_upstream_exhausted(gen, upstream_exhausted: asyncio.Event) -> list: + received = [] + while not upstream_exhausted.is_set(): + received.append(await gen.__anext__()) + for _ in range(25): + await asyncio.sleep(0) + assert len(received) <= len(COMPLETE_STREAM_EVENTS) + return received + + +@pytest.mark.asyncio +async def test_normal_end_parks_deferred_logging_even_when_sentinel_enqueue_backpressured(monkeypatch): + """ + Regression: with a full relay queue at end of stream, the pump suspends + while enqueueing the end-of-stream sentinel, and a client that then drains + the whole tail tears the relay down (setting ``client_detached``) before + the pump resumes. That teardown is a normally completed response, not a + disconnect: billing must still park for the proxy's post-response hook + (preserving post_call decoration such as guardrail_information) instead of + enqueueing immediately through the teardown path. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 2) + worker = _RecordingLoggingWorker() + monkeypatch.setattr(streaming_iterator_module, "GLOBAL_LOGGING_WORKER", worker) + + dispatched = [] + + async def _deferred_stream_complete(logging_coroutine): + dispatched.append(logging_coroutine) + logging_coroutine.close() + + iterator = _make_iterator("test_sentinel_backpressure_normal_end") + iterator.litellm_logging_obj._on_deferred_stream_complete = _deferred_stream_complete + + upstream_exhausted = asyncio.Event() + gen = _backpressured_wrapper(iterator, upstream_exhausted) + received = await _drain_with_pauses_until_upstream_exhausted(gen, upstream_exhausted) + + while True: + try: + received.append(await gen.__anext__()) + except StopAsyncIteration: + break + + for _ in range(100): + if worker.enqueued or getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None): + break + await asyncio.sleep(0.01) + + assert len(received) == len(COMPLETE_STREAM_EVENTS) + assert worker.enqueued == [], "fully delivered stream billed through the teardown path" + assert dispatched == [] + parked = getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None) + assert parked is not None, "pump never parked deferred billing" + parked[0].close() + + +@pytest.mark.asyncio +async def test_relay_teardown_dispatches_deferred_billing_when_sentinel_never_consumed(monkeypatch): + """ + Regression: when the pump has parked deferred billing but its end-of-stream + sentinel never fits in the full relay queue (the client disconnects without + draining the tail), the proxy's post-response hook never fires. Exactly one + of the relay teardown or the pump's fallback must dispatch the parked + billing, or the request logs no spend at all. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 2) + worker = _RecordingLoggingWorker() + monkeypatch.setattr(streaming_iterator_module, "GLOBAL_LOGGING_WORKER", worker) + + dispatched = [] + deferred_fired = asyncio.Event() + + def _deferred_stream_complete(logging_coroutine): + dispatched.append(logging_coroutine) + + async def _consume(): + logging_coroutine.close() + deferred_fired.set() + + return _consume() + + iterator = _make_iterator("test_sentinel_never_consumed_dispatch") + iterator.litellm_logging_obj._on_deferred_stream_complete = _deferred_stream_complete + + upstream_exhausted = asyncio.Event() + gen = _backpressured_wrapper(iterator, upstream_exhausted) + await _drain_with_pauses_until_upstream_exhausted(gen, upstream_exhausted) + + for _ in range(100): + if getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None) is not None: + break + await asyncio.sleep(0.01) + + await gen.aclose() + + for _ in range(100): + if dispatched: + break + await asyncio.sleep(0.01) + + assert len(dispatched) == 1, "parked billing was never dispatched" + assert getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None) is None + assert getattr(iterator.litellm_logging_obj, "_on_deferred_stream_complete", None) is None + assert len(worker.enqueued) == 1, "teardown billing enqueued alongside the deferred dispatch" + await worker.enqueued[0] + assert deferred_fired.is_set() From e34f43328c8f6e0bd0df10747f2453ac0433b684 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 31 Aug 2026 17:33:12 -0700 Subject: [PATCH 120/120] fix(proxy): ship psycopg so partitioned SpendLogs detection actually runs (#38994) ProxyExtrasDBManager.spend_logs_is_partitioned() (#38452) silently returns False when psycopg can't be imported, and psycopg was never added to the extra_proxy install, so every production image lacks it. Schema reconciliation then generates the unfiltered primary-key rewrite against a genuinely partitioned LiteLLM_SpendLogs and Postgres rejects it, exactly the failure the fix was meant to prevent. Ships psycopg via extra_proxy and logs a warning when it's still missing instead of failing silently. --- .../litellm_proxy_extras/utils.py | 7 ++++++ pyproject.toml | 5 +++++ .../test_litellm_proxy_extras_utils.py | 22 +++++++++++++++++++ uv.lock | 4 ++++ 4 files changed, 38 insertions(+) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index b2dc0a52c8f..b8032dd0d28 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -512,6 +512,13 @@ class ProxyExtrasDBManager: try: import psycopg except ImportError: + logger.warning( + "psycopg is not installed; skipping the LiteLLM_SpendLogs " + "partition check. If this table is partitioned (see " + "db_scripts/partition_spend_logs.sql), schema reconciliation " + "will try to rewrite its primary key and fail. Install the " + "litellm[extra_proxy] extra, which now includes psycopg." + ) return False cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url) diff --git a/pyproject.toml b/pyproject.toml index 34c1fec1c11..96dcf1121bc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -91,6 +91,11 @@ cli = [ ] extra_proxy = [ "prisma>=0.11.0,<1.0", + # Used by ProxyExtrasDBManager.spend_logs_is_partitioned() to detect a + # partitioned LiteLLM_SpendLogs and keep schema reconciliation from + # fighting its composite primary key. + "psycopg>=3.2,<4.0", + "psycopg-binary>=3.2,<4.0", "azure-identity>=1.25.2,<2.0", "azure-keyvault-secrets>=4.10.0,<5.0", # Not in PyPI proxy extra. diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 498d0cb4723..b3d457707b8 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -681,3 +681,25 @@ class TestSpendLogsPartitionDetectionSchemaScope: def test_only_partitioned_relations_match(self, monkeypatch): query, _ = self._detect(monkeypatch, "postgresql://u:p@localhost:5432/db") assert "pg_partitioned_table" in query + + +class TestSpendLogsPartitionDetectionMissingPsycopg: + """psycopg ships in the `extra_proxy` install, but a stripped-down image + can still lack it. When it does, detection must fail closed to False + (never crash the migration path) and say so loudly, because a silent + False here is what let a genuinely partitioned LiteLLM_SpendLogs hit the + unfiltered primary-key rewrite in production.""" + + def test_missing_psycopg_returns_false(self, monkeypatch): + monkeypatch.setitem(sys.modules, "psycopg", None) + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:5432/db") + assert ProxyExtrasDBManager.spend_logs_is_partitioned() is False + + def test_missing_psycopg_logs_a_warning(self, monkeypatch, caplog): + monkeypatch.setitem(sys.modules, "psycopg", None) + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:5432/db") + with caplog.at_level("WARNING", logger="litellm_proxy_extras"): + ProxyExtrasDBManager.spend_logs_is_partitioned() + assert any( + "psycopg is not installed" in record.message for record in caplog.records + ) diff --git a/uv.lock b/uv.lock index 8ef72116466..7df1a35b28a 100644 --- a/uv.lock +++ b/uv.lock @@ -4306,6 +4306,8 @@ extra-proxy = [ { name = "google-cloud-iam" }, { name = "google-cloud-kms" }, { name = "prisma" }, + { name = "psycopg" }, + { name = "psycopg-binary" }, { name = "redisvl" }, { name = "resend" }, ] @@ -4544,6 +4546,8 @@ requires-dist = [ { name = "polars", marker = "extra == 'proxy'", specifier = ">=1.38.1,<2.0" }, { name = "prisma", marker = "extra == 'extra-proxy'", specifier = ">=0.11.0,<1.0" }, { name = "prometheus-client", marker = "extra == 'proxy-runtime'", specifier = ">=0.20.0,<1.0" }, + { name = "psycopg", marker = "extra == 'extra-proxy'", specifier = ">=3.2,<4.0" }, + { name = "psycopg-binary", marker = "extra == 'extra-proxy'", specifier = ">=3.2,<4.0" }, { name = "pydantic", specifier = ">=2.10.0,<3.0.0" }, { name = "pydantic-settings", specifier = ">=2.14.1,<3.0" }, { name = "pyjwt", marker = "extra == 'proxy'", specifier = ">=2.13.0,<3.0" },