From cb94567ed7f964c217120beab67e5e2c36738928 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Tue, 13 Oct 2026 11:03:58 -0700 Subject: [PATCH 001/325] feat(vertex-ai): add gemini-4-argon-cyber-permissive (#46249) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 58 +++++++++++++++++++ model_prices_and_context_window.json | 58 +++++++++++++++++++ 2 files changed, 116 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b34dd9fa764..8977c5dd88e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -72689,5 +72689,63 @@ "supports_vision": true, "supports_web_search": true, "web_search_billing_unit": "per_query" + }, + "vertex_ai/gemini-4-argon-cyber-permissive": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "cache_read_input_token_cost_flex": 1e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "google_maps_grounding_cost_per_query": 0.014, + "input_cost_per_audio_token": 4e-06, + "input_cost_per_audio_token_batches": 2e-06, + "input_cost_per_audio_token_priority": 7.2e-06, + "input_cost_per_token": 4e-06, + "input_cost_per_token_batches": 2e-06, + "input_cost_per_token_flex": 2e-06, + "input_cost_per_token_priority": 7.2e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_reasoning_token": 2e-05, + "output_cost_per_token": 2e-05, + "output_cost_per_token_batches": 1e-05, + "output_cost_per_token_flex": 1e-05, + "output_cost_per_token_priority": 3.6e-05, + "prompt_cache_min_tokens": 6144.0, + "regional_endpoint_uplift_multiplier": 1.1, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b34dd9fa764..8977c5dd88e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -72689,5 +72689,63 @@ "supports_vision": true, "supports_web_search": true, "web_search_billing_unit": "per_query" + }, + "vertex_ai/gemini-4-argon-cyber-permissive": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "cache_read_input_token_cost_flex": 1e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "google_maps_grounding_cost_per_query": 0.014, + "input_cost_per_audio_token": 4e-06, + "input_cost_per_audio_token_batches": 2e-06, + "input_cost_per_audio_token_priority": 7.2e-06, + "input_cost_per_token": 4e-06, + "input_cost_per_token_batches": 2e-06, + "input_cost_per_token_flex": 2e-06, + "input_cost_per_token_priority": 7.2e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_reasoning_token": 2e-05, + "output_cost_per_token": 2e-05, + "output_cost_per_token_batches": 1e-05, + "output_cost_per_token_flex": 1e-05, + "output_cost_per_token_priority": 3.6e-05, + "prompt_cache_min_tokens": 6144.0, + "regional_endpoint_uplift_multiplier": 1.1, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" } } From fb051f9647b655abc14fe6e5f5c0ae130e87bace Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 13 Oct 2026 18:12:54 +0000 Subject: [PATCH 002/325] chore: bump litellm-enterprise 0.1.76 -> 0.1.77, litellm-proxy-extras 0.4.108 -> 0.4.109, litellm 1.106.0 -> 1.107.0 (#46239) * chore: bump litellm 1.106.0 -> 1.107.0 * chore: bump litellm-enterprise 0.1.76 -> 0.1.77, litellm-proxy-extras 0.4.108 -> 0.4.109 --------- Co-authored-by: Yuneng Jiang --- enterprise/pyproject.toml | 4 ++-- litellm-proxy-extras/pyproject.toml | 4 ++-- pyproject.toml | 8 ++++---- uv.lock | 6 +++--- 4 files changed, 11 insertions(+), 11 deletions(-) diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 78862d4c942..4ffb7169f24 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.76" +version = "0.1.77" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.76" +version = "0.1.77" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 68efeac71ae..61d56c6b301 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.108" +version = "0.4.109" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -30,7 +30,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.108" +version = "0.4.109" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/pyproject.toml b/pyproject.toml index 0ed8b8757ac..0be8adb7337 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.106.0" +version = "1.107.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.15" @@ -77,8 +77,8 @@ proxy = [ "mcp>=2.2.0,<3", "httpx2>=2.5.0,<3", "pydantic>=2.12.0,<3", - "litellm-proxy-extras==0.4.108", - "litellm-enterprise==0.1.76", + "litellm-proxy-extras==0.4.109", + "litellm-enterprise==0.1.77", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", @@ -372,7 +372,7 @@ litellm-enterprise = { workspace = true } members = ["enterprise", "litellm-proxy-extras"] [tool.commitizen] -version = "1.106.0" +version = "1.107.0" version_files = [ "pyproject.toml:^version", ] diff --git a/uv.lock b/uv.lock index 471460a32b5..05331266b3e 100644 --- a/uv.lock +++ b/uv.lock @@ -4500,7 +4500,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.106.0" +version = "1.107.0" source = { editable = "." } dependencies = [ { name = "aiohttp" }, @@ -4996,12 +4996,12 @@ provides-extras = ["test"] [[package]] name = "litellm-enterprise" -version = "0.1.76" +version = "0.1.77" source = { editable = "enterprise" } [[package]] name = "litellm-proxy-extras" -version = "0.4.108" +version = "0.4.109" source = { editable = "litellm-proxy-extras" } dependencies = [ { name = "psycopg" }, From 784581bace35819000712a6307df327a73a4c6cd Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 13 Oct 2026 11:42:39 -0700 Subject: [PATCH 003/325] fix(proxy)!: return provider-matching 4xx for the LIT-8062 status code follow-ups (#45537) * fix(proxy): return provider-matching 4xx for the LIT-8062 status code follow-ups - OCR multipart errors are 400 invalid_request_error on the file param - a missing model is a 400 on chat, responses and messages even with a wildcard deployment - responses max_output_tokens <= 0 reaches the provider instead of being raised to 16 - daily activity rejects end_date before start_date with a 400 - vector store file attributes are sent in full so the provider enforces its key limit - router-wide defaults no longer satisfy a param the Router takes positionally, or a model the router does not serve - a JSON null user_config is treated as absent * chore: keep the type gates at their merge-base counts * fix(proxy): keep fixed-target wildcards serving a missing model and close router-default gaps - a missing model is rejected only when the wildcard would forward it into the target - moderation only counts router defaults for models get_model_list lists - user_config defaults and deployment params never satisfy a positionally bound param * test(integration): pin the fixed-target wildcard path and no provider call on positional-param 400s * chore: validate pattern-matched deployments instead of reading untyped dicts * fix(proxy): keep router defaults counting for a2a agent models * test: drop the signature-walk test and type the new test helpers * test(integration): audit matrix for the status-code follow-ups Drives every mode the fix branches on through a real proxy, Postgres, Redis and the scripted upstream: positional params against router, deployment and user_config defaults; moderation listed vs unlisted; unknown and a2a models; every user_config shape on chat, responses and messages, streaming and SDK clients; forwarding vs fixed-target wildcards for missing, null and empty models; Responses max_output_tokens edges; OCR parse errors; vector store file attributes; every daily activity route; and a provider outage mid burst plus a killed worker on a two-worker proxy. Product gaps found are skipped as BUG. * style: ruff format the audit test files * test: name the auth pre-check as the source of the JSON array 500 * test: build the audit helpers without mutation and share one sentinel * test: build every audit request body as one literal * fix(proxy): keep router defaults and model-less dispatch for requests a fallback chain serves * fix(proxy): count request, key and team fallbacks when an unlisted model inherits router defaults * fix(proxy): skip key and team fallbacks on the client-credential dispatch path * fix(proxy): keep router defaults when the router's own * fallback serves an unlisted model * fix(proxy): count client-style fallbacks, forward an empty model, report invalid usage chat dates - The router-default pre-check now treats client-style request, key and team fallbacks ([{"model": ...}]) as serving an unlisted model, as the router does - The forwarding-wildcard gate only rejects a missing or null model; an empty model reaches the provider as before (vLLM serves its loaded model for it) - The /usage/ai/chat tools hand the model the date-range error instead of an empty dataset or a generic fetch error * test(usage): type the invalid-date tool test * fix(proxy): keep rejecting an empty model behind a forwarding wildcard Key and team model checks skip an empty model, so forwarding it let a key restricted to another model reach a server's default model (vLLM serves its loaded model for ""). Pin that with a restricted-key integration cell. * fix(usage): report the query layer's date 400 to the usage chat model Checking dates before the fetch refused formats the global usage tool's query accepts (2026/10/01 returned spend on main). Catch the daily activity query's own 400 instead, so each tool rejects exactly what it rejected before and only the message changes. * fix(proxy): restore three model-less and unlisted-model paths main served An empty model is looked up as sent, so next to a forwarding */* wildcard it still reaches the fixed * deployment it matches. A null model with a prompt_id is left to the prompt manager, which may pick the model. A Claude Code subagent turn on a session-bound auto/complexity router counts as served, so the router default max_tokens still applies. * fix(proxy): only skip the model check for a prompt_id when a prompt manager is registered Without a callback prompt manager nothing can pick the model, so a null model with a prompt_id behind a forwarding wildcard gets the same 400 as any other null model. * fix(proxy): leave a model-less Claude Code subagent turn to its session router A subagent turn on a session bound to an auto, complexity, adaptive or quality router is answered by that router whatever model it names, so a null or missing model behind a forwarding wildcard is no longer rejected. An empty model stays rejected. * fix(proxy): skip the wildcard model check only for a bound Claude Code session Read the session's router binding, as the router does, instead of only checking that the request could resume one, so a subagent turn without a model on an unbound session is still rejected before a forwarding wildcard sends it upstream. * refactor(proxy): stop predicting routing in the proxy pre-checks - drop the wildcard missing-model gate: a model-less request behind a forwarding wildcard goes to the router and gets the provider's own status, as on main - drop the "model the router serves" check: router defaults count for any named model again, except moderation, where the Router merges them only for listed models - keep the positional-param, moderation and null user_config fixes * test(ocr): cover the form, size, empty body, invalid JSON and file-document 400s --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- .../llms/openai/responses/transformation.py | 6 +- .../vector_store_files/transformation.py | 18 +- .../common_daily_activity.py | 2 + .../management_endpoints/team_endpoints.py | 2 - .../usage_endpoints/ai_usage_chat.py | 21 +- litellm/proxy/ocr_endpoints/endpoints.py | 26 +- litellm/proxy/route_llm_request.py | 40 +- .../compatibility/_status_code_audit.py | 280 ++++++++++++ .../test_status_code_followups.py | 374 +++++++++++++++ .../test_status_code_followups_chaos.py | 283 ++++++++++++ .../test_status_code_followups_defaults.py | 391 ++++++++++++++++ .../test_status_code_followups_requests.py | 425 ++++++++++++++++++ .../test_status_code_followups_usage_chat.py | 125 ++++++ .../test_openai_responses_transformation.py | 7 +- ...penai_vector_store_files_transformation.py | 26 ++ .../test_common_daily_activity.py | 9 + .../usage_endpoints/test_ai_usage_chat.py | 131 ++++++ .../proxy/ocr_endpoints/test_endpoints.py | 86 +++- tests/unit/proxy/test_route_llm_request.py | 130 ++++++ 19 files changed, 2345 insertions(+), 37 deletions(-) create mode 100644 tests/integration/compatibility/_status_code_audit.py create mode 100644 tests/integration/compatibility/test_status_code_followups.py create mode 100644 tests/integration/compatibility/test_status_code_followups_chaos.py create mode 100644 tests/integration/compatibility/test_status_code_followups_defaults.py create mode 100644 tests/integration/compatibility/test_status_code_followups_requests.py create mode 100644 tests/integration/compatibility/test_status_code_followups_usage_chat.py diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index ec4f0c04475..abd3d03b139 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -183,10 +183,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): OpenAI's Responses API rejects max_output_tokens below 16 for every model (not gpt-5 specific), so a client like Claude Code that sends a max_tokens=1 - warmup probe on model switch would otherwise 400. Values that are None or - already at/above the minimum are returned unchanged. + warmup probe on model switch would otherwise 400. Values that are None, + already at/above the minimum, or not positive are returned unchanged. """ - if isinstance(max_output_tokens, int) and max_output_tokens < OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS: + if isinstance(max_output_tokens, int) and 0 < max_output_tokens < OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS: return OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS return max_output_tokens diff --git a/litellm/llms/openai/vector_store_files/transformation.py b/litellm/llms/openai/vector_store_files/transformation.py index 8519b5f4bc4..d5f15d639d6 100644 --- a/litellm/llms/openai/vector_store_files/transformation.py +++ b/litellm/llms/openai/vector_store_files/transformation.py @@ -20,13 +20,17 @@ from litellm.types.vector_store_files import ( VectorStoreFileObject, VectorStoreFileUpdateRequest, ) -from litellm.utils import add_openai_metadata def _clean_dict(source: Mapping[str, object]) -> dict[str, object]: return {k: v for k, v in source.items() if v is not None} +def _visible_attributes(attributes: Mapping[str, object]) -> Mapping[str, str]: + """The caller's string attributes, all of them, so the provider enforces its own key limit.""" + return {str(k): v for k, v in attributes.items() if k != "hidden_params" and isinstance(v, str)} + + class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): ASSISTANTS_HEADER_KEY = "OpenAI-Beta" ASSISTANTS_HEADER_VALUE = "assistants=v2" @@ -106,11 +110,7 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): payload: Final[dict[str, object]] = _clean_dict(dict(create_request)) attributes: Final = payload.get("attributes") if isinstance(attributes, dict): - filtered_attributes: Final = add_openai_metadata(attributes) - if filtered_attributes is not None: - payload["attributes"] = filtered_attributes - else: - payload.pop("attributes", None) + payload["attributes"] = _visible_attributes(attributes) url: Final = api_base return url, payload @@ -211,11 +211,7 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): payload: Final[dict[str, object]] = dict(update_request) attributes: Final = payload.get("attributes") if isinstance(attributes, dict): - filtered_attributes: Final = add_openai_metadata(attributes) - if filtered_attributes is not None: - payload["attributes"] = filtered_attributes - else: - payload.pop("attributes", None) + payload["attributes"] = _visible_attributes(attributes) encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") return f"{api_base}/{encoded_file_id}", payload diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 28dfbb09eab..f948a49441a 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -90,6 +90,8 @@ def parse_canonical_date_range(start_date: str | None, end_date: str | None) -> end: Final = parse_canonical_date(end_date) if start is None or end is None: return InvalidDateRange(reason="start_date and end_date must be valid YYYY-MM-DD dates") + if end < start: + return InvalidDateRange(reason="end_date must be on or after start_date") return CanonicalDateRange(start=start, end=end) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 544bc471ce5..6a4c8ac2c9e 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -6833,8 +6833,6 @@ def aggregated_date_range_error(start_date: str | None, end_date: str | None) -> date_range: Final = parse_canonical_date_range(start_date, end_date) if isinstance(date_range, InvalidDateRange): return date_range.reason - if date_range.end < date_range.start: - return "end_date must be on or after start_date" if (date_range.end - date_range.start).days > _MAX_AGGREGATED_RANGE_DAYS: return f"Date range must be at most {_MAX_AGGREGATED_RANGE_DAYS} days" return None diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py index 3c2300f14fd..7a45d0438de 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -497,12 +497,29 @@ async def _execute_tool_call( user_id: str | None, is_admin: bool, ) -> str: - """Run a single tool and return the summarised result text.""" + """Run a single tool and return the summarised result text, or the reason the query + layer rejected its dates so the model can correct them.""" kwargs: Final = _resolve_fetch_kwargs(fn_name, fn_args, user_id, is_admin) - raw_data: Final = await handler["fetch"](**kwargs) + try: + raw_data: Final = await handler["fetch"](**kwargs) + except HTTPException as e: + reason: Final = _invalid_date_reason(e) + if reason is None: + raise + return f"Invalid date range: {reason}" return handler["summarise"](raw_data) +def _invalid_date_reason(error: HTTPException) -> str | None: + """The reason in the daily activity query layer's 400 for invalid or reversed dates.""" + detail: Final = cast(object, error.detail) # cast-ok: starlette types detail as str; the query layer passes a dict + match detail: + case {"error": str() as reason} if error.status_code == 400: + return reason + case _: + return None + + async def _process_tool_call( tc: ChatCompletionMessageToolCall, chat_messages: list[Mapping[str, object]], diff --git a/litellm/proxy/ocr_endpoints/endpoints.py b/litellm/proxy/ocr_endpoints/endpoints.py index 94fe9354386..62fde60dcac 100644 --- a/litellm/proxy/ocr_endpoints/endpoints.py +++ b/litellm/proxy/ocr_endpoints/endpoints.py @@ -17,6 +17,7 @@ from litellm.llms.base_llm.ocr.transformation import ( parse_ocr_request_format, ) from litellm.proxy._types import * +from litellm.proxy._types import ProxyException from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -32,6 +33,15 @@ class _NamedUpload(io.BytesIO): self.name = name +class OCRRequestError(ProxyException): + """A malformed OCR request, answered with a 400 before any provider call.""" + + def __init__(self, message: str, param: str | None = None) -> None: + super().__init__( # pyright: ignore[reportUnknownMemberType] # ProxyException.__init__ takes an untyped provider_specific_fields dict + message=message, type="invalid_request_error", param=param, code=400 + ) + + def _build_document_from_upload( file_content: bytes, filename: str | None, @@ -99,7 +109,7 @@ async def _parse_multipart_form(request: Request) -> dict[str, object]: try: form: Final = await request.form() except Exception as e: - raise ValueError( + raise OCRRequestError( f"Failed to parse multipart form data: {e}. " "When using curl with --form/-F, do NOT set the Content-Type header " "manually — curl will set it automatically with the required boundary." @@ -110,7 +120,7 @@ async def _parse_multipart_form(request: Request) -> dict[str, object]: # depending on middleware; check both via isinstance (FastAPI's UploadFile # is a subclass of Starlette's) and fall back to duck-type check. if uploaded_file is None or (not isinstance(uploaded_file, UploadFile) and not hasattr(uploaded_file, "read")): - raise ValueError("Multipart OCR request must include a 'file' field with the document to process") + raise OCRRequestError("Multipart OCR request must include a 'file' field with the document to process", "file") uploaded_file = cast(UploadFile, uploaded_file) @@ -118,9 +128,9 @@ async def _parse_multipart_form(request: Request) -> dict[str, object]: await uploaded_file.seek(0) file_content: Final = await uploaded_file.read(_MAX_FILE_BYTES + 1) if not file_content: - raise ValueError("Uploaded file is empty") + raise OCRRequestError("Uploaded file is empty", "file") if len(file_content) > _MAX_FILE_BYTES: - raise ValueError("OCR file exceeds the size limit") + raise OCRRequestError("OCR file exceeds the size limit", "file") document: Final = _build_document_from_upload( file_content=file_content, @@ -197,7 +207,7 @@ async def _parse_ocr_request_body(request: Request) -> dict[str, object]: ) return await _parse_multipart_form(request) - raise ValueError( + raise OCRRequestError( "Empty request body. For file uploads, use multipart/form-data content type " "with a file field. When using curl with --form/-F, do NOT set the Content-Type " "header manually." @@ -206,7 +216,7 @@ async def _parse_ocr_request_body(request: Request) -> dict[str, object]: try: data: Final = orjson.loads(body) except orjson.JSONDecodeError as e: - raise ValueError( + raise OCRRequestError( f"Invalid JSON in request body: {e}. " "Ensure the request body is valid JSON with Content-Type: application/json, " "or use multipart/form-data for file uploads." @@ -220,7 +230,7 @@ async def _parse_ocr_request_body(request: Request) -> dict[str, object]: # File uploads must go through multipart/form-data instead. doc: Final = data.get("document") if isinstance(data, dict) else None if isinstance(doc, dict) and doc.get("type") == "file": - raise ValueError( + raise OCRRequestError( "document type 'file' is not supported through the JSON API. " "To upload a local file, use multipart/form-data with a 'file' field. " "For JSON requests, use 'document_url' or 'image_url' document types." @@ -237,7 +247,7 @@ async def _parse_ocr_request_body(request: Request) -> dict[str, object]: for url_field in ("document_url", "image_url"): url_value = doc.get(url_field) if isinstance(url_value, str) and url_value.startswith("reducto://"): - raise ValueError( + raise OCRRequestError( "reducto:// file IDs are not accepted through the proxy " "OCR API; upload the file in the same request via " "multipart/form-data with a 'file' field, or pass an " diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index d8b8a68c8f9..6f93bcd82dc 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -246,13 +246,18 @@ def _find_missing_required_body_param( if not missing_present_params: return None candidate_litellm_params: Final = _candidate_deployment_litellm_params(data, llm_router) - router_default_litellm_params: Final = _router_default_litellm_params(route_type, data, llm_router) + router_default_litellm_params: Final = _router_default_litellm_params( + route_type, data, llm_router, model_listed=bool(candidate_litellm_params) + ) missing_param: Final = next( ( param for param in missing_present_params if router_default_litellm_params.get(param) is None - and not any(deployment_params.get(param) is not None for deployment_params in candidate_litellm_params) + and ( + route_type in _ROUTE_TYPES_WITH_POSITIONAL_REQUIRED_PARAM + or not any(deployment_params.get(param) is not None for deployment_params in candidate_litellm_params) + ) ), None, ) @@ -261,19 +266,34 @@ def _find_missing_required_body_param( return MissingBodyParam(name=missing_param, model_deployments_loaded=bool(candidate_litellm_params)) -_ROUTE_TYPES_WITHOUT_ROUTER_DEFAULTS_MERGE: Final[frozenset[str]] = frozenset( - {"asearch", "acreate_agent", "acreate_eval", "acreate_run"} +# The Router method takes the route's required param positionally, so it binds before +# any default or deployment param is merged in, and only the request body can supply it. +_ROUTE_TYPES_WITH_POSITIONAL_REQUIRED_PARAM: Final[frozenset[str]] = frozenset( + {"aimage_generation", "aspeech", "atext_completion", "atranscription"} ) +_ROUTE_TYPES_WITHOUT_ROUTER_DEFAULTS_MERGE: Final[frozenset[str]] = ( + frozenset({"asearch", "acreate_agent", "acreate_eval", "acreate_run"}) | _ROUTE_TYPES_WITH_POSITIONAL_REQUIRED_PARAM +) + +# The Router merges defaults on these only when `get_model_list` lists the model. +_ROUTE_TYPES_MERGING_ONLY_LISTED_MODELS: Final[frozenset[str]] = frozenset({"amoderation"}) + def _router_default_litellm_params( route_type: str, data: Mapping[str, object], llm_router: LitellmRouter | None, + *, + model_listed: bool, ) -> Mapping[str, object]: # Mirror exactly the defaults the dispatching router will merge at dispatch time: - # user_config requests dispatch on their own throwaway Router, and the listed route - # types (plus model-less direct dispatch) never pass through the router's merge. + # the listed route types either never pass through the router's merge or take the + # param positionally, which binds before the merge; user_config requests dispatch on + # their own throwaway Router; model-less direct dispatch never merges; and moderation + # merges only for a model `get_model_list` lists. + if route_type in _ROUTE_TYPES_WITHOUT_ROUTER_DEFAULTS_MERGE: + return {} user_config: Final[Mapping[str, object] | None] = ( data.get("user_config") if isinstance(data.get("user_config"), Mapping) else None ) @@ -281,7 +301,11 @@ def _router_default_litellm_params( defaults: Final[Mapping[str, object] | None] = user_config.get("default_litellm_params") return defaults if isinstance(defaults, Mapping) else {} model_name: Final = data.get("model") - if route_type in _ROUTE_TYPES_WITHOUT_ROUTER_DEFAULTS_MERGE or not isinstance(model_name, str) or not model_name: + if ( + not isinstance(model_name, str) + or not model_name + or (route_type in _ROUTE_TYPES_MERGING_ONLY_LISTED_MODELS and not model_listed) + ): return {} router_defaults: Final[Mapping[str, object] | None] = getattr(llm_router, "default_litellm_params", None) return router_defaults if isinstance(router_defaults, Mapping) else {} @@ -628,6 +652,8 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr raise_if_mock_testing_params_disallowed(data, allowed=mock_testing_params_allowed()) data.pop("enable_tag_filtering", None) + if "user_config" in data and data["user_config"] is None: + data.pop("user_config") # pyright: ignore[reportUnknownMemberType] # route_request takes an untyped dict body team_id: Final = get_team_id_from_data(data) router_model_names: Final = llm_router.model_names if llm_router is not None else [] diff --git a/tests/integration/compatibility/_status_code_audit.py b/tests/integration/compatibility/_status_code_audit.py new file mode 100644 index 00000000000..59513891ec0 --- /dev/null +++ b/tests/integration/compatibility/_status_code_audit.py @@ -0,0 +1,280 @@ +"""Shared payloads and wire helpers for the status-code follow-up audit cells. + +Every cell drives a real proxy against the scripted upstream and reads the outbound provider body back from the +upstream's observation feed, so a cell asserts both what the caller got and what the provider received.""" + +from __future__ import annotations + +import json +import os +import uuid +from collections.abc import Iterator, Mapping, Sequence +from contextlib import contextmanager +from pathlib import Path +from typing import Final + +import httpx +from integration._support.client import JSON_OBJECT, Gateway, Scenario, gateway_from_environment, object_value +from integration._support.process import owned_proxy_process +from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario +from pydantic import JsonValue + +from tests.integration.cost_calculation.cost_tracking_case import ( + JsonResponse, + SseResponse, + StoredResponse, +) + +ROUTER_DEFAULTS: Final[dict[str, JsonValue]] = { + "prompt": "router default prompt", + "input": "router default input", + "max_tokens": 32, +} +CHAT: Final[dict[str, JsonValue]] = { + "id": "chatcmpl-$UNIQUE_ID", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "scripted"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} +RESPONSE: Final[dict[str, JsonValue]] = { + "id": "resp_$UNIQUE_ID", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_$UNIQUE_ID", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "scripted", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, +} +MESSAGE: Final[dict[str, JsonValue]] = { + "id": "msg_$UNIQUE_ID", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "scripted"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, +} +MODERATION: Final[dict[str, JsonValue]] = { + "id": "modr-$UNIQUE_ID", + "model": "omni-moderation-latest", + "results": [{"flagged": False, "categories": {}, "category_scores": {}}], +} +CHAT_FRAMES: Final = ( + 'data: {"id":"chatcmpl-$UNIQUE_ID","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",' + '"choices":[{"index":0,"delta":{"role":"assistant","content":"streamed "},"finish_reason":null}]}', + 'data: {"id":"chatcmpl-$UNIQUE_ID","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",' + '"choices":[{"index":0,"delta":{"content":"response"},"finish_reason":null}]}', + 'data: {"id":"chatcmpl-$UNIQUE_ID","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",' + '"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", +) +RESPONSES_FRAMES: Final = ( + "event: response.created\n" + 'data: {"type":"response.created","sequence_number":0,"response":{"id":"resp_$UNIQUE_ID","object":"response",' + '"created_at":1,"status":"in_progress","model":"gpt-4o-mini","output":[],"usage":null}}', + "event: response.output_text.delta\n" + 'data: {"type":"response.output_text.delta","sequence_number":1,"item_id":"msg_$UNIQUE_ID","output_index":0,' + '"content_index":0,"delta":"streamed response"}', + "event: response.completed\n" + 'data: {"type":"response.completed","sequence_number":2,"response":{"id":"resp_$UNIQUE_ID","object":"response",' + '"created_at":1,"status":"completed","model":"gpt-4o-mini","output":[{"type":"message","id":"msg_$UNIQUE_ID",' + '"status":"completed","role":"assistant","content":[{"type":"output_text","text":"streamed response",' + '"annotations":[]}]}],"usage":{"input_tokens":1,"output_tokens":2,"total_tokens":3}}}', +) +MESSAGE_FRAMES: Final = ( + "event: message_start\n" + 'data: {"type":"message_start","message":{"id":"msg_$UNIQUE_ID","type":"message","role":"assistant",' + '"model":"claude-haiku-4-5","content":[],"stop_reason":null,"stop_sequence":null,' + '"usage":{"input_tokens":1,"output_tokens":0}}}', + "event: content_block_start\n" + 'data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}', + "event: content_block_delta\n" + 'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"streamed response"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + "event: message_delta\n" + 'data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":2}}', + 'event: message_stop\ndata: {"type":"message_stop"}', +) +UNSET: Final = object() +USER_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Hello"}] +STREAM_TERMINALS: Final[Mapping[str, str]] = { + "chat": "data: [DONE]", + "responses": '"type":"response.completed"', + "messages": "event: message_stop", +} +ENDPOINT_PATHS: Final[Mapping[str, str]] = { + "chat": "/v1/chat/completions", + "responses": "/v1/responses", + "messages": "/v1/messages", +} +ENDPOINT_BODIES: Final[Mapping[str, dict[str, JsonValue]]] = { + "chat": {"messages": USER_MESSAGES}, + "responses": {"input": "Hello", "store": False}, + "messages": {"messages": USER_MESSAGES, "max_tokens": 16}, +} + + +def json_response(body: dict[str, JsonValue], status: int = 200) -> JsonResponse: + return JsonResponse(content_type="application/json", body=body, status=status) + + +def sse_response(frames: tuple[str, ...]) -> SseResponse: + return SseResponse(content_type="text/event-stream", frames=frames) + + +def openai_error(message: str, param: str) -> JsonResponse: + return json_response( + {"error": {"message": message, "type": "invalid_request_error", "param": param, "code": None}}, 400 + ) + + +class Upstream: + """Reads the scripted upstream's destructive observation feed and keeps every record it drained.""" + + def __init__(self, url: str) -> None: + self.url = url.rstrip("/") + self.items: tuple[dict[str, JsonValue], ...] = () + + def drain(self) -> None: + with httpx.Client(timeout=10, trust_env=False) as client: + payload: Final = JSON_OBJECT.validate_python( + client.get(f"{self.url}/__observations?include_method=true").json() + ) + requests: Final = payload.get("requests") + assert isinstance(requests, list) + self.items = (*self.items, *(object_value(item) for item in requests if isinstance(item, dict))) + + def calls(self, identity: str) -> tuple[dict[str, JsonValue], ...]: + self.drain() + return tuple( + item + for item in self.items + if f"/{identity}/" in str(item.get("path")) and not str(item.get("path", "")).endswith("/models") + ) + + +def drain_rig_upstream() -> None: + """Drops every observation earlier cells left on the shared upstream, so a cell reads only its own calls.""" + Upstream(os.environ["INTEGRATION_UPSTREAM_URL"]).drain() + + +def assert_no_provider_call(gateway: Gateway, *identities: str) -> None: + upstream: Final = Upstream(gateway.upstream_url) + assert {identity: upstream.calls(identity) for identity in identities} == dict.fromkeys(identities, ()) + + +def one_outbound(gateway: Gateway, identity: str) -> dict[str, JsonValue]: + calls: Final = Upstream(gateway.upstream_url).calls(identity) + assert len(calls) == 1, calls + return object_value(calls[0]["body"]) + + +def post( + gateway: Gateway, + path: str, + body: Mapping[str, JsonValue], + *, + key: str | None = None, + headers: Mapping[str, str] | None = None, +) -> httpx.Response: + return gateway.client.post( + path, + json=dict(body), + headers={"Authorization": f"Bearer {gateway.key if key is None else key}", **(headers or {})}, + timeout=30, + ) + + +def stream_lines(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[int, tuple[str, ...]]: + with gateway.client.stream( + "POST", path, json=dict(body), headers={"Authorization": f"Bearer {gateway.key}"}, timeout=30 + ) as response: + return response.status_code, tuple(response.iter_lines()) + + +def error_body(response: httpx.Response) -> dict[str, JsonValue]: + """The error object of an OpenAI-shaped or Anthropic-shaped error response.""" + return object_value(JSON_OBJECT.validate_python(response.json())["error"]) + + +def invalid_request(response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 400, response.text + error: Final = error_body(response) + assert error.get("type") == "invalid_request_error", response.text + return error + + +def register(scenario: Scenario, prefix: str, response: StoredResponse) -> tuple[str, ScenarioHandle]: + identity: Final = f"{prefix}-{uuid.uuid4().hex}" + handle: Final = register_scenario(identity, response) + scenario.cleanups.callback(delete_scenario, handle) + return identity, handle + + +def deployment( + scenario: Scenario, name: str, model: str, response: StoredResponse, **extra: JsonValue +) -> tuple[str, dict[str, JsonValue]]: + identity, handle = register(scenario, name, response) + return identity, { + "model_name": name, + "litellm_params": {"model": model, "api_base": handle.api_base(), "api_key": identity, **extra}, + } + + +@contextmanager +def owned_gateway( + directory: Path, + config: Mapping[str, JsonValue], + *, + workers: int = 1, + environment: Mapping[str, str] | None = None, +) -> Iterator[Gateway]: + """A proxy booted from this checkout with `config` and extra `environment`, sharing the rig's database, Redis + and upstream.""" + path: Final = directory / "config.yaml" + path.write_text(json.dumps(dict(config)), encoding="utf-8") + with ( + gateway_from_environment() as gateway, + owned_proxy_process(gateway, directory, environment or {}, config=path, workers=workers) as owned, + ): + yield Gateway(owned.gateway.client, owned.gateway.key, gateway.upstream_url) + + +def stream_finished(endpoint: str, lines: Sequence[str]) -> bool: + """Whether the stream reached the endpoint's terminal frame.""" + return any(STREAM_TERMINALS[endpoint] in line for line in lines) + + +def _frame_text(frame: Mapping[str, JsonValue]) -> str: + """The text one chat, responses or messages delta frame carries, or "" for any other frame.""" + choices: Final = frame.get("choices") + first: Final = choices[0] if isinstance(choices, list) and choices else None + chat_delta: Final = first.get("delta") if isinstance(first, dict) else None + if isinstance(chat_delta, dict) and isinstance(chat_delta.get("content"), str): + return str(chat_delta["content"]) + delta: Final = frame.get("delta") + if frame.get("type") == "response.output_text.delta" and isinstance(delta, str): + return delta + if frame.get("type") == "content_block_delta" and isinstance(delta, dict) and delta.get("type") == "text_delta": + return str(delta.get("text")) + return "" + + +def assembled_text(lines: Sequence[str]) -> str: + """The text a chat, responses or messages stream carried, joined across its delta frames.""" + return "".join( + _frame_text(JSON_OBJECT.validate_json(line.removeprefix("data: "))) + for line in lines + if line.startswith("data: ") and line != "data: [DONE]" + ) diff --git a/tests/integration/compatibility/test_status_code_followups.py b/tests/integration/compatibility/test_status_code_followups.py new file mode 100644 index 00000000000..466eeeddde1 --- /dev/null +++ b/tests/integration/compatibility/test_status_code_followups.py @@ -0,0 +1,374 @@ +from __future__ import annotations + +import json +import uuid +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, +) +from integration._support.process import owned_proxy_process +from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario +from pydantic import JsonValue + +from tests.integration.cost_calculation.cost_tracking_case import BinaryResponse, JsonResponse + +_ROUTER_DEFAULTS: Final[dict[str, JsonValue]] = { + "prompt": "router default prompt", + "input": "router default input", + "max_tokens": 32, +} +_MESSAGE: Final[dict[str, JsonValue]] = { + "id": "msg-$UNIQUE_ID", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "scripted"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, +} +_CHAT: Final[dict[str, JsonValue]] = { + "id": "chatcmpl-$UNIQUE_ID", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "scripted"}, "finish_reason": "stop"}, + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} +_RESPONSE: Final[dict[str, JsonValue]] = { + "id": "resp_$UNIQUE_ID", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, +} +_VECTOR_STORE_FILE: Final[dict[str, JsonValue]] = { + "id": "file-audit", + "object": "vector_store.file", + "created_at": 1, + "usage_bytes": 0, + "vector_store_id": "vs-audit", + "status": "completed", + "last_error": None, + "attributes": {}, +} +_SEVENTEEN_ATTRIBUTES: Final[dict[str, JsonValue]] = {f"key_{index:02d}": "v" for index in range(17)} + + +def _openai_error(message: str, param: str) -> JsonResponse: + return JsonResponse( + content_type="application/json", + status=400, + body={"error": {"message": message, "type": "invalid_request_error", "param": param, "code": None}}, + ) + + +class _Observations: + def __init__(self, url: str) -> None: + self.url = url.rstrip("/") + self.items: tuple[dict[str, JsonValue], ...] = () + + def read(self) -> tuple[dict[str, JsonValue], ...]: + with httpx.Client(timeout=10, trust_env=False) as client: + payload: Final = JSON_OBJECT.validate_python( + client.get(f"{self.url}/__observations?include_method=true").json() + ) + requests: Final = payload.get("requests") + assert isinstance(requests, list) + self.items = (*self.items, *(object_value(item) for item in requests if isinstance(item, dict))) + return self.items + + def provider_calls(self, identity: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + item + for item in self.items + if f"/{identity}/" in str(item.get("path")) + and not (item.get("method") == "GET" and str(item.get("path", "")).endswith(("/v1/models", "/models"))) + ) + + +def _scripted(scenario: Scenario, prefix: str, response: JsonResponse | BinaryResponse) -> tuple[str, ScenarioHandle]: + identity: Final = f"{prefix}-{uuid.uuid4().hex}" + handle: Final = register_scenario(identity, response) + scenario.cleanups.callback(delete_scenario, handle) + return identity, handle + + +def _outbound(gateway: Gateway, identity: str) -> dict[str, JsonValue]: + observations: Final = _Observations(gateway.upstream_url) + eventually(observations.read, lambda _items: len(observations.provider_calls(identity)) == 1, seconds=20) + return object_value(observations.provider_calls(identity)[0]["body"]) + + +def _assert_no_provider_call(gateway: Gateway, identity: str) -> None: + observations: Final = _Observations(gateway.upstream_url) + observations.read() + assert observations.provider_calls(identity) == () + + +def _assert_no_call_to_deployment(gateway: Gateway, model_name: str) -> None: + observations: Final = _Observations(gateway.upstream_url) + observations.read() + calls: Final = tuple( + f"{item.get('method')} {item.get('path')}" + for item in observations.items + if f"/{model_name}-" in str(item.get("path")) and not str(item.get("path")).endswith("/models") + ) + assert calls == (), calls + + +def _post(gateway: Gateway, path: str, body: dict[str, JsonValue]) -> httpx.Response: + return gateway.client.post(path, json=body, headers={"Authorization": f"Bearer {gateway.key}"}) + + +def _invalid_request(response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 400, response.text + error: Final = object_value(JSON_OBJECT.validate_python(response.json())["error"]) + assert error.get("type") == "invalid_request_error", response.text + return error + + +def _write_config(path: Path, config: dict[str, JsonValue]) -> Path: + path.write_text(json.dumps(config), encoding="utf-8") + return path + + +@pytest.fixture(scope="module") +def router_default_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + """A proxy whose router-wide defaults name positional and route-required params, + with client-side `user_config` allowed so a null value can be sent.""" + directory: Final = tmp_path_factory.mktemp("router-defaults") + with gateway_from_environment() as gateway, gateway.scenario() as scenario: + image_identity, image_handle = _scripted( + scenario, "followup-image", JsonResponse(content_type="application/json", body={"created": 1, "data": []}) + ) + text_identity, text_handle = _scripted( + scenario, "followup-text", JsonResponse(content_type="application/json", body={"choices": []}) + ) + speech_identity, speech_handle = _scripted( + scenario, "followup-speech", BinaryResponse(content_type="audio/mpeg", length=16) + ) + message_identity, message_handle = _scripted( + scenario, "followup-message", JsonResponse(content_type="application/json", body=_MESSAGE) + ) + config: Final = _write_config( + directory / "router-defaults.yaml", + { + "model_list": [ + { + "model_name": "followup-image", + "litellm_params": { + "model": "openai/gpt-image-1", + "api_base": image_handle.api_base(), + "api_key": image_identity, + }, + }, + { + "model_name": "followup-text", + "litellm_params": { + "model": "openai/gpt-3.5-turbo-instruct", + "api_base": text_handle.api_base(), + "api_key": text_identity, + }, + }, + { + "model_name": "followup-speech", + "litellm_params": { + "model": "openai/gpt-4o-mini-tts", + "api_base": speech_handle.api_base(), + "api_key": speech_identity, + }, + }, + { + "model_name": "followup-message", + "litellm_params": { + "model": "anthropic/claude-haiku-4-5", + "api_base": message_handle.api_base(), + "api_key": message_identity, + }, + }, + ], + "router_settings": {"default_litellm_params": _ROUTER_DEFAULTS}, + "general_settings": {"allow_client_side_credentials": True}, + }, + ) + with owned_proxy_process(gateway, directory, {}, config=config) as owned: + yield Gateway(owned.gateway.client, owned.gateway.key, gateway.upstream_url) + + +def test_ocr_multipart_without_file_returns_400(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + identity, handle = _scripted( + scenario, "followup-ocr", JsonResponse(content_type="application/json", body={"pages": []}) + ) + model: Final = scenario.model(model="mistral/mistral-ocr-latest", api_base=handle.api_base(), api_key=identity) + response: Final = gateway.client.post( + "/v1/ocr", files={"model": (None, model)}, headers={"Authorization": f"Bearer {gateway.key}"} + ) + error: Final = _invalid_request(response) + assert error.get("param") == "file", response.text + _assert_no_provider_call(gateway, identity) + + +def test_ocr_empty_file_returns_400(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + identity, handle = _scripted( + scenario, "followup-ocr-empty", JsonResponse(content_type="application/json", body={"pages": []}) + ) + model: Final = scenario.model(model="mistral/mistral-ocr-latest", api_base=handle.api_base(), api_key=identity) + response: Final = gateway.request_multipart( + "/v1/ocr", {"model": model}, {"file": ("empty.pdf", b"", "application/pdf")} + ) + error: Final = _invalid_request(response) + assert error.get("message") == "Uploaded file is empty", response.text + _assert_no_provider_call(gateway, identity) + + +@pytest.mark.parametrize("max_output_tokens", (-1, 0), ids=("negative", "zero")) +def test_responses_non_positive_max_output_tokens_returns_provider_400( + gateway: Gateway, max_output_tokens: int +) -> None: + message: Final = f"Invalid 'max_output_tokens': integer below minimum value. Expected a value >= 16, but got {max_output_tokens} instead." + with gateway.scenario() as scenario: + identity, handle = _scripted(scenario, "followup-max-output", _openai_error(message, "max_output_tokens")) + model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=handle.api_base(), api_key=identity) + response: Final = _post( + gateway, + "/v1/responses", + {"model": model, "input": "Say pong", "store": False, "max_output_tokens": max_output_tokens}, + ) + assert response.status_code == 400, response.text + assert message in response.text, response.text + assert _outbound(gateway, identity).get("max_output_tokens") == max_output_tokens + + +def test_responses_small_positive_max_output_tokens_is_still_raised_to_minimum(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + identity, handle = _scripted( + scenario, "followup-max-output-min", JsonResponse(content_type="application/json", body=_RESPONSE) + ) + model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=handle.api_base(), api_key=identity) + response: Final = _post( + gateway, "/v1/responses", {"model": model, "input": "Say pong", "store": False, "max_output_tokens": 1} + ) + assert response.status_code == 200, response.text + assert _outbound(gateway, identity).get("max_output_tokens") == 16 + + +def test_team_daily_activity_end_before_start_returns_400(gateway: Gateway) -> None: + response: Final = gateway.request( + "GET", + "/team/daily/activity", + params={"start_date": "2026-10-01", "end_date": "2026-09-24", "page": "1"}, + ) + assert response.status_code == 400, response.text + assert response.json() == {"detail": {"error": "end_date must be on or after start_date"}}, response.text + + +def _registered_store(gateway: Gateway, scenario: Scenario, identity: str, handle: ScenarioHandle) -> str: + store: Final = f"vs-{uuid.uuid4().hex}" + model: Final = scenario.model(model="openai/text-embedding-3-small", api_base=handle.api_base(), api_key=identity) + gateway.post( + "/vector_store/new", + { + "vector_store_id": store, + "custom_llm_provider": "openai", + "litellm_params": {"model": model, "api_base": handle.api_base(), "api_key": identity}, + }, + ) + scenario.cleanups.callback(gateway.post, "/vector_store/delete", {"vector_store_id": store}) + return store + + +@pytest.mark.parametrize( + ("path", "body"), + ( + ("/v1/vector_stores/{store}/files/file-audit", {"attributes": _SEVENTEEN_ATTRIBUTES}), + ("/v1/vector_stores/{store}/files", {"file_id": "file-audit", "attributes": _SEVENTEEN_ATTRIBUTES}), + ), + ids=("update", "create"), +) +def test_vector_store_file_attributes_over_provider_limit_reach_provider( + gateway: Gateway, path: str, body: dict[str, JsonValue] +) -> None: + message: Final = ( + "Invalid 'attributes': too many properties. " + "Expected an object with at most 16 properties, but got an object with 17 properties instead." + ) + with gateway.scenario() as scenario: + identity, handle = _scripted(scenario, "followup-attributes", _openai_error(message, "attributes")) + store: Final = _registered_store(gateway, scenario, identity, handle) + response: Final = _post(gateway, path.format(store=store), body) + assert response.status_code == 400, response.text + assert message in response.text, response.text + assert _outbound(gateway, identity).get("attributes") == _SEVENTEEN_ATTRIBUTES + + +@pytest.mark.parametrize( + ("path", "body", "missing"), + ( + ("/v1/images/generations", {"model": "followup-image"}, "prompt"), + ("/v1/completions", {"model": "followup-text"}, "prompt"), + ("/v1/audio/speech", {"model": "followup-speech", "voice": "alloy"}, "input"), + ), + ids=("image-prompt", "text-completion-prompt", "speech-input"), +) +def test_router_default_for_positional_param_returns_400( + router_default_proxy: Gateway, path: str, body: dict[str, JsonValue], missing: str +) -> None: + response: Final = _post(router_default_proxy, path, body) + error: Final = _invalid_request(response) + assert error.get("param") == missing, response.text + _assert_no_call_to_deployment(router_default_proxy, str(body["model"])) + + +def test_router_default_for_model_outside_router_returns_400(router_default_proxy: Gateway) -> None: + response: Final = _post(router_default_proxy, "/v1/moderations", {"model": "omni-moderation-2024-09-26"}) + error: Final = _invalid_request(response) + assert error.get("param") == "input", response.text + + +@pytest.mark.parametrize( + ("body", "expected_max_tokens"), + ( + ({"messages": [{"role": "user", "content": "router default"}]}, 32), + ({"messages": [{"role": "user", "content": "explicit"}], "max_tokens": 8}, 8), + ), + ids=("router-default", "explicit"), +) +def test_null_user_config_is_treated_as_absent( + router_default_proxy: Gateway, body: dict[str, JsonValue], expected_max_tokens: int +) -> None: + observations: Final = _Observations(router_default_proxy.upstream_url) + observations.read() + before: Final = len(observations.items) + response: Final = _post( + router_default_proxy, "/v1/messages", {"model": "followup-message", "user_config": None, **body} + ) + assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_python(response.json()).get("type") == "message", response.text + eventually( + observations.read, + lambda items: any("followup-message-" in str(item.get("path")) for item in items[before:]), + seconds=20, + ) + outbound: Final = next( + object_value(item["body"]) + for item in observations.items[before:] + if "followup-message-" in str(item.get("path")) + ) + assert outbound.get("max_tokens") == expected_max_tokens diff --git a/tests/integration/compatibility/test_status_code_followups_chaos.py b/tests/integration/compatibility/test_status_code_followups_chaos.py new file mode 100644 index 00000000000..1cb7f45e681 --- /dev/null +++ b/tests/integration/compatibility/test_status_code_followups_chaos.py @@ -0,0 +1,283 @@ +"""The provider goes down in the middle of a concurrent burst across /v1/chat/completions, /v1/messages and +/v1/responses, streaming and not, on a two-worker proxy: every failed call reaches the caller as an error status, +another key on another deployment keeps working, and once the provider is back every call serves and lands in the +spend log exactly once.""" + +from __future__ import annotations + +import json +import os +import re +import signal +import uuid +from collections.abc import Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.database import read_rows +from integration._support.process import UpstreamSlot, owned_proxy_process, owned_upstream +from integration._support.upstream import register_scenario +from integration.compatibility._status_code_audit import ( + CHAT, + CHAT_FRAMES, + ENDPOINT_BODIES, + ENDPOINT_PATHS, + MESSAGE, + MESSAGE_FRAMES, + RESPONSE, + RESPONSES_FRAMES, + Upstream, + json_response, + sse_response, + stream_finished, +) +from pydantic import JsonValue + +from tests.integration.cost_calculation.cost_tracking_case import StoredResponse + +pytestmark: Final = pytest.mark.timeout(480) + +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_MARKER_HEADER: Final = "x-litellm-spend-logs-metadata" +_KINDS: Final[dict[str, tuple[str, bool, str, str, StoredResponse]]] = { + "chat": ("chat", False, "chaos-chat", "openai/gpt-4o-mini", json_response(CHAT)), + "chatstream": ("chat", True, "chaos-chat-stream", "openai/gpt-4o-mini", sse_response(CHAT_FRAMES)), + "responses": ("responses", False, "chaos-responses", "openai/gpt-4o-mini", json_response(RESPONSE)), + "responsesstream": ( + "responses", + True, + "chaos-responses-stream", + "openai/gpt-4o-mini", + sse_response(RESPONSES_FRAMES), + ), + "messages": ("messages", False, "chaos-message", "anthropic/claude-haiku-4-5", json_response(MESSAGE)), + "messagesstream": ( + "messages", + True, + "chaos-message-stream", + "anthropic/claude-haiku-4-5", + sse_response(MESSAGE_FRAMES), + ), +} +_BURST: Final = 30 + + +@dataclass(frozen=True, slots=True) +class _Chaos: + gateway: Gateway + log: Path + slot: UpstreamSlot + identities: Mapping[str, str] + stable_identity: str + burst_key: str + other_key: str + + +@dataclass(frozen=True, slots=True) +class _Outcome: + marker: str + kind: str + status: int | None + text: str + finished: bool = False + + +def _register_on_slot(slot: UpstreamSlot, identities: Mapping[str, str]) -> None: + for kind, identity in identities.items(): + register_scenario(identity, _KINDS[kind][4], control_url=slot.url) + + +@pytest.fixture(scope="module") +def chaos(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Chaos]: + directory: Final = tmp_path_factory.mktemp("audit-chaos") + with ( + owned_upstream(directory) as slot, + gateway_from_environment() as gateway, + gateway.scenario() as scenario, + ): + identities: Final = {kind: f"chaos-{kind}-{uuid.uuid4().hex}" for kind in _KINDS} + _register_on_slot(slot, identities) + stable_identity: Final = f"chaos-stable-{uuid.uuid4().hex}" + stable_handle: Final = register_scenario(stable_identity, json_response(CHAT)) + model_list: Final[list[JsonValue]] = [ + *( + { + "model_name": _KINDS[kind][2], + "litellm_params": { + "model": _KINDS[kind][3], + "api_base": f"{slot.url}/{identity}", + "api_key": identity, + }, + } + for kind, identity in identities.items() + ), + { + "model_name": "chaos-stable", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": stable_handle.api_base(), + "api_key": stable_identity, + }, + }, + ] + config: Final = directory / "chaos.yaml" + config.write_text( + json.dumps({"model_list": model_list, "router_settings": {"num_retries": 0, "disable_cooldowns": True}}), + encoding="utf-8", + ) + burst_key: Final = scenario.key(key_alias=f"chaos-burst-{uuid.uuid4().hex}") + other_key: Final = scenario.key(key_alias=f"chaos-other-{uuid.uuid4().hex}", models=["chaos-stable"]) + with owned_proxy_process(gateway, directory, {}, config=config, workers=2) as owned: + eventually(lambda: _worker_pids(owned.log), lambda pids: len(pids) == 2, seconds=120) + yield _Chaos( + Gateway(owned.gateway.client, owned.gateway.key, gateway.upstream_url), + owned.log, + slot, + identities, + stable_identity, + burst_key, + other_key, + ) + + +def _worker_pids(log: Path) -> tuple[int, ...]: + return tuple(int(found[1]) for found in _STARTED_WORKER.finditer(log.read_text())) + + +def _call(chaos: _Chaos, kind: str, marker: str, *, key: str | None = None) -> _Outcome: + endpoint, stream, model_name, _target, _response = _KINDS[kind] + body: Final[dict[str, JsonValue]] = { + **ENDPOINT_BODIES[endpoint], + **({"stream": True} if stream else {}), + "model": model_name, + **({"input": marker} if endpoint == "responses" else {"messages": [{"role": "user", "content": marker}]}), + } + headers: Final = { + "Authorization": f"Bearer {chaos.burst_key if key is None else key}", + _MARKER_HEADER: json.dumps({"marker": marker}), + } + try: + with httpx.Client(base_url=str(chaos.gateway.client.base_url), trust_env=False, timeout=60) as client: + with client.stream("POST", ENDPOINT_PATHS[endpoint], json=body, headers=headers) as response: + lines: Final = tuple(response.iter_lines()) + status: Final = response.status_code + except httpx.TransportError as error: + return _Outcome(marker, kind, None, repr(error)) + return _Outcome(marker, kind, status, "\n".join(lines), not stream or stream_finished(endpoint, lines)) + + +def _burst(chaos: _Chaos, phase: str) -> tuple[_Outcome, ...]: + kinds: Final = tuple(_KINDS) + plan: Final = tuple( + (kinds[index % len(kinds)], f"{phase}-{kinds[index % len(kinds)]}-{index}-{uuid.uuid4().hex}") + for index in range(_BURST) + ) + with ThreadPoolExecutor(max_workers=12) as pool: + return tuple(pool.map(lambda item: _call(chaos, item[0], item[1]), plan)) + + +def _upstream_hits(slot: UpstreamSlot, markers: frozenset[str]) -> Mapping[str, int]: + upstream: Final = Upstream(slot.url) + upstream.drain() + bodies: Final = tuple(json.dumps(item.get("body")) for item in upstream.items) + return {marker: sum(marker in body for body in bodies) for marker in markers} + + +def _spend_markers(key: str, markers: frozenset[str]) -> Mapping[str, tuple[str, ...]]: + rows: Final = read_rows( + "SELECT status, metadata->'spend_logs_metadata'->>'marker' AS marker " + 'FROM "LiteLLM_SpendLogs" WHERE api_key = %s', + (sha256(key.encode()).hexdigest(),), + ) + return { + marker: found + for marker in markers + if (found := tuple(str(row["status"]) for row in rows if str(row["marker"]) == marker)) + } + + +def _assert_served(outcomes: tuple[_Outcome, ...]) -> None: + assert [(outcome.kind, outcome.status, outcome.finished) for outcome in outcomes] == [ + (outcome.kind, 200, True) for outcome in outcomes + ], [(outcome.kind, outcome.status, outcome.text[:300]) for outcome in outcomes if outcome.status != 200] + + +def _stable_call(chaos: _Chaos) -> httpx.Response: + return chaos.gateway.client.post( + "/v1/chat/completions", + json={"model": "chaos-stable", "messages": [{"role": "user", "content": "other key"}]}, + headers={"Authorization": f"Bearer {chaos.other_key}"}, + timeout=30, + ) + + +def _settled(chaos: _Chaos) -> bool: + return all(_call(chaos, kind, f"settle-{kind}-{uuid.uuid4().hex}").status == 200 for kind in _KINDS) + + +def _stop_mid_burst(chaos: _Chaos) -> tuple[_Outcome, ...]: + with ThreadPoolExecutor(max_workers=1) as stopper: + stopping: Final = stopper.submit(chaos.slot.stop) + during: Final = _burst(chaos, "during") + stopping.result(timeout=120) + return during + + +def _restart(chaos: _Chaos) -> None: + chaos.slot.start() + _register_on_slot(chaos.slot, chaos.identities) + eventually(lambda: _settled(chaos), lambda ok: ok, seconds=90) + + +def test_provider_outage_mid_burst_errors_reach_callers_and_recovery_serves_once(chaos: _Chaos) -> None: + before: Final = _burst(chaos, "before") + _assert_served(before) + + during: Final = _stop_mid_burst(chaos) + outage: Final = _burst(chaos, "outage") + for outcome in outage: + assert outcome.status is not None and outcome.status >= 500, (outcome.kind, outcome.status, outcome.text) + assert '"error"' in outcome.text, (outcome.kind, outcome.text) + for outcome in during: + assert (outcome.status == 200 and outcome.finished) or ( + outcome.status is not None and outcome.status >= 500 and '"error"' in outcome.text + ), (outcome.kind, outcome.status, outcome.text[:300]) + assert chaos.gateway.request("GET", "/health/liveliness").status_code == 200 + stable: Final = _stable_call(chaos) + assert stable.status_code == 200, stable.text + + _restart(chaos) + after: Final = _burst(chaos, "after") + _assert_served(after) + after_markers: Final = frozenset(outcome.marker for outcome in after) + assert _upstream_hits(chaos.slot, after_markers) == dict.fromkeys(after_markers, 1) + + served: Final = frozenset(outcome.marker for outcome in (*before, *after)) + landed: Final = eventually( + lambda: _spend_markers(chaos.burst_key, served), lambda found: frozenset(found) == served, seconds=120 + ) + assert landed == dict.fromkeys(served, ("success",)), landed + assert len(_worker_pids(chaos.log)) >= 2 + + +def _tolerant_status(chaos: _Chaos, kind: str) -> int | None: + return _call(chaos, kind, f"kill-probe-{kind}-{uuid.uuid4().hex}").status + + +def test_killed_worker_leaves_the_sibling_serving_and_the_respawn_serving(chaos: _Chaos) -> None: + workers: Final = _worker_pids(chaos.log) + assert len(workers) >= 2, workers + os.kill(workers[-1], signal.SIGKILL) + eventually(lambda: _tolerant_status(chaos, "chat"), lambda status: status == 200, seconds=60) + _assert_served(_burst(chaos, "killed")) + stable: Final = _stable_call(chaos) + assert stable.status_code == 200, stable.text + eventually(lambda: _worker_pids(chaos.log), lambda pids: len(pids) == len(workers) + 1, seconds=180) + eventually(lambda: _settled(chaos), lambda ok: ok, seconds=120) + _assert_served(_burst(chaos, "respawned")) diff --git a/tests/integration/compatibility/test_status_code_followups_defaults.py b/tests/integration/compatibility/test_status_code_followups_defaults.py new file mode 100644 index 00000000000..a0c73d12dec --- /dev/null +++ b/tests/integration/compatibility/test_status_code_followups_defaults.py @@ -0,0 +1,391 @@ +"""Router-wide, deployment and `user_config` defaults against the params the Router binds positionally, the +moderation merge, unknown and A2A models, and every `user_config` shape on the three LLM endpoints.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from typing import Final + +import anthropic +import httpx +import openai +import pytest +from integration._support.client import Gateway, gateway_from_environment +from integration.compatibility._status_code_audit import ( + UNSET, + CHAT, + CHAT_FRAMES, + ENDPOINT_BODIES, + ENDPOINT_PATHS, + MESSAGE, + MESSAGE_FRAMES, + MODERATION, + RESPONSE, + RESPONSES_FRAMES, + ROUTER_DEFAULTS, + USER_MESSAGES, + assembled_text, + assert_no_provider_call, + drain_rig_upstream, + deployment, + error_body, + invalid_request, + json_response, + one_outbound, + owned_gateway, + post, + sse_response, + stream_finished, + stream_lines, +) +from pydantic import JsonValue + +from tests.integration.cost_calculation.cost_tracking_case import BinaryResponse, StoredResponse + +pytestmark: Final = pytest.mark.timeout(180) + + +@pytest.fixture(autouse=True) +def _drained_upstream() -> None: + drain_rig_upstream() + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + identities: Mapping[str, str] + deployments: Mapping[str, dict[str, JsonValue]] + + +_SPECS: Final[tuple[tuple[str, str, StoredResponse, dict[str, JsonValue]], ...]] = ( + ("audit-image", "openai/gpt-image-1", json_response({"created": 1, "data": []}), {}), + ("audit-image-dep", "openai/gpt-image-1", json_response({"created": 1, "data": []}), {"prompt": "deployment"}), + ("audit-text", "openai/gpt-3.5-turbo-instruct", json_response({"choices": []}), {}), + ("audit-text-dep", "openai/gpt-3.5-turbo-instruct", json_response({"choices": []}), {"prompt": "deployment"}), + ("audit-speech", "openai/gpt-4o-mini-tts", BinaryResponse(content_type="audio/mpeg", length=16), {}), + ( + "audit-speech-dep", + "openai/gpt-4o-mini-tts", + BinaryResponse(content_type="audio/mpeg", length=16), + {"input": "deployment"}, + ), + ("audit-transcribe", "openai/whisper-1", json_response({"text": "scripted"}), {}), + ("audit-moderation", "openai/omni-moderation-latest", json_response(MODERATION), {}), + ("audit-chat", "openai/gpt-4o-mini", json_response(CHAT), {}), + ("audit-chat-stream", "openai/gpt-4o-mini", sse_response(CHAT_FRAMES), {}), + ("audit-responses", "openai/gpt-4o-mini", json_response(RESPONSE), {}), + ("audit-responses-stream", "openai/gpt-4o-mini", sse_response(RESPONSES_FRAMES), {}), + ("audit-message", "anthropic/claude-haiku-4-5", json_response(MESSAGE), {}), + ("audit-message-stream", "anthropic/claude-haiku-4-5", sse_response(MESSAGE_FRAMES), {}), +) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + """A proxy whose router-wide defaults name every positional and route-required param, with client-side + `user_config` allowed, and one deployment per route (plus `-dep` copies that carry the param themselves).""" + with gateway_from_environment() as gateway, gateway.scenario() as scenario: + built: Final = tuple( + deployment(scenario, name, model, response, **extra) for name, model, response, extra in _SPECS + ) + identities: Final = {name: identity for (name, *_), (identity, _) in zip(_SPECS, built, strict=True)} + deployments: Final = {name: entry for (name, *_), (_, entry) in zip(_SPECS, built, strict=True)} + config: Final[dict[str, JsonValue]] = { + "model_list": [entry for _, entry in built], + "router_settings": {"default_litellm_params": ROUTER_DEFAULTS, "num_retries": 0}, + "general_settings": {"allow_client_side_credentials": True}, + } + with owned_gateway(tmp_path_factory.mktemp("audit-defaults"), config) as owned: + yield _Rig(owned, identities, deployments) + + +_POSITIONAL: Final[dict[str, tuple[str, str, str, str, dict[str, JsonValue]]]] = { + "image": ("/v1/images/generations", "/image/generations", "audit-image", "prompt", {}), + "text-completion": ("/v1/completions", "/completions", "audit-text", "prompt", {}), + "speech": ("/v1/audio/speech", "/audio/speech", "audit-speech", "input", {"voice": "alloy"}), +} + + +def _user_config(rig: _Rig, name: str, defaults: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {"model_list": [rig.deployments[name]], "default_litellm_params": dict(defaults)} + + +@pytest.mark.parametrize("route", tuple(_POSITIONAL)) +@pytest.mark.parametrize("source", ("deployment-param", "user-config-default", "user-config-null")) +def test_positional_param_from_any_default_source_returns_400(rig: _Rig, route: str, source: str) -> None: + path, error_route, name, param, extra = _POSITIONAL[route] + user_config: Final[dict[str, JsonValue]] = ( + {"user_config": _user_config(rig, name, {param: "user config default"})} + if source == "user-config-default" + else {"user_config": None} + if source == "user-config-null" + else {} + ) + body: Final[dict[str, JsonValue]] = { + "model": f"{name}-dep" if source == "deployment-param" else name, + **extra, + **user_config, + } + response: Final = post(rig.gateway, path, body) + error: Final = invalid_request(response) + assert error.get("param") == param, response.text + assert error.get("message") == f"{error_route}: Missing required parameter: '{param}'.", response.text + assert_no_provider_call(rig.gateway, rig.identities[name], rig.identities[f"{name}-dep"]) + + +@pytest.mark.parametrize("route", tuple(_POSITIONAL)) +def test_positional_param_in_the_body_reaches_upstream(rig: _Rig, route: str) -> None: + pytest.skip( + "BUG: a router-wide default_litellm_params value replaces the caller's own positional param " + "(prompt/input) on its way to the provider; same on the merge base" + ) + path, _error_route, name, param, extra = _POSITIONAL[route] + response: Final = post(rig.gateway, path, {"model": name, param: f"explicit {route}", **extra}) + assert response.status_code == 200, response.text + assert one_outbound(rig.gateway, rig.identities[name]).get(param) == f"explicit {route}" + + +def test_transcription_without_file_is_rejected_by_the_form_before_routing(rig: _Rig) -> None: + response: Final = rig.gateway.client.post( + "/v1/audio/transcriptions", + data={"model": "audit-transcribe"}, + files={"note": ("note.txt", b"not audio", "text/plain")}, + headers={"Authorization": f"Bearer {rig.gateway.key}"}, + ) + assert response.status_code == 422, response.text + detail: Final = response.json()["detail"] + assert [entry["loc"] for entry in detail] == [["body", "file"]], response.text + assert_no_provider_call(rig.gateway, rig.identities["audit-transcribe"]) + + +def test_transcription_with_file_reaches_upstream(rig: _Rig) -> None: + response: Final = rig.gateway.client.post( + "/v1/audio/transcriptions", + data={"model": "audit-transcribe"}, + files={"file": ("clip.wav", b"RIFF0000WAVEfmt ", "audio/wav")}, + headers={"Authorization": f"Bearer {rig.gateway.key}"}, + ) + assert response.status_code == 200, response.text + assert response.json()["text"] == "scripted" + outbound: Final = one_outbound(rig.gateway, rig.identities["audit-transcribe"]) + assert outbound.get("model") == "whisper-1" + assert outbound.get("file") == {"filename": "clip.wav", "content_type": "audio/x-wav"} + + +@pytest.mark.parametrize( + ("body", "expected_input"), + (({}, "router default input"), ({"input": "explicit input"}, "explicit input")), + ids=("router-default", "explicit"), +) +def test_listed_moderation_model_merges_the_router_default( + rig: _Rig, body: dict[str, JsonValue], expected_input: str +) -> None: + response: Final = post(rig.gateway, "/v1/moderations", {"model": "audit-moderation", **body}) + assert response.status_code == 200, response.text + assert one_outbound(rig.gateway, rig.identities["audit-moderation"]).get("input") == expected_input + + +@pytest.mark.parametrize( + ("path", "body", "message"), + ( + pytest.param( + "/v1/images/generations", + {"model": "audit-unknown-model"}, + "/image/generations: Missing required parameter: 'prompt'.", + id="image-unknown-model", + ), + pytest.param( + "/v1/images/generations", + {"model": "a2a/audit-unknown-agent"}, + "/image/generations: Missing required parameter: 'prompt'.", + id="image-a2a-model", + ), + pytest.param( + "/v1/messages", + {"model": "a2a/audit-unknown-agent", "messages": USER_MESSAGES}, + "anthropic_messages: Invalid model name passed in model=a2a/audit-unknown-agent. " + "Call `/v1/models` to view available models for your key.", + id="messages-a2a-model", + ), + ), +) +def test_model_outside_the_router_with_a_router_default_param( + rig: _Rig, path: str, body: dict[str, JsonValue], message: str +) -> None: + response: Final = post(rig.gateway, path, body) + assert response.status_code == 400, response.text + error: Final = error_body(response) + assert error.get("type") == "invalid_request_error", response.text + assert error.get("message") == message, response.text + + +@pytest.mark.parametrize( + ("user_config", "message"), + ( + pytest.param(None, None, id="null"), + pytest.param(UNSET, None, id="missing"), + pytest.param({}, "There are no healthy deployments for this model", id="empty-object"), + pytest.param("", "Invalid request format: 'str' object has no attribute 'items'", id="empty-string"), + pytest.param(7, "Invalid request format: 'int' object has no attribute 'items'", id="int"), + pytest.param([], "Invalid request format: 'list' object has no attribute 'items'", id="list"), + pytest.param("x" * 5000, "Invalid request format: 'str' object has no attribute 'items'", id="5kb-string"), + ), +) +def test_chat_user_config_shapes(rig: _Rig, user_config: object, message: str | None) -> None: + body: Final[dict[str, JsonValue]] = { + "model": "audit-chat", + "messages": USER_MESSAGES, + **({} if user_config is UNSET else {"user_config": user_config}), # pyright: ignore[reportAssignmentType] # the sad shapes are deliberately not JSON objects + } + response: Final = post(rig.gateway, "/v1/chat/completions", body) + if message is None: + assert response.status_code == 200, response.text + assert one_outbound(rig.gateway, rig.identities["audit-chat"]).get("max_tokens") == 32 + return + assert response.status_code == 400, response.text + assert message in str(error_body(response).get("message")), response.text + assert_no_provider_call(rig.gateway, rig.identities["audit-chat"]) + + +def test_messages_empty_user_config_does_not_inherit_router_defaults(rig: _Rig) -> None: + response: Final = post( + rig.gateway, "/v1/messages", {"model": "audit-message", "messages": USER_MESSAGES, "user_config": {}} + ) + assert response.status_code == 400, response.text + assert error_body(response).get("message") == "anthropic_messages: Missing required parameter: 'max_tokens'." + assert_no_provider_call(rig.gateway, rig.identities["audit-message"]) + + +@pytest.mark.parametrize( + ("endpoint", "model", "stream", "expected_max_tokens"), + ( + pytest.param("chat", "audit-chat", False, 32, id="chat"), + pytest.param("chat", "audit-chat-stream", True, 32, id="chat-stream"), + pytest.param("responses", "audit-responses", False, None, id="responses"), + pytest.param("responses", "audit-responses-stream", True, None, id="responses-stream"), + pytest.param("messages", "audit-message-stream", True, 32, id="messages-stream"), + ), +) +def test_null_user_config_serves_every_endpoint( + rig: _Rig, endpoint: str, model: str, stream: bool, expected_max_tokens: int | None +) -> None: + body: Final[dict[str, JsonValue]] = { + **{key: value for key, value in ENDPOINT_BODIES[endpoint].items() if key != "max_tokens"}, + "model": model, + "user_config": None, + **({"stream": True} if stream else {}), + } + if stream: + status, lines = stream_lines(rig.gateway, ENDPOINT_PATHS[endpoint], body) + assert status == 200, lines + assert stream_finished(endpoint, lines), lines + assert assembled_text(lines) == "streamed response", lines + else: + response: Final = post(rig.gateway, ENDPOINT_PATHS[endpoint], body) + assert response.status_code == 200, response.text + assert "scripted" in response.text, response.text + outbound: Final = one_outbound(rig.gateway, rig.identities[model]) + assert "user_config" not in outbound, outbound + assert outbound.get("max_tokens") == expected_max_tokens, outbound + + +def test_null_user_config_without_a_valid_key_is_401(rig: _Rig) -> None: + response: Final = post( + rig.gateway, + "/v1/chat/completions", + {"model": "audit-chat", "messages": USER_MESSAGES, "user_config": None}, + key="sk-audit-not-a-key", + ) + assert response.status_code == 401, response.text + assert_no_provider_call(rig.gateway, rig.identities["audit-chat"]) + + +def _openai_chat(rig: _Rig, asynchronous: bool) -> str | None: + base_url: Final = f"{rig.gateway.client.base_url}/v1" + if not asynchronous: + with openai.OpenAI(base_url=base_url, api_key=rig.gateway.key, max_retries=0) as client: + completion: Final = client.chat.completions.create( + model="audit-chat", + messages=[{"role": "user", "content": "Hello"}], + extra_body={"user_config": None}, + ) + return completion.choices[0].message.content + + async def call() -> str | None: + async with openai.AsyncOpenAI(base_url=base_url, api_key=rig.gateway.key, max_retries=0) as client: + completion: Final = await client.chat.completions.create( + model="audit-chat", + messages=[{"role": "user", "content": "Hello"}], + extra_body={"user_config": None}, + ) + return completion.choices[0].message.content + + return asyncio.run(call()) + + +def _anthropic_message(rig: _Rig, asynchronous: bool) -> str: + base_url: Final = str(rig.gateway.client.base_url) + if not asynchronous: + with anthropic.Anthropic(base_url=base_url, api_key=rig.gateway.key, max_retries=0) as client: + message: Final = client.messages.create( + model="audit-message", + max_tokens=8, + messages=[{"role": "user", "content": "Hello"}], + extra_body={"user_config": None}, + ) + block: Final = message.content[0] + assert isinstance(block, anthropic.types.TextBlock), message + return block.text + + async def call() -> str: + async with anthropic.AsyncAnthropic(base_url=base_url, api_key=rig.gateway.key, max_retries=0) as client: + message: Final = await client.messages.create( + model="audit-message", + max_tokens=8, + messages=[{"role": "user", "content": "Hello"}], + extra_body={"user_config": None}, + ) + block: Final = message.content[0] + assert isinstance(block, anthropic.types.TextBlock), message + return block.text + + return asyncio.run(call()) + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +def test_null_user_config_through_the_openai_sdk(rig: _Rig, asynchronous: bool) -> None: + assert _openai_chat(rig, asynchronous) == "scripted" + assert one_outbound(rig.gateway, rig.identities["audit-chat"]).get("max_tokens") == 32 + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +def test_null_user_config_through_the_anthropic_sdk(rig: _Rig, asynchronous: bool) -> None: + assert _anthropic_message(rig, asynchronous) == "scripted" + assert one_outbound(rig.gateway, rig.identities["audit-message"]).get("max_tokens") == 8 + + +def test_null_user_config_without_the_client_side_opt_in_is_rejected(gateway: Gateway) -> None: + response: Final = gateway.client.post( + "/v1/chat/completions", + json={"model": "audit-not-routed", "messages": USER_MESSAGES, "user_config": None}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) + assert response.status_code == 401, response.text + assert "user_config is not allowed in request body" in response.text, response.text + + +def test_raw_httpx_async_null_user_config(rig: _Rig) -> None: + async def call() -> httpx.Response: + async with httpx.AsyncClient(base_url=str(rig.gateway.client.base_url), trust_env=False) as client: + return await client.post( + "/v1/responses", + json={"model": "audit-responses", "input": "Hello", "store": False, "user_config": None}, + headers={"Authorization": f"Bearer {rig.gateway.key}"}, + timeout=30, + ) + + response: Final = asyncio.run(call()) + assert response.status_code == 200, response.text + assert one_outbound(rig.gateway, rig.identities["audit-responses"]).get("input") == "Hello" diff --git a/tests/integration/compatibility/test_status_code_followups_requests.py b/tests/integration/compatibility/test_status_code_followups_requests.py new file mode 100644 index 00000000000..cffaf09ec55 --- /dev/null +++ b/tests/integration/compatibility/test_status_code_followups_requests.py @@ -0,0 +1,425 @@ +"""OCR parse errors, Responses `max_output_tokens`, vector store file attributes and the daily activity date range, +driven through the shared rig proxy with the provider body read back from the scripted upstream.""" + +from __future__ import annotations + +import asyncio +import itertools +import uuid +from collections.abc import Callable +from typing import Final + +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario +from integration._support.upstream import ScenarioHandle +from integration.compatibility._status_code_audit import ( + UNSET, + RESPONSE, + RESPONSES_FRAMES, + assembled_text, + assert_no_provider_call, + drain_rig_upstream, + invalid_request, + json_response, + one_outbound, + openai_error, + post, + register, + sse_response, + stream_finished, + stream_lines, +) +from pydantic import JsonValue + +from tests.integration.cost_calculation.cost_tracking_case import StoredResponse + +pytestmark: Final = pytest.mark.timeout(180) + + +@pytest.fixture(autouse=True) +def _drained_upstream() -> None: + drain_rig_upstream() + + +_OCR: Final[dict[str, JsonValue]] = { + "pages": [{"index": 0, "markdown": "scripted page", "images": [], "dimensions": None}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1}, +} +_VECTOR_STORE_FILE: Final[dict[str, JsonValue]] = { + "id": "file-audit", + "object": "vector_store.file", + "created_at": 1, + "usage_bytes": 0, + "vector_store_id": "vs-audit", + "status": "completed", + "last_error": None, + "attributes": {}, +} + + +def _ocr_model(scenario: Scenario) -> tuple[str, str]: + identity, handle = register(scenario, "audit-ocr", json_response(_OCR)) + return scenario.model(model="mistral/mistral-ocr-latest", api_base=handle.api_base(), api_key=identity), identity + + +def _raw(gateway: Gateway, path: str, content: bytes, content_type: str, *, key: str | None = None) -> httpx.Response: + return gateway.client.post( + path, + content=content, + headers={"Authorization": f"Bearer {gateway.key if key is None else key}", "Content-Type": content_type}, + timeout=60, + ) + + +@pytest.mark.parametrize( + ("path", "body", "message", "param"), + ( + pytest.param( + "/v1/ocr", + {"document": {"type": "file", "file": "local.pdf"}}, + "document type 'file' is not supported through the JSON API.", + None, + id="document-type-file", + ), + pytest.param( + "/v1/ocr", + {"document": {"type": "document_url", "document_url": "reducto://file-123"}}, + "reducto:// file IDs are not accepted through the proxy OCR API;", + None, + id="reducto-document-url", + ), + pytest.param( + "/v1/ocr", + {"document": {"type": "image_url", "image_url": "reducto://file-123"}}, + "reducto:// file IDs are not accepted through the proxy OCR API;", + None, + id="reducto-image-url", + ), + pytest.param( + "/ocr", + {"document": {"type": "file", "file": "local.pdf"}}, + "document type 'file' is not supported through the JSON API.", + None, + id="unversioned-route", + ), + ), +) +def test_ocr_unsupported_json_document_returns_400( + gateway: Gateway, path: str, body: dict[str, JsonValue], message: str, param: str | None +) -> None: + with gateway.scenario() as scenario: + model, identity = _ocr_model(scenario) + response: Final = post(gateway, path, {"model": model, **body}) + error: Final = invalid_request(response) + assert str(error.get("message")).startswith(message), response.text + assert error.get("param") == param, response.text + assert_no_provider_call(gateway, identity) + + +@pytest.mark.parametrize( + ("content", "content_type", "message", "param"), + ( + pytest.param(b"", "application/json", "Empty request body.", None, id="empty-json-body"), + pytest.param(b"{not json", "application/json", "Invalid JSON payload", "request_body", id="invalid-json"), + pytest.param(b"hello", "text/plain", "Invalid JSON payload", "request_body", id="text-plain"), + pytest.param( + b"--x\r\n", + "multipart/form-data", + "Invalid form payload: 400: Missing boundary", + "request_body", + id="no-boundary", + ), + ), +) +def test_ocr_unparseable_body_returns_400( + gateway: Gateway, content: bytes, content_type: str, message: str, param: str | None +) -> None: + response: Final = _raw(gateway, "/v1/ocr", content, content_type) + error: Final = invalid_request(response) + assert str(error.get("message")).startswith(message), response.text + assert error.get("param") == param, response.text + + +def test_ocr_json_array_body_returns_400(gateway: Gateway) -> None: + pytest.skip("BUG: a JSON array body answers 500 on every route, from the auth pre-checks; same on base") + response: Final = _raw(gateway, "/v1/ocr", b"[1,2]", "application/json") + assert invalid_request(response).get("param") == "request_body", response.text + + +def test_ocr_file_over_the_size_limit_returns_400(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model, identity = _ocr_model(scenario) + response: Final = gateway.request_multipart( + "/v1/ocr", {"model": model}, {"file": ("big.pdf", b"%PDF" + b"0" * (50 * 1024 * 1024), "application/pdf")} + ) + error: Final = invalid_request(response) + assert error.get("message") == "OCR file exceeds the size limit", response.text + assert error.get("param") == "file", response.text + assert_no_provider_call(gateway, identity) + + +def test_ocr_multipart_without_a_key_is_401(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model, identity = _ocr_model(scenario) + response: Final = gateway.client.post( + "/v1/ocr", files={"model": (None, model)}, headers={"Authorization": "Bearer sk-audit-not-a-key"} + ) + assert response.status_code == 401, response.text + assert_no_provider_call(gateway, identity) + + +def test_ocr_document_url_reaches_the_provider(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model, identity = _ocr_model(scenario) + document: Final[dict[str, JsonValue]] = {"type": "document_url", "document_url": "https://example.com/a.pdf"} + response: Final = post(gateway, "/v1/ocr", {"model": model, "document": document}) + assert response.status_code == 200, response.text + assert response.json()["pages"][0]["markdown"] == "scripted page" + assert one_outbound(gateway, identity).get("document") == document + + +def test_ocr_multipart_file_reaches_the_provider_as_a_data_uri(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model, identity = _ocr_model(scenario) + response: Final = gateway.request_multipart( + "/v1/ocr", {"model": model}, {"file": ("doc.pdf", b"%PDF-1.4 audit", "application/pdf")} + ) + assert response.status_code == 200, response.text + document: Final = one_outbound(gateway, identity).get("document") + assert isinstance(document, dict), document + assert document.get("type") == "document_url", document + assert str(document.get("document_url")).startswith("data:application/pdf;base64,"), document + + +def _responses_model(scenario: Scenario, response: StoredResponse) -> tuple[str, str]: + identity, handle = register(scenario, "audit-max-output", response) + return scenario.model(model="openai/gpt-4o-mini", api_base=handle.api_base(), api_key=identity), identity + + +@pytest.mark.parametrize( + ("sent", "forwarded"), + ( + pytest.param(1, 16, id="1"), + pytest.param(15, 16, id="15"), + pytest.param(16, 16, id="16"), + pytest.param(4096, 4096, id="4096"), + pytest.param(True, 16, id="true"), + pytest.param(None, UNSET, id="null"), + pytest.param(UNSET, UNSET, id="absent"), + pytest.param("8", "8", id="string"), + pytest.param(0.5, 0.5, id="float"), + ), +) +def test_responses_max_output_tokens_forwarding(gateway: Gateway, sent: object, forwarded: object) -> None: + with gateway.scenario() as scenario: + model, identity = _responses_model(scenario, json_response(RESPONSE)) + body: Final[dict[str, JsonValue]] = { + "model": model, + "input": "Say pong", + "store": False, + **({} if sent is UNSET else {"max_output_tokens": sent}), # pyright: ignore[reportAssignmentType] # the parametrized shapes are JSON values + } + response: Final = post(gateway, "/v1/responses", body) + assert response.status_code == 200, response.text + outbound: Final = one_outbound(gateway, identity) + assert outbound.get("max_output_tokens", UNSET) == forwarded, outbound + + +_BELOW_MINIMUM: Final = ( + "Invalid 'max_output_tokens': integer below minimum value. Expected a value >= 16, but got {} instead." +) + + +@pytest.mark.parametrize("sent", (0, -1), ids=("zero", "negative")) +def test_responses_non_positive_max_output_tokens_stream_returns_the_provider_400(gateway: Gateway, sent: int) -> None: + message: Final = _BELOW_MINIMUM.format(sent) + with gateway.scenario() as scenario: + model, identity = _responses_model(scenario, openai_error(message, "max_output_tokens")) + status, lines = stream_lines( + gateway, + "/v1/responses", + {"model": model, "input": "Say pong", "store": False, "stream": True, "max_output_tokens": sent}, + ) + assert status == 400, lines + assert any(message in line for line in lines), lines + outbound: Final = one_outbound(gateway, identity) + assert (outbound.get("max_output_tokens"), outbound.get("stream")) == (sent, True), outbound + + +def test_responses_small_max_output_tokens_stream_is_raised_to_the_minimum(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model, identity = _responses_model(scenario, sse_response(RESPONSES_FRAMES)) + status, lines = stream_lines( + gateway, + "/v1/responses", + {"model": model, "input": "Say pong", "store": False, "stream": True, "max_output_tokens": 1}, + ) + assert status == 200, lines + assert stream_finished("responses", lines), lines + assert assembled_text(lines) == "streamed response", lines + assert one_outbound(gateway, identity).get("max_output_tokens") == 16 + + +def _sdk_responses( + gateway: Gateway, model: str, sent: int, *, asynchronous: bool, stream: bool +) -> Callable[[], object]: + base_url: Final = f"{gateway.client.base_url}/v1" + + def sync_call() -> object: + with openai.OpenAI(base_url=base_url, api_key=gateway.key, max_retries=0) as client: + return client.responses.create( + model=model, input="Say pong", store=False, max_output_tokens=sent, stream=stream + ) + + async def async_call() -> object: + async with openai.AsyncOpenAI(base_url=base_url, api_key=gateway.key, max_retries=0) as client: + return await client.responses.create( + model=model, input="Say pong", store=False, max_output_tokens=sent, stream=stream + ) + + return (lambda: asyncio.run(async_call())) if asynchronous else sync_call + + +@pytest.mark.parametrize( + ("sent", "asynchronous", "stream"), + ((0, False, False), (0, True, False), (-1, False, True), (-1, True, True)), + ids=("zero-sync", "zero-async", "negative-sync-stream", "negative-async-stream"), +) +def test_responses_non_positive_max_output_tokens_through_the_openai_sdk( + gateway: Gateway, sent: int, asynchronous: bool, stream: bool +) -> None: + message: Final = _BELOW_MINIMUM.format(sent) + with gateway.scenario() as scenario: + model, identity = _responses_model(scenario, openai_error(message, "max_output_tokens")) + call: Final = _sdk_responses(gateway, model, sent, asynchronous=asynchronous, stream=stream) + with pytest.raises(openai.BadRequestError) as raised: + call() + assert raised.value.status_code == 400 + assert message in raised.value.response.text + assert one_outbound(gateway, identity).get("max_output_tokens") == sent + + +def _registered_store(gateway: Gateway, scenario: Scenario, identity: str, handle: ScenarioHandle) -> str: + store: Final = f"vs-{uuid.uuid4().hex}" + model: Final = scenario.model(model="openai/text-embedding-3-small", api_base=handle.api_base(), api_key=identity) + gateway.post( + "/vector_store/new", + { + "vector_store_id": store, + "custom_llm_provider": "openai", + "litellm_params": {"model": model, "api_base": handle.api_base(), "api_key": identity}, + }, + ) + scenario.cleanups.callback(gateway.post, "/vector_store/delete", {"vector_store_id": store}) + return store + + +_SIXTEEN: Final[dict[str, JsonValue]] = {f"key_{index:02d}": f"value {index}" for index in range(16)} +_ATTRIBUTE_CASES: Final = ( + pytest.param(_SIXTEEN, _SIXTEEN, id="sixteen"), + pytest.param({}, {}, id="empty"), + pytest.param({"team": "a", "hidden_params": "internal"}, {"team": "a"}, id="hidden-params"), + pytest.param({"long": "x" * 5000}, {"long": "x" * 5000}, id="5kb-value"), +) + + +@pytest.mark.parametrize("operation", ("create", "update")) +@pytest.mark.parametrize(("attributes", "forwarded"), _ATTRIBUTE_CASES) +def test_vector_store_file_attributes_reach_the_provider( + gateway: Gateway, operation: str, attributes: dict[str, JsonValue], forwarded: dict[str, JsonValue] +) -> None: + with gateway.scenario() as scenario: + identity, handle = register(scenario, "audit-attributes", json_response(_VECTOR_STORE_FILE)) + store: Final = _registered_store(gateway, scenario, identity, handle) + path: Final = ( + f"/v1/vector_stores/{store}/files" + if operation == "create" + else f"/v1/vector_stores/{store}/files/file-audit" + ) + body: Final[dict[str, JsonValue]] = { + "attributes": attributes, + **({"file_id": "file-audit"} if operation == "create" else {}), + } + response: Final = post(gateway, path, body) + assert response.status_code == 200, response.text + assert one_outbound(gateway, identity).get("attributes") == forwarded + + +@pytest.mark.parametrize("operation", ("create", "update")) +def test_vector_store_file_number_and_boolean_attributes_reach_the_provider(gateway: Gateway, operation: str) -> None: + pytest.skip( + "BUG: number and boolean vector store file attributes (valid OpenAI attribute values) are dropped before " + "the provider call, so they are silently lost; same on base" + ) + attributes: Final[dict[str, JsonValue]] = {"pages": 3, "public": True, "team": "a"} + with gateway.scenario() as scenario: + identity, handle = register(scenario, "audit-attributes-typed", json_response(_VECTOR_STORE_FILE)) + store: Final = _registered_store(gateway, scenario, identity, handle) + path: Final = ( + f"/v1/vector_stores/{store}/files" + if operation == "create" + else f"/v1/vector_stores/{store}/files/file-audit" + ) + body: Final[dict[str, JsonValue]] = { + "attributes": attributes, + **({"file_id": "file-audit"} if operation == "create" else {}), + } + response: Final = post(gateway, path, body) + assert response.status_code == 200, response.text + assert one_outbound(gateway, identity).get("attributes") == attributes + + +_ENTITIES: Final = ("agent", "customer", "organization", "tag", "team", "user") +_SUFFIXES: Final = ( + "", + "/aggregated", + "/aggregated/keys", + "/aggregated/model_top_keys", + "/aggregated/search", + "/export", +) +_DAILY_ROUTES: Final = tuple( + f"/{entity}/daily/activity{suffix}" + for entity, suffix in itertools.product(_ENTITIES, _SUFFIXES) + if f"/{entity}/daily/activity{suffix}" != "/organization/daily/activity" +) + ("/user/daily/activity/aggregated/cache_leakage_keys",) +_DATE_CASES: Final[dict[str, tuple[str, str, int, str | None]]] = { + "reversed": ("2026-10-01", "2026-09-24", 400, "end_date must be on or after start_date"), + "equal": ("2026-10-01", "2026-10-01", 200, None), + "bad-month": ("2026-13-01", "2026-10-01", 400, "start_date and end_date must be valid YYYY-MM-DD dates"), +} + + +def _route_params(route: str) -> dict[str, str]: + if route.endswith("/model_top_keys"): + return {"model_group": "audit-model"} + if route.endswith("/search"): + return {"search": "audit"} + return {} + + +@pytest.mark.parametrize("case", tuple(_DATE_CASES)) +@pytest.mark.parametrize("route", _DAILY_ROUTES) +def test_daily_activity_date_range(gateway: Gateway, route: str, case: str) -> None: + start, end, status, error = _DATE_CASES[case] + response: Final = gateway.request( + "GET", route, params={"start_date": start, "end_date": end, **_route_params(route)} + ) + assert response.status_code == status, response.text + if error is not None: + assert response.json() == {"detail": {"error": error}}, response.text + + +@pytest.mark.parametrize("case", ("reversed", "bad-month")) +def test_gateway_daily_activity_date_range(gateway: Gateway, case: str) -> None: + pytest.skip( + "BUG: /gateway/daily/activity answers 200 with empty totals for a reversed or invalid date range " + "instead of the 400 every other daily activity route gives; same on base" + ) + start, end, status, error = _DATE_CASES[case] + response: Final = gateway.request("GET", "/gateway/daily/activity", params={"start_date": start, "end_date": end}) + assert response.status_code == status, response.text + assert response.json() == {"detail": {"error": error}}, response.text diff --git a/tests/integration/compatibility/test_status_code_followups_usage_chat.py b/tests/integration/compatibility/test_status_code_followups_usage_chat.py new file mode 100644 index 00000000000..6d9d9507158 --- /dev/null +++ b/tests/integration/compatibility/test_status_code_followups_usage_chat.py @@ -0,0 +1,125 @@ +"""The /usage/ai/chat tools hand the model the reason the daily activity query rejected its dates, the same check +the daily activity routes answer 400 with, instead of a generic fetch error.""" + +from __future__ import annotations + +import json +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import JSON_OBJECT, Gateway, object_value +from integration._support.upstream import delete_scenario, register_scenario +from integration.compatibility._status_code_audit import Upstream, json_response, owned_gateway +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(180) + +_TOOLS: Final = ("get_usage_data", "get_team_usage_data", "get_tag_usage_data") + + +@dataclass(frozen=True, slots=True) +class _UsageChat: + gateway: Gateway + identity: str + + +@pytest.fixture(scope="module") +def usage_chat(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_UsageChat]: + """A proxy whose usage chat model (`openai/...`, called straight through litellm) is the scripted upstream.""" + identity: Final = f"audit-usage-chat-{uuid.uuid4().hex}" + handle: Final = register_scenario(identity, json_response({})) + try: + with owned_gateway( + tmp_path_factory.mktemp("audit-usage-chat"), + {"model_list": [], "router_settings": {"num_retries": 0}}, + environment={"OPENAI_API_BASE": handle.api_base(), "OPENAI_API_KEY": identity}, + ) as owned: + yield _UsageChat(owned, identity) + finally: + delete_scenario(handle) + + +def _tool_calls(start_date: str, end_date: str) -> dict[str, JsonValue]: + arguments: Final = json.dumps({"start_date": start_date, "end_date": end_date}) + return { + "id": "chatcmpl-usage", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": f"call-{tool}", "type": "function", "function": {"name": tool, "arguments": arguments}} + for tool in _TOOLS + ], + }, + "finish_reason": "tool_calls", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + + +def _tool_results(usage_chat: _UsageChat, start_date: str, end_date: str) -> tuple[dict[str, str], tuple[str, ...]]: + """Runs one chat turn whose model calls every tool with these dates; returns the tool messages the follow-up + model call carried, keyed by tool, and the tool statuses the SSE stream reported.""" + register_scenario(usage_chat.identity, json_response(_tool_calls(start_date, end_date))) + Upstream(usage_chat.gateway.upstream_url).drain() + with usage_chat.gateway.client.stream( + "POST", + "/usage/ai/chat", + json={"model": "openai/gpt-4o-mini", "messages": [{"role": "user", "content": "How much did we spend?"}]}, + headers={"Authorization": f"Bearer {usage_chat.gateway.key}"}, + timeout=60, + ) as response: + assert response.status_code == 200 + events: Final = tuple( + object_value(JSON_OBJECT.validate_json(line.removeprefix("data: "))) + for line in response.iter_lines() + if line.startswith("data: ") + ) + calls: Final = Upstream(usage_chat.gateway.upstream_url).calls(usage_chat.identity) + assert len(calls) == 2, calls + messages: Final = object_value(calls[1]["body"]).get("messages") + assert isinstance(messages, list), calls[1] + tool_messages: Final = { + str(message["tool_call_id"]).removeprefix("call-"): str(message["content"]) + for message in messages + if isinstance(message, dict) and message.get("role") == "tool" + } + statuses: Final = tuple(str(event.get("status")) for event in events if event.get("type") == "tool_call") + return tool_messages, statuses + + +_REVERSED: Final = "Invalid date range: end_date must be on or after start_date" +_NOT_CANONICAL: Final = "Invalid date range: start_date and end_date must be valid YYYY-MM-DD dates" +_NO_TEAM_DATA: Final = "No Team usage data found for the given date range." +_NO_TAG_DATA: Final = "No Tag usage data found for the given date range." + + +@pytest.mark.parametrize( + ("start_date", "end_date", "team_message", "tag_message"), + ( + pytest.param("2026-10-09", "2026-10-01", _REVERSED, _REVERSED, id="reversed"), + pytest.param("2026-13-01", "2026-10-09", _NOT_CANONICAL, _NOT_CANONICAL, id="invalid-month"), + pytest.param("2001/01/01", "2001-01-09", _NOT_CANONICAL, _NOT_CANONICAL, id="slashes"), + pytest.param("2001-01-01", "2001-01-09", _NO_TEAM_DATA, _NO_TAG_DATA, id="valid"), + ), +) +def test_usage_chat_tools_answer_with_the_query_layers_date_check( + usage_chat: _UsageChat, start_date: str, end_date: str, team_message: str, tag_message: str +) -> None: + """The team and tag tools give the model the daily activity routes' 400 reason; the global tool, whose query + takes these dates without that check, still answers with its usage summary.""" + tool_messages, statuses = _tool_results(usage_chat, start_date, end_date) + assert tool_messages["get_usage_data"].startswith("Total Spend: $"), tool_messages + assert tool_messages["get_team_usage_data"] == team_message, tool_messages + assert tool_messages["get_tag_usage_data"] == tag_message, tool_messages + assert statuses == ("running", "complete") * len(_TOOLS), statuses diff --git a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py index 73332728964..2121e04dcac 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py @@ -145,11 +145,14 @@ class TestOpenAIResponsesAPIConfig: (17, 17), (256, 256), (None, None), + (0, 0), + (-1, -1), ], ) def test_enforce_min_max_output_tokens(self, value, expected): - """Below the minimum clamps to 16; the boundary, larger values, and None - are returned unchanged so no previously-valid request regresses.""" + """Positive values below the minimum clamp to 16; the boundary, larger values, + None, and non-positive values are returned unchanged so the provider rejects + an invalid value instead of the proxy silently rewriting it.""" assert self.config._enforce_min_max_output_tokens(value) == expected def validate_responses_api_request_params(self, params, expected_fields): diff --git a/tests/unit/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py b/tests/unit/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py index aebd93d062e..089b092dfa6 100644 --- a/tests/unit/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py +++ b/tests/unit/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py @@ -1,3 +1,5 @@ +from typing import Final + import httpx import pytest @@ -63,6 +65,30 @@ def test_transform_create_request(config: OpenAIVectorStoreFilesConfig): assert payload["attributes"]["key"] == "value" +_SEVENTEEN_ATTRIBUTES: Final[dict[str, str]] = {f"key_{index}": f"value_{index}" for index in range(17)} + + +def test_transform_create_request_keeps_every_attribute(config: OpenAIVectorStoreFilesConfig): + _, payload = config.transform_create_vector_store_file_request( + vector_store_id="vs_123", + create_request={"file_id": "file-abc", "attributes": {**_SEVENTEEN_ATTRIBUTES, "hidden_params": "x"}}, + api_base="https://api.example.com/v1/vector_stores/vs_123/files", + ) + + assert payload["attributes"] == _SEVENTEEN_ATTRIBUTES + + +def test_transform_update_request_keeps_every_attribute(config: OpenAIVectorStoreFilesConfig): + _, payload = config.transform_update_vector_store_file_request( + vector_store_id="vs_123", + file_id="file-abc", + update_request={"attributes": {**_SEVENTEEN_ATTRIBUTES, "hidden_params": "x"}}, + api_base="https://api.example.com/v1/vector_stores/vs_123/files", + ) + + assert payload["attributes"] == _SEVENTEEN_ATTRIBUTES + + def test_transform_list_request(config: OpenAIVectorStoreFilesConfig): api_base = "https://api.example.com/v1/vector_stores/vs_123/files" url, params = config.transform_list_vector_store_files_request( diff --git a/tests/unit/proxy/management_endpoints/test_common_daily_activity.py b/tests/unit/proxy/management_endpoints/test_common_daily_activity.py index e6a6680d3e4..01d2f7007a1 100644 --- a/tests/unit/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/unit/proxy/management_endpoints/test_common_daily_activity.py @@ -2513,6 +2513,15 @@ def test_parse_canonical_date_range_reports_missing_then_malformed_dates() -> No ) +def test_parse_canonical_date_range_rejects_end_before_start() -> None: + assert parse_canonical_date_range("2026-09-26", "2026-09-24") == InvalidDateRange( + reason="end_date must be on or after start_date" + ) + assert parse_canonical_date_range("2026-09-24", "2026-09-24") == CanonicalDateRange( + start=date(2026, 9, 24), end=date(2026, 9, 24) + ) + + @pytest.mark.asyncio @pytest.mark.parametrize("start_date", ("2026-9-24", "２０２６-09-24", "2026-09-4")) async def test_get_daily_activity_rejects_non_canonical_dates_before_querying(start_date: str) -> None: diff --git a/tests/unit/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py b/tests/unit/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py index e5616a1975d..4daddb89350 100644 --- a/tests/unit/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py +++ b/tests/unit/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py @@ -3,9 +3,12 @@ Tests for AI Usage Chat module. """ import json +from collections.abc import AsyncIterator +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import ( TOOL_HANDLERS, @@ -317,6 +320,71 @@ class TestStreamUsageAiChat: assert len(chunk_events) >= 1 assert "Engineering" in chunk_events[0]["content"] + @pytest.mark.asyncio + @pytest.mark.parametrize( + "reason", + [ + "end_date must be on or after start_date", + "start_date and end_date must be valid YYYY-MM-DD dates", + ], + ids=["reversed", "invalid-format"], + ) + @pytest.mark.parametrize( + "tool_name", ["get_usage_data", "get_team_usage_data", "get_tag_usage_data"] + ) + async def test_invalid_tool_dates_are_reported_to_the_model( + self, tool_name: str, reason: str + ) -> None: + fetch: Final = AsyncMock( + side_effect=HTTPException(status_code=400, detail={"error": reason}) + ) + + tool_message, statuses = await _run_one_tool( + tool_name, "2025-01-31", "2025-01-01", fetch + ) + + fetch.assert_awaited_once() + assert tool_message == f"Invalid date range: {reason}" + assert statuses == ["running", "complete"] + + @pytest.mark.asyncio + async def test_tool_dates_the_query_layer_accepts_are_queried_as_sent( + self, + ) -> None: + fetch: Final = AsyncMock(return_value=SAMPLE_AGGREGATED_RESPONSE) + + tool_message, statuses = await _run_one_tool( + "get_usage_data", "2025/01/01", "2025-01-31T00:00:00", fetch + ) + + fetch.assert_awaited_once_with( + start_date="2025/01/01", end_date="2025-01-31T00:00:00" + ) + assert tool_message.startswith("Total Spend: $") + assert statuses == ["running", "complete"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "error", + [ + HTTPException(status_code=500, detail={"error": "db down"}), + HTTPException(status_code=400, detail="page and page_size must be >= 1"), + ], + ids=["server-error", "other-400"], + ) + async def test_other_tool_errors_stay_a_generic_fetch_error( + self, error: HTTPException + ) -> None: + tool_message, statuses = await _run_one_tool( + "get_team_usage_data", + "2025-01-01", + "2025-01-31", + AsyncMock(side_effect=error), + ) + + assert tool_message == "Error fetching team usage data. Please try again." + assert statuses == ["running", "error"] + @pytest.mark.asyncio async def test_stream_handles_error(self): with patch( @@ -524,3 +592,66 @@ class TestUsageAiChatKeepalive: assert b": ping\n\n" not in chunks assert chunks[0].startswith(b'data: {"type": "status"') assert chunks[-1] == b'data: {"type": "done"}\n\n' + + +async def _run_one_tool( + tool_name: str, start_date: str, end_date: str, fetch: AsyncMock +) -> tuple[str, list[object]]: + """Runs one admin chat turn whose model calls `tool_name` with these dates; returns the + tool message the follow-up model call carried and the tool statuses streamed.""" + arguments: Final = json.dumps({"start_date": start_date, "end_date": end_date}) + mock_tool_call: Final = MagicMock() + mock_tool_call.id = "call_dates" + mock_tool_call.function.name = tool_name + mock_tool_call.function.arguments = arguments + + mock_first_response: Final = MagicMock() + mock_first_response.choices = [MagicMock()] + mock_first_response.choices[0].message.tool_calls = [mock_tool_call] + mock_first_response.choices[0].message.model_dump.return_value = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_dates", + "type": "function", + "function": {"name": tool_name, "arguments": arguments}, + } + ], + } + + async def mock_stream() -> AsyncIterator[MagicMock]: + chunk: Final = MagicMock() + chunk.choices = [MagicMock()] + chunk.choices[0].delta.content = "Done." + yield chunk + + with ( + patch( + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm" + ) as mock_litellm, + patch.dict(TOOL_HANDLERS[tool_name], {"fetch": fetch}), + ): + mock_litellm.acompletion = AsyncMock( + side_effect=[mock_first_response, mock_stream()] + ) + events: Final[tuple[dict[str, object], ...]] = tuple( + [ + json.loads(event.replace("data: ", "").strip()) + async for event in stream_usage_ai_chat( + messages=[{"role": "user", "content": "How much did we spend?"}], + model="gpt-4o-mini", + is_admin=True, + ) + ] + ) + final_messages: Final = mock_litellm.acompletion.call_args_list[1].kwargs[ + "messages" + ] + + tool_messages: Final = [m for m in final_messages if m["role"] == "tool"] + assert len(tool_messages) == 1, tool_messages + assert tool_messages[0]["tool_call_id"] == "call_dates" + return str(tool_messages[0]["content"]), [ + e["status"] for e in events if e["type"] == "tool_call" + ] diff --git a/tests/unit/proxy/ocr_endpoints/test_endpoints.py b/tests/unit/proxy/ocr_endpoints/test_endpoints.py index 2b4f370ddee..6764e0cb434 100644 --- a/tests/unit/proxy/ocr_endpoints/test_endpoints.py +++ b/tests/unit/proxy/ocr_endpoints/test_endpoints.py @@ -3,14 +3,16 @@ Tests for the proxy OCR endpoint helpers that select the response format (`x-req-format: native | litellm`) and return the provider's native payload. """ +import io +from typing import Final from unittest.mock import AsyncMock, MagicMock import orjson import pytest -from fastapi import HTTPException +from fastapi import HTTPException, UploadFile from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse -from litellm.proxy.ocr_endpoints.endpoints import _native_response, _parse_ocr_request +from litellm.proxy.ocr_endpoints.endpoints import OCRRequestError, _native_response, _parse_ocr_request AZURE_NATIVE_OPERATION = { "status": "succeeded", @@ -145,3 +147,83 @@ def test_upload_octet_stream_content_type_falls_back_to_filename_inference(): assert "mime_type" not in document assert document["file"].name == "receipt.pdf" + + +def _multipart_request(form: dict[str, object]) -> MagicMock: + request: Final = MagicMock() + request.headers = {"content-type": "multipart/form-data; boundary=x"} + request.form = AsyncMock(return_value=form) + return request + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "form, message", + [ + ({"model": "mistral-ocr"}, "must include a 'file' field"), + ( + {"model": "mistral-ocr", "file": UploadFile(file=io.BytesIO(b""), filename="empty.pdf")}, + "Uploaded file is empty", + ), + ], + ids=["missing-file", "empty-file"], +) +async def test_malformed_multipart_upload_is_a_400_on_the_file_param(form: dict[str, object], message: str) -> None: + with pytest.raises(OCRRequestError) as exc_info: + await _parse_ocr_request(_multipart_request(form)) + + assert exc_info.value.code == "400" + assert exc_info.value.type == "invalid_request_error" + assert exc_info.value.param == "file" + assert message in exc_info.value.message + + +@pytest.mark.asyncio +async def test_unparseable_multipart_form_is_a_400() -> None: + request: Final = _multipart_request({}) + request.form = AsyncMock(side_effect=ValueError("missing boundary")) + + with pytest.raises(OCRRequestError) as exc_info: + await _parse_ocr_request(request) + + assert exc_info.value.code == "400" + assert exc_info.value.type == "invalid_request_error" + assert "Failed to parse multipart form data: missing boundary" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_multipart_file_over_the_size_limit_is_a_400_on_the_file_param(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("litellm.proxy.ocr_endpoints.endpoints._MAX_FILE_BYTES", 4) + form: Final = {"model": "mistral-ocr", "file": UploadFile(file=io.BytesIO(b"12345"), filename="big.pdf")} + + with pytest.raises(OCRRequestError) as exc_info: + await _parse_ocr_request(_multipart_request(form)) + + assert exc_info.value.code == "400" + assert exc_info.value.param == "file" + assert exc_info.value.message == "OCR file exceeds the size limit" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body, message", + [ + (b"", "Empty request body."), + (b"{not json", "Invalid JSON in request body:"), + ( + orjson.dumps({"model": "mistral-ocr", "document": {"type": "file", "file": "/etc/passwd"}}), + "document type 'file' is not supported through the JSON API.", + ), + ], + ids=["empty-body", "invalid-json", "document-type-file"], +) +async def test_malformed_json_ocr_request_is_a_400(body: bytes, message: str) -> None: + request: Final = _json_request({}, {}) + request.body = AsyncMock(return_value=body) + + with pytest.raises(OCRRequestError) as exc_info: + await _parse_ocr_request(request) + + assert exc_info.value.code == "400" + assert exc_info.value.type == "invalid_request_error" + assert exc_info.value.message.startswith(message) diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index a4b34825d95..67f29220cb6 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1760,3 +1760,133 @@ async def test_route_request_router_settings_override_skips_null_fields(): assert "context_window_fallbacks" not in call_kwargs assert "num_retries" not in call_kwargs assert "model_group_retry_policy" not in call_kwargs + + +def _router_with_defaults(default_litellm_params: dict[str, object]): + import litellm + + return litellm.Router( + model_list=[ + { + "model_name": "served", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, + "model_info": {"id": "served-deployment"}, + } + ], + default_litellm_params=default_litellm_params, + ) + + +@pytest.mark.parametrize( + "route_type, param", + [ + ("aimage_generation", "prompt"), + ("aspeech", "input"), + ("atext_completion", "prompt"), + ("atranscription", "file"), + ], +) +def test_router_default_for_positional_param_still_raises(route_type: str, param: str) -> None: + from litellm.proxy.route_llm_request import ( + ProxyMissingRequiredParamError, + raise_if_required_body_param_missing, + ) + + with pytest.raises(ProxyMissingRequiredParamError) as exc_info: + raise_if_required_body_param_missing( + route_type=route_type, + data={"model": "served"}, + llm_router=_router_with_defaults({param: "supplied-by-router-default"}), + ) + + assert exc_info.value.param == param + + +@pytest.mark.parametrize( + "route_type, param", + [ + ("aimage_generation", "prompt"), + ("aspeech", "input"), + ], +) +def test_user_config_default_for_positional_param_still_raises(route_type: str, param: str) -> None: + from litellm.proxy.route_llm_request import ( + ProxyMissingRequiredParamError, + raise_if_required_body_param_missing, + ) + + with pytest.raises(ProxyMissingRequiredParamError) as exc_info: + raise_if_required_body_param_missing( + route_type=route_type, + data={"model": "img", "user_config": {"model_list": [], "default_litellm_params": {param: "x"}}}, + llm_router=None, + ) + + assert exc_info.value.param == param + + +def test_deployment_param_for_positional_param_still_raises() -> None: + import litellm + from litellm.proxy.route_llm_request import ( + ProxyMissingRequiredParamError, + raise_if_required_body_param_missing, + ) + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "img", + "litellm_params": {"model": "openai/gpt-image-1", "api_key": "test-key", "prompt": "a cat"}, + } + ] + ) + + with pytest.raises(ProxyMissingRequiredParamError) as exc_info: + raise_if_required_body_param_missing(route_type="aimage_generation", data={"model": "img"}, llm_router=router) + + assert exc_info.value.param == "prompt" + + +@pytest.mark.parametrize( + "model", + [ + pytest.param("omni-moderation-2024-09-26", id="unknown-model"), + pytest.param("openai/gpt-4o-mini", id="deployment-name"), + pytest.param("served-deployment", id="deployment-id"), + ], +) +def test_router_default_ignored_for_moderation_model_the_router_does_not_list(model: str) -> None: + from litellm.proxy.route_llm_request import ( + ProxyMissingRequiredParamError, + raise_if_required_body_param_missing, + ) + + with pytest.raises(ProxyMissingRequiredParamError) as exc_info: + raise_if_required_body_param_missing( + route_type="amoderation", + data={"model": model}, + llm_router=_router_with_defaults({"input": "supplied"}), + ) + + assert exc_info.value.param == "input" + + +@pytest.mark.asyncio +async def test_null_user_config_is_treated_as_absent() -> None: + router: Final = _router_with_defaults({}) + + response: Final = await ( + await route_request( + { + "model": "served", + "messages": [{"role": "user", "content": "hi"}], + "mock_response": "from the main router", + "user_config": None, + }, + router, + None, + "acompletion", + ) + ) + + assert response.choices[0].message.content == "from the main router" From 719c44db3f2ec5a13d51fbf29524e001720306be Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 13 Oct 2026 12:08:01 -0700 Subject: [PATCH 004/325] fix(proxy): raise on provider count_tokens errors instead of falling back to tiktoken (#46121) Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yuneng-jiang --- litellm/proxy/proxy_server.py | 17 ++--- ...anthropic_count_tokens_system_lift_wire.py | 31 ++++---- .../test_bedrock_mantle_count_tokens_wire.py | 71 ++++++++++--------- .../test_vertex_partner_count_tokens_wire.py | 49 +++++++------ tests/unit/proxy/test_proxy_token_counter.py | 55 +++++++------- 5 files changed, 110 insertions(+), 113 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c2ca045236c..60b9ade82c0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14116,19 +14116,12 @@ async def _try_provider_token_count( code=status_code, ) if result is not None and result.error is True: - if litellm.disable_token_counter is True: - raise ProxyException( - message=result.error_message or "Token counting failed", - type="token_counting_error", - param="model", - code=result.status_code or 500, - ) - verbose_proxy_logger.warning( - "Provider token counting failed (%s): %s. Falling back to local tokenizer.", - result.status_code, - result.error_message, + raise ProxyException( + message=result.error_message or "Token counting failed", + type="token_counting_error", + param="model", + code=result.status_code or 500, ) - return None return result diff --git a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py index 13091ccd14a..a4b5c15ad54 100644 --- a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py +++ b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py @@ -398,24 +398,24 @@ def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place( ) -def test_messages_count_tokens_falls_back_locally_when_every_message_is_system( +def test_messages_count_tokens_surfaces_the_peer_error_when_every_message_is_system( deployment: _Deployment, gateway: Gateway ) -> None: with _listening(deployment) as peer: body: Final[dict[str, JsonValue]] = {"model": deployment.model, "messages": [LEADING]} - local: Final = _local_count(gateway, body) response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": local}, response.text + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text assert _bodies(peer, deployment.provider) == ({"model": _MODEL, "messages": [], "system": [LIFTED]},) -def test_messages_count_tokens_leaves_the_request_untouched_for_a_non_text_system( +def test_messages_count_tokens_surfaces_the_peer_error_for_a_non_text_system( deployment: _Deployment, gateway: Gateway ) -> None: with _listening(deployment) as peer: - local: Final = _local_count(gateway, count_request(deployment.model, STRING_CASE)) response: Final = _count(gateway, {**count_request(deployment.model, STRING_CASE), "system": 5}) - assert _payload(response) == {"input_tokens": local}, response.text + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text assert _bodies(peer, deployment.provider) == ({"model": _MODEL, "messages": [LEADING, USER], "system": 5},) @@ -428,14 +428,14 @@ def test_messages_count_tokens_repeated_request_lifts_each_time(deployment: _Dep assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * 2 -def test_messages_count_tokens_answers_a_leading_system_without_content_before_any_peer_call( +def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_without_content( deployment: _Deployment, gateway: Gateway ) -> None: with _listening(deployment) as peer: body: Final[dict[str, JsonValue]] = {"model": deployment.model, "messages": [{"role": "system"}, USER]} - local: Final = _local_count(gateway, body) response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": local}, response.text + assert 500 <= response.status_code < 600, response.text + assert "input_tokens" not in response.text, response.text assert peer.drain() == () follow_up: Final = _count(gateway, count_request(deployment.model, STRING_CASE)) assert _payload(follow_up) == {"input_tokens": _COUNT}, follow_up.text @@ -458,14 +458,14 @@ def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value( @pytest.mark.parametrize("status", [400, 403, 404, 500, 503]) -def test_messages_count_tokens_falls_back_locally_when_the_peer_rejects_the_lifted_body( +def test_messages_count_tokens_surfaces_the_peer_status_when_the_peer_rejects_the_lifted_body( deployment: _Deployment, gateway: Gateway, status: int ) -> None: with _listening(deployment, _rejecting(status)) as peer: body: Final = count_request(deployment.model, STRING_CASE) - local: Final = _local_count(gateway, body) response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": local}, response.text + assert response.status_code == status, response.text + assert "input_tokens" not in response.text, response.text assert _bodies(peer, deployment.provider) == (_STRING_BODY,) @@ -478,14 +478,13 @@ def test_messages_count_tokens_unauthenticated_request_never_reaches_the_peer( assert peer.drain() == () -def test_peer_outage_between_concurrent_waves_falls_back_then_recovers( +def test_peer_outage_between_concurrent_waves_errors_then_recovers( deployment: _Deployment, gateway: Gateway ) -> None: with ExitStack() as stack: clients: Final = _clients(stack, str(gateway.client.base_url), 8) pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) body: Final = count_request(deployment.model, STRING_CASE) - local: Final = _local_count(gateway, body) def count(client: httpx.Client) -> tuple[int, JsonValue]: return _counted_on(client, gateway.key, body) @@ -493,7 +492,9 @@ def test_peer_outage_between_concurrent_waves_falls_back_then_recovers( with _listening(deployment) as peer: assert tuple(pool.map(count, clients)) == ((200, _COUNT),) * len(clients) assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * len(clients) - assert tuple(pool.map(count, clients)) == ((200, local),) * len(clients) + outage: Final = tuple(pool.map(count, clients)) + for status, tokens in outage: + assert 500 <= status < 600 and tokens is None, outage with _listening(deployment) as revived: assert tuple(pool.map(count, clients)) == ((200, _COUNT),) * len(clients) assert _bodies(revived, deployment.provider) == (_STRING_BODY,) * len(clients) diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index 7f59b818cba..a5ab085c492 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -522,7 +522,7 @@ def test_model_the_runtime_counts_never_reaches_mantle(counting_proxy: OwnedProx assert mantle.drain() == () -def test_non_claude_model_the_runtime_cannot_count_falls_back_locally_without_mantle( +def test_non_claude_model_the_runtime_cannot_count_surfaces_the_error_without_mantle( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -533,7 +533,8 @@ def test_non_claude_model_the_runtime_cannot_count_falls_back_locally_without_ma ): model: Final = _deployment(scenario, runtime.url, _NOVA) response: Final = _count(gateway, _bare(model)) - assert _payload(response) == {"input_tokens": _local_count(gateway, _bare(model))}, response.text + assert response.status_code == 400, response.text + assert _UNSUPPORTED in response.text and "input_tokens" not in response.text, response.text assert _runtime_count_targets(runtime) == (f"/model/{_NOVA}/count-tokens",) assert mantle.drain() == () @@ -551,17 +552,19 @@ def test_runtime_403_on_a_claude_model_never_reaches_mantle(counting_proxy: Owne "messages": [{"role": "user", "content": "status=403 Count this message"}], } response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": _local_count(gateway, body)}, response.text + assert response.status_code == 403, response.text + assert "input_tokens" not in response.text, response.text assert len(_runtime_count_targets(runtime)) == 1 assert mantle.drain() == () -def test_unreachable_runtime_falls_back_locally_without_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: +def test_unreachable_runtime_surfaces_the_error_without_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: gateway: Final = counting_proxy.gateway with wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario: model: Final = _deployment(scenario, _closed_port_url()) response: Final = _count(gateway, _full(model)) - assert _payload(response) == {"input_tokens": _local_count(gateway, _full(model))}, response.text + assert 500 <= response.status_code < 600, response.text + assert "input_tokens" not in response.text, response.text assert mantle.drain() == () @@ -770,7 +773,7 @@ def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place_for_mant ) -def test_messages_count_tokens_falls_back_locally_when_every_message_is_system_for_mantle( +def test_messages_count_tokens_surfaces_the_mantle_error_when_every_message_is_system( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -781,14 +784,14 @@ def test_messages_count_tokens_falls_back_locally_when_every_message_is_system_f ): model: Final = _deployment(scenario, runtime.url) body: Final[dict[str, JsonValue]] = {"model": model, "messages": [LEADING]} - local: Final = _local_count(gateway, body) response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": local}, response.text + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_mantle_body(messages=[], system=[LIFTED]),) -def test_messages_count_tokens_leaves_a_leading_system_in_place_beside_a_non_text_system_for_mantle( +def test_messages_count_tokens_surfaces_the_mantle_error_beside_a_non_text_system( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -798,9 +801,9 @@ def test_messages_count_tokens_leaves_a_leading_system_in_place_beside_a_non_tex gateway.scenario() as scenario, ): model: Final = _deployment(scenario, runtime.url) - local: Final = _local_count(gateway, count_request(model, STRING_CASE)) response: Final = _count(gateway, {**count_request(model, STRING_CASE), "system": 5}) - assert _payload(response) == {"input_tokens": local}, response.text + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_mantle_body(messages=[LEADING, USER], system=5),) @@ -821,7 +824,7 @@ def test_messages_count_tokens_repeated_request_lifts_the_leading_system_each_ti assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) * 2 -def test_messages_count_tokens_answers_a_leading_system_without_content_before_any_mantle_call( +def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_without_content( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -832,9 +835,9 @@ def test_messages_count_tokens_answers_a_leading_system_without_content_before_a ): model: Final = _deployment(scenario, runtime.url) body: Final[dict[str, JsonValue]] = {"model": model, "messages": [{"role": "system"}, USER]} - local: Final = _local_count(gateway, body) response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": local}, response.text + assert 500 <= response.status_code < 600, response.text + assert "input_tokens" not in response.text, response.text assert mantle.drain() == () follow_up: Final = _count(gateway, count_request(model, STRING_CASE)) assert _payload(follow_up) == {"input_tokens": _MANTLE_COUNT}, follow_up.text @@ -889,7 +892,7 @@ def test_disabled_token_counter_counts_a_leading_system_through_mantle( assert len(_runtime_count_targets(runtime)) == 1 -def test_mantle_outage_between_concurrent_lifted_waves_falls_back_then_recovers( +def test_mantle_outage_between_concurrent_lifted_waves_errors_then_recovers( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -900,7 +903,6 @@ def test_mantle_outage_between_concurrent_lifted_waves_falls_back_then_recovers( scenario: Final = stack.enter_context(gateway.scenario()) model: Final = _deployment(scenario, runtime.url) body: Final = count_request(model, STRING_CASE) - local: Final = _local_count(gateway, body) def count(client: httpx.Client) -> tuple[int, JsonValue]: return _counted_on(client, gateway.key, body) @@ -908,7 +910,9 @@ def test_mantle_outage_between_concurrent_lifted_waves_falls_back_then_recovers( with wire_server(_mantle(_strict), port=mantle_port) as mantle: assert tuple(pool.map(count, clients)) == ((200, _MANTLE_COUNT),) * len(clients) assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) * len(clients) - assert tuple(pool.map(count, clients)) == ((200, local),) * len(clients) + outage: Final = tuple(pool.map(count, clients)) + for status, tokens in outage: + assert 500 <= status < 600 and tokens is None, outage with wire_server(_mantle(_strict), port=mantle_port) as revived: assert tuple(pool.map(count, clients)) == ((200, _MANTLE_COUNT),) * len(clients) assert _mantle_bodies(revived) == (_MANTLE_LIFTED,) * len(clients) @@ -916,7 +920,7 @@ def test_mantle_outage_between_concurrent_lifted_waves_falls_back_then_recovers( @pytest.mark.parametrize("status", [400, 403, 404, 500, 503]) -def test_messages_count_tokens_falls_back_locally_when_mantle_errors( +def test_messages_count_tokens_surfaces_the_mantle_status_when_mantle_errors( counting_proxy: OwnedProxy, mantle_port: int, status: int ) -> None: gateway: Final = counting_proxy.gateway @@ -927,7 +931,8 @@ def test_messages_count_tokens_falls_back_locally_when_mantle_errors( ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, _full(model)) - assert _payload(response) == {"input_tokens": _local_count(gateway, _full(model))}, response.text + assert response.status_code == status, response.text + assert "input_tokens" not in response.text, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_MANTLE_FULL,) @@ -937,7 +942,7 @@ def test_messages_count_tokens_falls_back_locally_when_mantle_errors( [b'{"inputTokens": 7}', b"not json at all", b"{}"], ids=["runtime_key", "not_json", "empty_object"], ) -def test_messages_count_tokens_falls_back_locally_when_mantle_answers_without_input_tokens( +def test_messages_count_tokens_errors_when_mantle_answers_without_input_tokens( counting_proxy: OwnedProxy, mantle_port: int, body: bytes ) -> None: gateway: Final = counting_proxy.gateway @@ -948,21 +953,22 @@ def test_messages_count_tokens_falls_back_locally_when_mantle_answers_without_in ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, _bare(model)) - assert _payload(response) == {"input_tokens": _local_count(gateway, _bare(model))}, response.text + assert 500 <= response.status_code < 600, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_MANTLE_BARE,) -def test_messages_count_tokens_falls_back_locally_when_mantle_is_unreachable(counting_proxy: OwnedProxy) -> None: +def test_messages_count_tokens_errors_when_mantle_is_unreachable(counting_proxy: OwnedProxy) -> None: gateway: Final = counting_proxy.gateway with wire_server(_runtime) as runtime, gateway.scenario() as scenario: model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, _full(model)) - assert _payload(response) == {"input_tokens": _local_count(gateway, _full(model))}, response.text + assert 500 <= response.status_code < 600, response.text + assert "input_tokens" not in response.text, response.text assert len(_runtime_count_targets(runtime)) == 1 -def test_messages_count_tokens_falls_back_locally_when_mantle_never_answers( +def test_messages_count_tokens_errors_when_mantle_never_answers( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -983,18 +989,18 @@ def test_messages_count_tokens_falls_back_locally_when_mantle_never_answers( scenario: Final = stack.enter_context(gateway.scenario()) model: Final = _deployment(scenario, runtime.url) body: Final = _full(model) - local: Final = _local_count(gateway, body) future: Final = pool.submit(_counted_on, client, gateway.key, body) eventually(held.qsize, lambda size: size == 1, seconds=30) assert gateway.request("GET", "/health/liveliness").status_code == 200 - assert _local_count(gateway, body) == local + assert _local_count(gateway, body) > 0 assert not future.done() - assert future.result(timeout=120) == (200, local) + status, tokens = future.result(timeout=120) + assert 500 <= status < 600 and tokens is None, (status, tokens) assert len(_runtime_count_targets(runtime)) == 1 assert len(mantle.drain()) == 1 -def test_messages_count_tokens_falls_back_locally_when_mantle_rejects_a_non_text_system( +def test_messages_count_tokens_surfaces_the_mantle_error_for_a_non_text_system( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -1005,7 +1011,8 @@ def test_messages_count_tokens_falls_back_locally_when_mantle_rejects_a_non_text ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": 5}) - assert _payload(response) == {"input_tokens": _local_count(gateway, _bare(model))}, response.text + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_mantle_body(system=5),) @@ -1174,7 +1181,7 @@ def test_disabled_token_counter_surfaces_the_mantle_error_instead_of_counting_lo assert len(_runtime_count_targets(runtime)) == 4 -def test_mantle_outage_between_concurrent_waves_falls_back_then_recovers( +def test_mantle_outage_between_concurrent_waves_errors_then_recovers( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -1185,7 +1192,6 @@ def test_mantle_outage_between_concurrent_waves_falls_back_then_recovers( scenario: Final = stack.enter_context(gateway.scenario()) model: Final = _deployment(scenario, runtime.url) body: Final = _full(model) - local: Final = _local_count(gateway, body) def generate_then_count(client: httpx.Client) -> tuple[int, int, JsonValue]: return _generated_then_counted(client, gateway.key, model, body) @@ -1197,7 +1203,8 @@ def test_mantle_outage_between_concurrent_waves_falls_back_then_recovers( assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _MANTLE_COUNT),) * len(clients) assert _mantle_bodies(mantle) == (_MANTLE_FULL,) * len(clients) outage: Final = tuple(pool.map(count_only, clients)) - assert outage == ((200, local),) * len(clients) + for status, tokens in outage: + assert 500 <= status < 600 and tokens is None, outage with wire_server(_mantle(), port=mantle_port) as revived: assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _MANTLE_COUNT),) * len(clients) assert _mantle_bodies(revived) == (_MANTLE_FULL,) * len(clients) diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py index 558ee3a004e..e181e38b1bf 100644 --- a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py @@ -344,15 +344,14 @@ def test_utils_token_counter_local_mode_never_calls_the_peer(gateway: Gateway) - assert wire.drain() == () -def test_utils_token_counter_falls_back_locally_when_peer_rejects_openai_tools(gateway: Gateway) -> None: +def test_utils_token_counter_surfaces_the_peer_error_for_openai_tools(gateway: Gateway) -> None: with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) body: Final = {**_bare(model), "tools": _OPENAI_TOOLS} response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "true"}) - payload: Final = _payload(response) + assert response.status_code == 400, response.text + assert "total_tokens" not in response.text, response.text assert _counted_bodies(wire) == ({**_PEER_BARE, "tools": _OPENAI_TOOLS},) - assert payload["total_tokens"] == _local_count(gateway, body), response.text - assert payload["tokenizer_type"] != "vertex_ai_partner_models", response.text def test_responses_input_tokens_counts_through_the_partner_peer(gateway: Gateway) -> None: @@ -365,7 +364,7 @@ def test_responses_input_tokens_counts_through_the_partner_peer(gateway: Gateway assert _counted_bodies(wire) == (_PEER_BARE,) -def test_responses_input_tokens_falls_back_locally_when_peer_rejects(gateway: Gateway) -> None: +def test_responses_input_tokens_surfaces_the_peer_error(gateway: Gateway) -> None: with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = gateway.request( @@ -373,11 +372,10 @@ def test_responses_input_tokens_falls_back_locally_when_peer_rejects(gateway: Ga "/v1/responses/input_tokens", {"model": model, "input": "Count this message", "instructions": "Be terse", "tools": _RESPONSES_TOOLS}, ) - payload: Final = _payload(response) + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text (sent,) = _counted_bodies(wire) assert sent["tools"] == _RESPONSES_TOOLS, sent - local: Final = _local_count(gateway, {"model": model, "messages": sent["messages"], "tools": _RESPONSES_TOOLS}) - assert payload == {"object": "response.input_tokens", "input_tokens": local}, response.text def test_gemini_count_tokens_route_reaches_the_partner_peer(gateway: Gateway) -> None: @@ -388,26 +386,26 @@ def test_gemini_count_tokens_route_reaches_the_partner_peer(gateway: Gateway) -> assert _counted_bodies(wire) == ({"model": _BACKEND, "messages": _GEMINI_MESSAGES},) -def test_gemini_count_tokens_route_falls_back_locally_when_peer_rejects(gateway: Gateway) -> None: +def test_gemini_count_tokens_route_surfaces_the_peer_error(gateway: Gateway) -> None: with wire_server(_peer(_rejecting(400))) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) - payload: Final = _payload(response) + assert response.status_code == 400, response.text + assert "totalTokens" not in response.text, response.text assert _counted_bodies(wire) == ({"model": _BACKEND, "messages": _GEMINI_MESSAGES},) - assert payload["totalTokens"] == _local_count(gateway, {"model": model, "messages": _GEMINI_MESSAGES}) @pytest.mark.parametrize("status", [400, 500, 503]) -def test_messages_count_tokens_falls_back_locally_when_peer_errors(gateway: Gateway, status: int) -> None: +def test_messages_count_tokens_surfaces_the_peer_error_when_peer_errors(gateway: Gateway, status: int) -> None: with wire_server(_peer(_rejecting(status))) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = _count(gateway, _full(model)) - payload: Final = _payload(response) + assert response.status_code == status, response.text + assert "input_tokens" not in response.text, response.text assert _counted_bodies(wire) == (_PEER_FULL,) - assert payload == {"input_tokens": _local_count(gateway, _full(model))}, response.text -def test_messages_count_tokens_falls_back_locally_when_token_endpoint_rejects(gateway: Gateway) -> None: +def test_messages_count_tokens_surfaces_the_error_when_token_endpoint_rejects(gateway: Gateway) -> None: def respond(request: Request) -> Reply: if request.target == "/_oauth/token": return Reply( @@ -421,17 +419,18 @@ def test_messages_count_tokens_falls_back_locally_when_token_endpoint_rejects(ga gateway, scenario, wire.url, vertex_credentials=service_account_json(_PROJECT, wire.url) ) response: Final = _count(gateway, _full(model)) - payload: Final = _payload(response) + assert 500 <= response.status_code < 600, response.text + assert "input_tokens" not in response.text, response.text targets: Final = frozenset(request.target for request in wire.drain()) assert targets == {"/_oauth/token"}, targets - assert payload == {"input_tokens": _local_count(gateway, _full(model))}, response.text -def test_messages_count_tokens_falls_back_locally_when_peer_is_unreachable(gateway: Gateway) -> None: +def test_messages_count_tokens_errors_when_peer_is_unreachable(gateway: Gateway) -> None: with gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, _closed_port_url()) response: Final = _count(gateway, _full(model)) - assert _payload(response) == {"input_tokens": _local_count(gateway, _full(model))}, response.text + assert 500 <= response.status_code < 600, response.text + assert "input_tokens" not in response.text, response.text @pytest.mark.parametrize( @@ -476,13 +475,13 @@ def test_messages_count_tokens_leaves_null_system_and_tools_out(gateway: Gateway assert _counted_bodies(wire) == (_PEER_BARE,) -def test_messages_count_tokens_falls_back_locally_when_peer_rejects_a_non_text_system(gateway: Gateway) -> None: +def test_messages_count_tokens_surfaces_the_peer_error_for_a_non_text_system(gateway: Gateway) -> None: with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = _count(gateway, {**_bare(model), "system": 5}) - payload: Final = _payload(response) + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text assert _counted_bodies(wire) == ({**_PEER_BARE, "system": 5},) - assert payload == {"input_tokens": _local_count(gateway, _bare(model))}, response.text def test_messages_count_tokens_forwards_a_5kb_system_verbatim(gateway: Gateway) -> None: @@ -578,7 +577,7 @@ def test_disabled_token_counter_surfaces_provider_failures_instead_of_counting_l assert len(_counted_bodies(wire)) == 2 -def test_peer_outage_between_concurrent_waves_falls_back_then_recovers(gateway: Gateway) -> None: +def test_peer_outage_between_concurrent_waves_errors_then_recovers(gateway: Gateway) -> None: with ExitStack() as stack: clients: Final = _clients(stack, str(gateway.client.base_url), 8) pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) @@ -587,7 +586,6 @@ def test_peer_outage_between_concurrent_waves_falls_back_then_recovers(gateway: port: Final = int(wire.url.rsplit(":", 1)[1]) model: Final = _deployment(gateway, scenario, wire.url) body: Final = _full(model) - local: Final = _local_count(gateway, body) def generate_then_count(client: httpx.Client) -> tuple[int, int, JsonValue]: return _generated_then_counted(client, gateway.key, model, body) @@ -595,7 +593,8 @@ def test_peer_outage_between_concurrent_waves_falls_back_then_recovers(gateway: assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _PEER_COUNT),) * len(clients) assert _counted_bodies(wire) == (_PEER_FULL,) * len(clients) outage: Final = tuple(pool.map(lambda client: _counted_on(client, gateway.key, body), clients)) - assert outage == ((200, local),) * len(clients) + for status, tokens in outage: + assert 500 <= status < 600 and tokens is None, outage with wire_server(_peer(), port=port) as revived: assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _PEER_COUNT),) * len(clients) assert _counted_bodies(revived) == (_PEER_FULL,) * len(clients) diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index e4c4ddce64f..c93c199cfa2 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -310,20 +310,16 @@ async def test_internal_token_counter_anthropic_provider_detection(): setattr(litellm.proxy.proxy_server, "llm_router", llm_router) # Test with is_direct_request=False (simulating call from Anthropic endpoint) - response = await token_counter( - request=TokenCountRequest( - model="claude-test", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, - ) - - print("Anthropic provider test response:", response) - - # Verify response structure - assert response.request_model == "claude-test" - assert response.model_used == "claude-3-sonnet-20240229" - assert response.total_tokens > 0 + # The provider counter is engaged: its API error surfaces instead of a local count + with pytest.raises(ProxyException) as exc_info: + await token_counter( + request=TokenCountRequest( + model="claude-test", + messages=[{"role": "user", "content": "hello"}], + ), + call_endpoint=True, + ) + assert exc_info.value.type == "token_counting_error" # Test with non-Anthropic provider llm_router = Router( @@ -1018,10 +1014,10 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): @pytest.mark.asyncio -async def test_proxy_token_counter_error_falls_back_when_enabled(): +async def test_proxy_token_counter_error_raises_when_enabled(): """ - Test that proxy token_counter falls back to local tokenizer when disable_token_counter=False - and provider returns an error response. + Test that proxy token_counter raises ProxyException when disable_token_counter=False + and the provider returns an error response, never falling back to the local tokenizer. """ # Create error response error_response = TokenCountResponse( @@ -1069,19 +1065,20 @@ async def test_proxy_token_counter_error_falls_back_when_enabled(): mock_get_provider_token_counter ) - # Should not raise, should fall back to local tokenizer - result = await token_counter( - request=TokenCountRequest( - model="claude-bedrock", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, - ) + with patch("litellm.token_counter") as mock_local_counter: + with pytest.raises(ProxyException) as exc_info: + await token_counter( + request=TokenCountRequest( + model="claude-bedrock", + messages=[{"role": "user", "content": "hello"}], + ), + call_endpoint=True, + ) - # Should have used the fallback tokenizer - assert result.error is False - assert result.total_tokens > 0 - assert result.tokenizer_type != "bedrock_api" + assert exc_info.value.code == "429" + assert "Rate limit exceeded" in exc_info.value.message + assert exc_info.value.type == "token_counting_error" + mock_local_counter.assert_not_called() finally: litellm.disable_token_counter = original_disable litellm.proxy.proxy_server._get_provider_token_counter = ( From 0b3fcc197dc8b703680c0ac7eb43070fa8fa76c9 Mon Sep 17 00:00:00 2001 From: "moyai-devin-berriai[bot]" <336287033+moyai-devin-berriai[bot]@users.noreply.github.com> Date: Tue, 13 Oct 2026 12:25:56 -0700 Subject: [PATCH 005/325] fix(logging): redact nested callback input and provider response fields (#45923) * fix(logging): redact nested callback input and provider response fields * fix(logging): keep shared request body for per-callback redaction and drop hidden provider fields The deprecated per-callback message_logging=False path ran perform_redaction on the shared model_call_details, which now also redacted the nested proxy_server_request body. The spend log Request column and previous_response_id replay read that body, so one opted-out callback blanked them for everyone. That path now keeps the shared body and hidden params, and redact_streaming_responses_for_custom_logger hands the opted-out callback its own redacted litellm_params. Redacted litellm_params also clear provider_specific_fields copied into metadata.hidden_params (Anthropic thinking text), in litellm_params and in the nested request body, and the nested body redacts system and instructions. * fix(logging): keep replayed instructions and router usage marker under redaction The shared request body keeps system and instructions again, so spend logs store them and previous_response_id follow-ups replay the real instructions. Per-callback copies still redact them. Provider-specific fields in metadata.hidden_params are cleared in place, so the metadata dict keeps its identity and the router's counted-usage marker stays visible to deployment_callback_on_success (no double tpm/rpm count). * fix(logging): never copy shared metadata dicts during redaction The shared redaction pass now leaves litellm_params metadata dicts as they are, so a hidden_params write from the other success handler between the in-place clear and the copy can no longer move the router usage marker to a stale dict. The key-presence predicate change is reverted as unneeded. * refactor(logging): redact nested request body and provider fields with the existing message rules Run perform_redaction on a copy of proxy_server_request.body, the same call the spend log already makes, and redact reasoning inside provider_specific_fields the way the message's own reasoning_content and thinking_blocks are redacted, instead of clearing every provider field. Opaque fields such as Gemini thought_signatures stay, and the metadata dicts and streaming objects are handled as on main. --------- Co-authored-by: moyai-devin-berriai[bot] <336287033+moyai-devin-berriai[bot]@users.noreply.github.com> Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- litellm/litellm_core_utils/redact_messages.py | 83 ++- tests/integration/_support/redaction_probe.py | 47 ++ .../test_custom_logger_redaction.py | 485 ++++++++++++++++++ .../test_redact_messages.py | 120 +++++ 4 files changed, 728 insertions(+), 7 deletions(-) create mode 100644 tests/integration/_support/redaction_probe.py create mode 100644 tests/integration/observability/test_custom_logger_redaction.py diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 7c6c39abb76..c2cca5739cf 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -108,6 +108,7 @@ def _redact_choice_content(choice): choice.message.reasoning_content = REDACTED_BY_LITELLM if hasattr(choice.message, "thinking_blocks"): choice.message.thinking_blocks = None + _redact_provider_specific_fields(choice.message) _redact_tool_calls(getattr(choice.message, "tool_calls", None)) _redact_function_call(getattr(choice.message, "function_call", None)) elif isinstance(choice, litellm.utils.StreamingChoices): @@ -117,10 +118,39 @@ def _redact_choice_content(choice): choice.delta.reasoning_content = REDACTED_BY_LITELLM if hasattr(choice.delta, "thinking_blocks"): choice.delta.thinking_blocks = None + _redact_provider_specific_fields(choice.delta) _redact_tool_calls(getattr(choice.delta, "tool_calls", None)) _redact_function_call(getattr(choice.delta, "function_call", None)) +def _as_mapping(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, Mapping) else None # pyright: ignore[reportUnknownVariableType] # keys are read one by one + + +def _redact_provider_specific_fields(message: object) -> None: + if (fields := _as_mapping(getattr(message, "provider_specific_fields", None))) is not None: + setattr(message, "provider_specific_fields", _redacted_provider_specific_fields(fields)) # noqa: B010 # typed as object + + +def _redact_provider_specific_fields_dict(choice: Mapping[str, object]) -> None: + for message in (choice.get("message"), choice.get("delta")): + if ( + (readable := _as_mapping(message)) is not None + and (fields := _as_mapping(readable.get("provider_specific_fields"))) is not None + and isinstance(message, dict) + ): + message["provider_specific_fields"] = _redacted_provider_specific_fields(fields) + + +def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[str, object]: + """Copy of provider_specific_fields with its reasoning redacted the same way as the message's own.""" + return { + **fields, + **({"reasoning_content": REDACTED_BY_LITELLM} if fields.get("reasoning_content") is not None else {}), + **{key: None for key in ("thinking_blocks", "reasoningContentBlocks", "citations") if key in fields}, + } + + def _redact_responses_api_output(output_items): """Helper to redact ResponsesAPIResponse output items.""" for output_item in output_items: @@ -228,6 +258,7 @@ def _redact_tool_calls_dict(message: Mapping[str, object]) -> None: def _redact_model_response_dict_choices(choices, redacted_str: str): for choice in choices: if isinstance(choice, dict): + _redact_provider_specific_fields_dict(choice) if "message" in choice and isinstance(choice["message"], dict): if choice["message"].get("content") is not None: choice["message"]["content"] = redacted_str @@ -261,8 +292,37 @@ def _redacted_baseline_metadata(metadata: Mapping[str, object]) -> Mapping[str, return {**metadata, "_autorouter_baseline_route": replace(route, request_parameters=None)} -def redacted_litellm_params(params: Mapping[str, object]) -> dict[str, object]: - request: Final = params.get("proxy_server_request") +def _redacted_proxy_server_request(request: Mapping[str, object], *, redact_body: bool) -> Mapping[str, object]: + """Copy of proxy_server_request; redact_body redacts its body the same way as the top-level request.""" + body: Final = _as_mapping(request.get("body")) + if not redact_body or body is None: + return without_classifier_audit(request) + redacted_body: Final = {**body} + perform_redaction(redacted_body, result=None) + return {**without_classifier_audit(request), "body": redacted_body} + + +def _redact_hidden_provider_specific_fields(params: Mapping[str, object]) -> None: + """Redact the provider_specific_fields copied into metadata hidden_params, keeping each metadata dict.""" + for metadata in (params.get("metadata"), params.get("litellm_metadata")): + if ( + (readable := _as_mapping(metadata)) is not None + and (hidden_params := _as_mapping(readable.get("hidden_params"))) is not None + and (fields := _as_mapping(hidden_params.get("provider_specific_fields"))) is not None + and isinstance(metadata, dict) + ): + metadata["hidden_params"] = { + **hidden_params, + "provider_specific_fields": _redacted_provider_specific_fields(fields), + } + + +def redacted_litellm_params(params: Mapping[str, object], *, redact_request_body: bool = True) -> dict[str, object]: + """ + Copy of litellm_params without classifier audit fields or baseline route request parameters. + redact_request_body also redacts the nested proxy_server_request body. + """ + request: Final = _as_mapping(params.get("proxy_server_request")) return { **params, **{ @@ -270,7 +330,11 @@ def redacted_litellm_params(params: Mapping[str, object]) -> dict[str, object]: for key, value in params.items() if key in ("metadata", "litellm_metadata") and isinstance(value, Mapping) }, - **({"proxy_server_request": without_classifier_audit(request)} if isinstance(request, Mapping) else {}), + **( + {"proxy_server_request": _redacted_proxy_server_request(request, redact_body=redact_request_body)} + if request is not None + else {} + ), } @@ -280,12 +344,17 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons redact_streaming_responses=False skips the in-place redaction of the shared streaming response entries; per-callback redaction hands each opted-out callback its own redacted - copy via redact_streaming_responses_for_custom_logger instead. + copy via redact_streaming_responses_for_custom_logger instead. It also leaves the shared + proxy request body and hidden provider fields as they are, since the spend log reads them. """ # Redact model_call_details - params: Final = model_call_details.get("litellm_params") - if isinstance(params, Mapping): - model_call_details["litellm_params"] = redacted_litellm_params(params) + params: Final = _as_mapping(model_call_details.get("litellm_params")) + if params is not None: + if redact_streaming_responses: + _redact_hidden_provider_specific_fields(params) + model_call_details["litellm_params"] = redacted_litellm_params( + params, redact_request_body=redact_streaming_responses + ) model_call_details["messages"] = [{"role": "user", "content": REDACTED_BY_LITELLM}] model_call_details["prompt"] = "" model_call_details["input"] = "" diff --git a/tests/integration/_support/redaction_probe.py b/tests/integration/_support/redaction_probe.py new file mode 100644 index 00000000000..2a344daa73c --- /dev/null +++ b/tests/integration/_support/redaction_probe.py @@ -0,0 +1,47 @@ +import json +import os +from collections.abc import Mapping +from datetime import datetime +from typing import Final + +import httpx +from pydantic import BaseModel + +from litellm.integrations.custom_logger import CustomLogger + + +def _payload(kwargs: Mapping[str, object], response_obj: object) -> str: + params: Final = kwargs.get("litellm_params") + request: Final = params.get("proxy_server_request") if isinstance(params, Mapping) else None + standard: Final = kwargs.get("standard_logging_object") + return json.dumps( + { + "request": request, + "response": response_obj.model_dump() if isinstance(response_obj, BaseModel) else response_obj, + "standard_response": standard.get("response") if isinstance(standard, Mapping) else None, + }, + default=str, + ) + + +class SyncRedactionProbe(CustomLogger): + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + with httpx.Client(trust_env=False) as client: + response: Final = client.post(os.environ["REDACTION_PROBE_URL"], content=_payload(kwargs, response_obj)) + response.raise_for_status() + + +class RedactionProbe(CustomLogger): + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + async with httpx.AsyncClient(trust_env=False) as client: + response: Final = await client.post( + os.environ["REDACTION_PROBE_URL"], content=_payload(kwargs, response_obj) + ) + response.raise_for_status() + + +probe: Final = RedactionProbe() diff --git a/tests/integration/observability/test_custom_logger_redaction.py b/tests/integration/observability/test_custom_logger_redaction.py new file mode 100644 index 00000000000..adc0eb9103c --- /dev/null +++ b/tests/integration/observability/test_custom_logger_redaction.py @@ -0,0 +1,485 @@ +import asyncio +import json +import os +import uuid +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment, held +from integration._support.database import read_rows +from integration._support.process import owned_proxy +from integration._support.redaction_probe import RedactionProbe, SyncRedactionProbe +from integration._support.wire import Reply, Request, Wire, wire_server +from redis import Redis + +import litellm +from litellm.constants import REDACTED_BY_LITELLM +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + +INPUT: Final = "private-input-redaction-canary" +REASONING: Final = "private-reasoning-redaction-canary" +OUTPUT: Final = "private-output-redaction-canary" +SIGNATURE: Final = "opaque-thought-signature-canary" +INSTRUCTIONS: Final = "Reply in French. redaction-instructions-canary" +BRIDGED: Final = "bridged-redaction-canary" +USAGE_DEPLOYMENT: Final = f"redaction-usage-{uuid.uuid4().hex}" +REDACT: Final = {"x-litellm-enable-message-redaction": "true"} +pytestmark = pytest.mark.timeout(240) + + +def _provider(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + body: Final = json.loads(request.body) + if request.target == "/v1/embeddings": + assert body["input"] == INPUT, body + return Reply( + body=json.dumps( + { + "object": "list", + "model": body["model"], + "data": [{"object": "embedding", "index": 0, "embedding": [0.25, 0.5]}], + "usage": {"prompt_tokens": 10, "total_tokens": 10}, + } + ).encode() + ) + if request.target == "/v1/completions": + assert body["prompt"] == INPUT, body + return Reply( + body=json.dumps( + { + "id": "cmpl-redaction-probe", + "object": "text_completion", + "created": 1, + "model": body["model"], + "choices": [{"text": OUTPUT, "index": 0, "finish_reason": "stop", "logprobs": None}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + if request.target == "/v1/responses": + assert body["input"] == INPUT, body + return Reply( + body=json.dumps( + { + "id": "resp-redaction-probe", + "object": "response", + "created_at": 1, + "model": body["model"], + "status": "completed", + "output": [ + { + "id": "msg-redaction-probe", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": OUTPUT, "annotations": []}], + } + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + if request.target == "/v1/messages": + if BRIDGED not in request.body.decode(): + assert body["messages"] == [{"role": "user", "content": INPUT}], body + return Reply( + body=json.dumps( + { + "id": "msg-redaction-probe", + "type": "message", + "role": "assistant", + "model": body["model"], + "content": [ + {"type": "thinking", "thinking": REASONING, "signature": "synthetic-signature"}, + {"type": "text", "text": OUTPUT}, + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + ).encode() + ) + assert body["messages"] == [{"role": "user", "content": INPUT}], body + payload: Final = { + "id": "chatcmpl-redaction-" + body.get("user", "probe"), + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": OUTPUT, + "provider_specific_fields": { + "reasoning_content": REASONING, + "thinking_blocks": [{"thinking": REASONING}], + "reasoningContentBlocks": [{"reasoningText": {"text": REASONING}}], + "citations": [{"cited_text": INPUT}], + "thought_signatures": [SIGNATURE], + }, + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + if body.get("stream"): + choice: Final = payload["choices"][0] + chunk: Final = { + **payload, + "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": choice["message"], "finish_reason": None}], + } + finish: Final = { + **payload, + "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + } + return Reply( + content_type="text/event-stream", + chunks=tuple(f"data: {json.dumps(event)}\n\n".encode() for event in (chunk, finish)) + + (b"data: [DONE]\n\n",), + ) + return Reply(body=json.dumps(payload).encode()) + + +@pytest.fixture(scope="module") +def redaction_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Wire, Wire]]: + tmp_path: Final = tmp_path_factory.mktemp("custom-logger-redaction") + with ( + wire_server(_provider) as provider, + wire_server(lambda request: Reply()) as sink, + gateway_from_environment() as rig, + ): + config: Final = tmp_path / "redaction.yaml" + config.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": name, + "litellm_params": { + "model": model, + "api_base": provider.url if model.startswith("anthropic/") else provider.url + "/v1", + "api_key": "synthetic-provider-key", + }, + } + for name, model in ( + ("redaction-probe", "openai/gpt-4o-mini"), + ("redaction-text", "openai/gpt-3.5-turbo-instruct"), + ("redaction-embedding", "openai/text-embedding-3-large"), + ("redaction-responses", "openai/gpt-5.6-sol"), + ("redaction-messages", "anthropic/claude-sonnet-5-5"), + ) + ] + + [ + { + "model_name": "redaction-usage", + "litellm_params": { + "model": "anthropic/claude-sonnet-5-5", + "api_base": provider.url, + "api_key": "synthetic-provider-key", + "rpm": 100000, + "tpm": 100000000, + }, + "model_info": {"id": USAGE_DEPLOYMENT}, + } + ], + "litellm_settings": { + "callbacks": ["integration._support.redaction_probe.probe"], + "enable_redis_auth_cache": True, + "cache": True, + "cache_params": { + "type": "redis", + "host": "os.environ/REDIS_HOST", + "port": "os.environ/REDIS_PORT", + }, + }, + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "disable_model_info_refresh": True, + "store_prompts_in_spend_logs": True, + }, + "router_settings": {"redis_host": "os.environ/REDIS_HOST", "redis_port": "os.environ/REDIS_PORT"}, + } + ) + ) + with owned_proxy( + rig, tmp_path, {"REDACTION_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}, config=config, workers=2 + ) as proxy: + yield proxy, provider, sink + + +@pytest.mark.parametrize("redact", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +def test_custom_logger_redacts_nested_input_and_provider_reasoning( + redaction_proxy: tuple[Gateway, Wire, Wire], redact: bool, stream: bool +) -> None: + proxy, provider, sink = redaction_proxy + response: Final = proxy.request( + "POST", + "/v1/chat/completions", + { + "model": "redaction-probe", + "messages": [{"role": "user", "content": INPUT}], + "stream": stream, + "cache": {"no-cache": True}, + }, + headers={"x-litellm-enable-message-redaction": "true"} if redact else {}, + ) + assert response.status_code == 200, response.text + first: Final = ( + json.loads(next(line[6:] for line in response.text.splitlines() if line.startswith("data: "))) + if stream + else response.json() + ) + message: Final = first["choices"][0]["delta" if stream else "message"] + assert message["content"] == OUTPUT, message + assert REASONING in json.dumps(message["provider_specific_fields"]), message + received: Final = eventually(sink.drain, lambda values: len(values) == 1, seconds=20) + callback: Final = json.loads(received[0].body) + body: Final = callback["request"]["body"] + assert "proxy_server_request" not in body, body + assert len(tuple(request for request in provider.drain() if request.method == "POST")) == 1 + logged: Final = json.dumps(callback) + assert SIGNATURE in logged, callback + if redact: + assert INPUT not in logged and REASONING not in logged and OUTPUT not in logged, callback + else: + assert INPUT in logged and REASONING in logged and OUTPUT in logged, callback + + +@pytest.mark.parametrize("redact", [False, True]) +@pytest.mark.parametrize( + ("route", "field", "model"), + [("/v1/completions", "prompt", "redaction-text"), ("/v1/embeddings", "input", "redaction-embedding")], +) +def test_custom_logger_redacts_prompt_and_embedding_input( + redaction_proxy: tuple[Gateway, Wire, Wire], route: str, field: str, model: str, redact: bool +) -> None: + proxy, provider, sink = redaction_proxy + response: Final = proxy.request( + "POST", + route, + {"model": model, field: INPUT, "cache": {"no-cache": True}}, + headers={"x-litellm-enable-message-redaction": "true"} if redact else {}, + ) + assert response.status_code == 200, response.text + if field == "input": + assert response.json()["data"][0]["embedding"] == [0.25, 0.5] + else: + assert response.json()["choices"][0]["text"] == OUTPUT + received: Final = eventually(sink.drain, bool, seconds=20) + assert len(received) == 1, received + callback: Final = json.loads(received[0].body) + assert "proxy_server_request" not in callback["request"]["body"] + assert len(tuple(request for request in provider.drain() if request.method == "POST")) == 1 + assert (INPUT in json.dumps(callback)) is not redact, callback + + +@pytest.mark.parametrize("redact", [False, True]) +@pytest.mark.parametrize("native_messages", [False, True]) +def test_native_messages_and_responses_logging_redaction( + redaction_proxy: tuple[Gateway, Wire, Wire], native_messages: bool, redact: bool +) -> None: + proxy, provider, sink = redaction_proxy + body: Final = ( + {"model": "redaction-messages", "messages": [{"role": "user", "content": INPUT}], "max_tokens": 64} + if native_messages + else {"model": "redaction-responses", "input": INPUT} + ) + response: Final = proxy.request( + "POST", + "/v1/messages" if native_messages else "/v1/responses", + {**body, "cache": {"no-cache": True}}, + headers={"x-litellm-enable-message-redaction": "true"} if redact else {}, + ) + assert response.status_code == 200, response.text + if native_messages: + assert response.json()["content"][1]["text"] == OUTPUT + assert response.json()["content"][0]["thinking"] == REASONING + else: + assert response.json()["output"][0]["content"][0]["text"] == OUTPUT + received: Final = eventually(sink.drain, bool, seconds=20) + assert len(received) == 1, received + logged: Final = received[0].body.decode() + assert len(tuple(request for request in provider.drain() if request.method == "POST")) == 1 + assert (INPUT in logged) is not redact, logged + assert (OUTPUT in logged) is not redact, logged + if native_messages: + assert (REASONING in logged) is not redact, logged + + +def test_concurrent_redacted_and_unredacted_callbacks_remain_isolated( + redaction_proxy: tuple[Gateway, Wire, Wire], +) -> None: + proxy, provider, sink = redaction_proxy + users: Final = tuple(f"redaction-chaos-{index}" for index in range(4)) + + def send(index: int): + return proxy.request( + "POST", + "/v1/chat/completions", + { + "model": "redaction-probe", + "messages": [{"role": "user", "content": INPUT}], + "user": users[index], + "cache": {"no-cache": True}, + }, + headers={"x-litellm-enable-message-redaction": "true"} if index % 2 else {}, + ) + + with ThreadPoolExecutor(max_workers=4) as pool: + responses: Final = tuple(pool.map(send, range(4))) + for response in responses: + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == OUTPUT + assert REASONING in json.dumps(response.json()["choices"][0]["message"]) + batches: Final[list[Request]] = [] # mutable-ok: drain consumes the sink queue, so polls retain previous events + + def observed() -> tuple[Request, ...]: + batches.extend(sink.drain()) + return tuple(batches) + + events: Final = eventually(observed, lambda values: len(values) == 4, seconds=20) + callbacks: Final = tuple(json.loads(event.body) for event in events) + assert {callback["request"]["body"]["user"] for callback in callbacks} == set(users) + for callback in callbacks: + index: Final = users.index(callback["request"]["body"]["user"]) + logged: Final = json.dumps(callback) + assert (INPUT in logged) is (index % 2 == 0), callback + assert (REASONING in logged) is (index % 2 == 0), callback + assert (OUTPUT in logged) is (index % 2 == 0), callback + assert len(tuple(request for request in provider.drain() if request.method == "POST")) == 4 + + +def _drained(wire: Wire, count: int) -> tuple[Request, ...]: + batches: Final[list[Request]] = [] # mutable-ok: drain consumes the queue, so polls retain previous events + + def observed() -> tuple[Request, ...]: + batches.extend(wire.drain()) + return tuple(batches) + + return eventually(observed, lambda values: len(values) >= count, seconds=30) + + +def test_redacted_responses_follow_up_replays_original_instructions( + redaction_proxy: tuple[Gateway, Wire, Wire], +) -> None: + proxy, provider, sink = redaction_proxy + end_user: Final = f"redaction-replay-{uuid.uuid4().hex}" + first: Final = proxy.request( + "POST", + "/v1/responses", + { + "model": "redaction-messages", + "instructions": INSTRUCTIONS, + "input": f"{BRIDGED} turn one", + "user": end_user, + "cache": {"no-cache": True}, + }, + headers=REDACT, + ) + assert first.status_code == 200, first.text + eventually( + lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE end_user = %s', (end_user,)), + lambda rows: len(rows) == 1, + seconds=70, + ) + assert len(_drained(provider, 1)) == 1 + follow_up: Final = proxy.request( + "POST", + "/v1/responses", + { + "model": "redaction-messages", + "previous_response_id": first.json()["id"], + "input": f"{BRIDGED} turn two", + "user": end_user, + "cache": {"no-cache": True}, + }, + headers=REDACT, + ) + assert follow_up.status_code == 200, follow_up.text + sent: Final = tuple(json.loads(request.body) for request in _drained(provider, 1)) + assert len(_drained(sink, 2)) == 2 + assert len(sent) == 1, sent + assert len(sent[0]["messages"]) == 3, sent[0] + system: Final = json.dumps(sent[0].get("system")) + assert INSTRUCTIONS in system and REDACTED_BY_LITELLM not in system, sent[0] + + +def test_redacted_router_usage_is_counted_once(redaction_proxy: tuple[Gateway, Wire, Wire]) -> None: + proxy, provider, sink = redaction_proxy + requests: Final = 48 + + def send(index: int): + return proxy.request( + "POST", + "/v1/chat/completions", + { + "model": "redaction-usage", + "messages": [{"role": "user", "content": f"{BRIDGED} usage {index}"}], + "cache": {"no-cache": True}, + }, + headers=REDACT, + ) + + with ThreadPoolExecutor(max_workers=16) as pool: + responses: Final = tuple(pool.map(send, range(requests))) + for response in responses: + assert response.status_code == 200, response.text + assert response.json()["usage"]["total_tokens"] == 15, response.json() + + def counted_tokens() -> int: + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache: + keys: Final = tuple(cache.scan_iter(f"global_router:{USAGE_DEPLOYMENT}:*:tpm:*")) + return sum(int(float(cache.get(key) or 0)) for key in keys) + + _drained(sink, requests) + held(counted_tokens, lambda total: total == 15 * requests, holding=3, seconds=30) + assert len(_drained(provider, requests)) == requests + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("redact", [False, True]) +async def test_sdk_success_callback_redaction_over_http( + monkeypatch: pytest.MonkeyPatch, asynchronous: bool, redact: bool +) -> None: + with wire_server(_provider) as provider, wire_server(lambda request: Reply()) as sink: + monkeypatch.setenv("REDACTION_PROBE_URL", sink.url) + messages: Final = [{"role": "user", "content": INPUT}] + request: Final = {"body": {"messages": messages, "model": "gpt-4o-mini"}} + callback: Final = RedactionProbe().async_log_success_event if asynchronous else SyncRedactionProbe() + parameters: Final = { + "model": "openai/gpt-4o-mini", + "messages": messages, + "api_base": provider.url + "/v1", + "api_key": "synthetic-sdk-provider-key", + "proxy_server_request": request, + "turn_off_message_logging": redact, + "success_callback": [callback], + } + response: Final = ( + await litellm.acompletion(**parameters) + if asynchronous + else await asyncio.to_thread(litellm.completion, **parameters) + ) + assert response.choices[0].message.content == OUTPUT + assert REASONING in json.dumps(response.choices[0].message.provider_specific_fields) + await GLOBAL_LOGGING_WORKER.flush() + delivered: Final = await asyncio.to_thread(eventually, sink.drain, bool, seconds=20) + assert len(delivered) == 1, delivered + logged: Final = delivered[0].body.decode() + assert request["body"]["messages"] == messages + if redact: + assert INPUT not in logged and REASONING not in logged and OUTPUT not in logged, logged + else: + assert INPUT in logged and REASONING in logged and OUTPUT in logged, logged diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 6fbcba170f1..d14f236f765 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -6,6 +6,7 @@ but litellm_params["litellm_metadata"] is None. """ import asyncio +import copy import importlib import json import os @@ -36,12 +37,131 @@ from litellm.responses.main import mock_responses_api_response from litellm.types.router import BaselineRouteStamp from litellm.types.utils import ( ModelResponse, + ModelResponseStream, ResponsesAPIResponse, StandardLoggingPayload, TextCompletionResponse, ) from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome +SECRET: Final = "private-redaction-canary" +SIGNATURE: Final = "opaque-thought-signature" + + +def _provider_fields() -> dict[str, object]: + return { + "reasoning_content": SECRET, + "thinking_blocks": [{"type": "thinking", "thinking": SECRET}], + "reasoningContentBlocks": [{"reasoningText": {"text": SECRET}}], + "citations": [{"cited_text": SECRET}], + "thought_signatures": [SIGNATURE], + } + + +def _request_params(metadata: dict[str, object]) -> dict[str, object]: + return { + "metadata": metadata, + "proxy_server_request": { + "url": "http://localhost/v1/chat/completions", + "body": { + "model": "redaction-model", + "system": "kept-system", + "messages": [{"role": "user", "content": SECRET}], + "metadata": metadata, + }, + }, + } + + +@pytest.mark.parametrize("shape", ["object", "dict_message", "dict_delta"]) +def test_redaction_reuses_message_rules_for_request_body_and_provider_fields(shape: str) -> None: + response: Final = ModelResponse( + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": SECRET, "provider_specific_fields": _provider_fields()}, + } + ] + ) + result: Final = ( + response + if shape == "object" + else {"choices": [{"index": 0, "delta": {"content": SECRET, "provider_specific_fields": _provider_fields()}}]} + if shape == "dict_delta" + else response.model_dump() + ) + original_result: Final = copy.deepcopy(result) + params: Final = _request_params({"session_id": "session"}) + original_params: Final = copy.deepcopy(params) + details: Final[dict[str, object]] = {"litellm_params": params} + + redacted: Final = perform_redaction(details, result) + + assert SECRET not in str(redacted) and SECRET not in str(details), (redacted, details) + assert SIGNATURE in str(redacted) + body: Final = details["litellm_params"]["proxy_server_request"]["body"] + assert body["messages"] == [{"role": "user", "content": "redacted-by-litellm"}] + assert body["system"] == "kept-system" and body["model"] == "redaction-model" + assert body["metadata"] is params["metadata"] + assert result == original_result and params == original_params + + +def test_redaction_redacts_reasoning_copied_into_hidden_params_in_place() -> None: + hidden_params: Final = {"model_id": "deployment", "provider_specific_fields": _provider_fields()} + original_hidden_params: Final = copy.deepcopy(hidden_params) + metadata: Final[dict[str, object]] = {"session_id": "session", "hidden_params": hidden_params} + params: Final = _request_params(metadata) + details: Final[dict[str, object]] = {"litellm_params": params} + + perform_redaction(details, None) + + assert details["litellm_params"]["metadata"] is metadata + assert SECRET not in str(details), details + assert metadata["hidden_params"]["model_id"] == "deployment" + assert metadata["hidden_params"]["provider_specific_fields"]["thought_signatures"] == [SIGNATURE] + assert hidden_params == original_hidden_params + + +def test_per_callback_redaction_leaves_shared_request_body_and_hidden_params() -> None: + metadata: Final[dict[str, object]] = {"hidden_params": {"provider_specific_fields": _provider_fields()}} + params: Final = _request_params(metadata) + original_params: Final = copy.deepcopy(params) + details: Final[dict[str, object]] = {"litellm_params": params} + + perform_redaction(details, None, redact_streaming_responses=False) + + assert details["litellm_params"] == original_params + + +def test_logger_level_redaction_redacts_its_copy_of_the_request_body() -> None: + params: Final = _request_params({"session_id": "session"}) + original_params: Final = copy.deepcopy(params) + details: Final[dict[str, object]] = {"litellm_params": params} + + redacted: Final = CustomLogger( + turn_off_message_logging=True + ).redact_standard_logging_payload_from_model_call_details(details) + + assert SECRET not in str(redacted["litellm_params"]), redacted + assert details["litellm_params"] == original_params + + +@pytest.mark.parametrize("key", ["complete_streaming_response", "async_complete_streaming_response"]) +def test_streaming_redaction_redacts_delta_provider_fields_without_touching_the_chunk_dict(key: str) -> None: + fields: Final = _provider_fields() + response: Final = ModelResponseStream( + choices=[ + litellm.utils.StreamingChoices(delta=litellm.utils.Delta(content=SECRET, provider_specific_fields=fields)) + ] + ) + details: Final[dict[str, object]] = {"stream": True, key: response} + + perform_redaction(details, None) + + assert SECRET not in str(response.choices[0].delta), response + assert response.choices[0].delta.provider_specific_fields["thought_signatures"] == [SIGNATURE] + assert fields == _provider_fields() + @pytest.fixture(autouse=True) def _reset_global_redaction(): From 1587faeeb99081af9d12a0070c106cbd93859ae6 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 13 Oct 2026 12:41:06 -0700 Subject: [PATCH 006/325] test(unit): wait on the work itself instead of fixed sleeps (#46261) * test(unit): wait on the work itself instead of fixed sleeps in 21 flaky files Replace sleeps that raced background logging, cache writes, refreshes and thread work with either awaiting that work directly (logging worker flush, the batch logger's flush_queue, the tasks the call spawned, thread joins) or an Event the test's own fake sets when it runs. Each wait keeps a 5s hang guard that never decides pass or fail. Tests that need a fake clock, real sockets, or measure latency are left unchanged, and no production code moves * test(unit): await the tasks each completion spawns instead of yielding once The logging tests yielded once with asyncio.sleep(0) so the task that enqueues the success log could run before the worker flush. That relied on the enqueue happening before the task's first await. Snapshot the running tasks around each acompletion call and await the ones it spawned, so the guarantee holds however many steps sit before the enqueue --- .../unit/a2a_protocol/test_cost_calculator.py | 20 ++- tests/unit/caching/test_redis_cache.py | 27 ++-- .../unit/integrations/datadog/test_datadog.py | 31 +++-- .../generic_api/test_generic_api_callback.py | 121 +++++++++++------- tests/unit/integrations/opik/test_opik.py | 44 +++++-- .../otel/test_otel_v2_destinations.py | 40 +++--- tests/unit/integrations/test_mlflow.py | 12 +- .../test_prometheus_metrics_endpoint.py | 24 ++-- tests/unit/integrations/test_s3_v2.py | 16 ++- .../test_interactions_streaming_iterator.py | 4 +- .../test_streaming_handler.py | 52 ++++---- .../unit/litellm_core_utils/test_url_utils.py | 27 ++-- .../test_vertex_gemini_gcs_uri_mime.py | 38 +++--- .../llms/vertex_ai/test_vertex_llm_base.py | 13 +- .../test_auth_cache_invalidation_pubsub.py | 18 +-- .../test_custom_code_bounded_execution.py | 28 +++- ...test_proxy_server_endpoints_and_startup.py | 50 ++++++-- .../test_responses_streaming_iterator.py | 11 +- .../unit/responses/test_streaming_iterator.py | 34 +++-- tests/unit/test_litellm_overhead.py | 3 +- tests/unit/test_utils.py | 27 ++-- 21 files changed, 393 insertions(+), 247 deletions(-) diff --git a/tests/unit/a2a_protocol/test_cost_calculator.py b/tests/unit/a2a_protocol/test_cost_calculator.py index 8d8ec815f3a..212af4a6760 100644 --- a/tests/unit/a2a_protocol/test_cost_calculator.py +++ b/tests/unit/a2a_protocol/test_cost_calculator.py @@ -117,6 +117,7 @@ class CostLogger(CustomLogger): def __init__(self): self.response_cost: Optional[float] = None + self.logged = asyncio.Event() super().__init__() async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -125,6 +126,7 @@ class CostLogger(CustomLogger): self.response_cost = ( slp.get("response_cost") if isinstance(slp, dict) else getattr(slp, "response_cost", None) ) + self.logged.set() @pytest.mark.asyncio @@ -157,7 +159,7 @@ async def test_asend_message_uses_cost_per_query(monkeypatch): cost_per_query=0.05, ) - await asyncio.sleep(0.1) + assert await asyncio.wait_for(cost_logger.logged.wait(), timeout=5) assert cost_logger.response_cost == 0.05 @@ -194,7 +196,7 @@ async def test_asend_message_uses_cost_per_query_from_litellm_params_dict(monkey }, ) - await asyncio.sleep(0.1) + assert await asyncio.wait_for(cost_logger.logged.wait(), timeout=5) assert cost_logger.response_cost == 0.5 @@ -206,6 +208,7 @@ class TokenAndCostLogger(CustomLogger): self.response_cost: Optional[float] = None self.prompt_tokens: Optional[int] = None self.completion_tokens: Optional[int] = None + self.logged = asyncio.Event() super().__init__() async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -220,6 +223,7 @@ class TokenAndCostLogger(CustomLogger): self.completion_tokens = ( slp.get("completion_tokens") if isinstance(slp, dict) else getattr(slp, "completion_tokens", None) ) + self.logged.set() @pytest.mark.asyncio @@ -257,7 +261,7 @@ async def test_asend_message_uses_input_output_cost_per_token(monkeypatch): output_cost_per_token=output_cost_per_token, ) - await asyncio.sleep(0.1) + assert await asyncio.wait_for(token_cost_logger.logged.wait(), timeout=5) # Get actual token counts from logger prompt_tokens = token_cost_logger.prompt_tokens @@ -290,11 +294,13 @@ class AgentIdLogger(CustomLogger): def __init__(self): self.agent_id: Optional[str] = None self.kwargs: Optional[dict] = None + self.logged = asyncio.Event() super().__init__() async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): self.kwargs = kwargs self.agent_id = kwargs.get("agent_id") + self.logged.set() @pytest.mark.asyncio @@ -329,7 +335,7 @@ async def test_asend_message_passes_agent_id_to_callback(monkeypatch): agent_id=test_agent_id, ) - await asyncio.sleep(0.1) + assert await asyncio.wait_for(agent_id_logger.logged.wait(), timeout=5) # Verify agent_id was passed to callback assert agent_id_logger.agent_id == test_agent_id, ( @@ -346,6 +352,7 @@ class MetadataLogger(CustomLogger): self.user_api_key: Optional[str] = None self.user_id: Optional[str] = None self.team_id: Optional[str] = None + self.logged = asyncio.Event() super().__init__() async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -354,6 +361,7 @@ class MetadataLogger(CustomLogger): self.user_api_key = self.metadata.get("user_api_key") self.user_id = self.metadata.get("user_api_key_user_id") self.team_id = self.metadata.get("user_api_key_team_id") + self.logged.set() @pytest.mark.asyncio @@ -396,7 +404,7 @@ async def test_asend_message_streaming_propagates_metadata(): ): chunks.append(chunk) - await asyncio.sleep(0.2) + assert await asyncio.wait_for(metadata_logger.logged.wait(), timeout=5) # Verify metadata was propagated to callback assert metadata_logger.user_api_key == "sk-test-key-hash-12345" @@ -439,7 +447,7 @@ async def test_asend_message_streaming_triggers_callbacks(): ): chunks.append(chunk) - await asyncio.sleep(0.2) + assert await asyncio.wait_for(callback_logger.logged.wait(), timeout=5) # Verify chunks were received assert len(chunks) == 2 diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 57c32ace509..0a041694d52 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -573,9 +573,10 @@ async def test_sync_batch_get_cache_reports_a_failed_read_from_a_running_loop( outage while the success path keeps reporting, so the dashboards read healthy. """ cache, service_logger = sync_batch_cache_with_service_logger + tasks_before_read = asyncio.all_tasks() assert cache.batch_get_cache(key_list=["lit6729"]) == {} - await asyncio.sleep(0.05) + await asyncio.wait_for(asyncio.gather(*(asyncio.all_tasks() - tasks_before_read)), timeout=5) assert service_logger.mock_testing_sync_failure_hook == 1 assert service_logger.mock_testing_async_failure_hook == 1 @@ -779,21 +780,29 @@ async def test_concurrent_success_is_not_cancelled_by_another_calls_failure(): # The failure has to land after both calls are already in flight, which is the only # ordering where a shared counter confuses the two. Failing before the healthy call # starts would leave its snapshot correct and prove nothing. - async def swallows_a_failure(): - await asyncio.sleep(0.02) - _record_swallowed_redis_failure(breaker, RedisConnectionError("redis unreachable")) + async def run_one_round() -> None: + healthy_call_in_flight = asyncio.Event() + failure_recorded = asyncio.Event() - async def succeeds_while_the_other_fails(): - await asyncio.sleep(0.05) - return "ok" + async def swallows_a_failure() -> None: + await asyncio.wait_for(healthy_call_in_flight.wait(), timeout=5) + _record_swallowed_redis_failure(breaker, RedisConnectionError("redis unreachable")) + failure_recorded.set() + + async def succeeds_while_the_other_fails() -> str: + healthy_call_in_flight.set() + await asyncio.wait_for(failure_recorded.wait(), timeout=5) + return "ok" - rounds = breaker.failure_threshold + 1 - for _ in range(rounds): await asyncio.gather( _run_under_circuit_breaker(breaker, "failing", swallows_a_failure), _run_under_circuit_breaker(breaker, "healthy", succeeds_while_the_other_fails), ) + rounds = breaker.failure_threshold + 1 + for _ in range(rounds): + await run_one_round() + assert breaker._failure_count < breaker.failure_threshold, "the healthy call must clear the streak" assert breaker.is_open() is False, "a Redis answering every round must stay in the pool" diff --git a/tests/unit/integrations/datadog/test_datadog.py b/tests/unit/integrations/datadog/test_datadog.py index eec0022c1db..14cd376cbc8 100644 --- a/tests/unit/integrations/datadog/test_datadog.py +++ b/tests/unit/integrations/datadog/test_datadog.py @@ -3,7 +3,7 @@ import json import os from datetime import datetime from pathlib import Path -from typing import Coroutine, Final, TypedDict +from typing import Awaitable, Coroutine, Final, TypedDict, TypeVar from unittest.mock import AsyncMock, patch import pytest @@ -41,6 +41,17 @@ STANDARD_START_TIME: Final[datetime] = datetime(2025, 1, 1) STANDARD_END_TIME: Final[datetime] = datetime(2025, 1, 1, 0, 0, 1) +_T = TypeVar("_T") + + +async def _await_with_spawned_tasks(call: Awaitable[_T]) -> _T: + before: Final = asyncio.all_tasks() + result: Final = await call + spawned: Final = asyncio.all_tasks() - before - {asyncio.current_task()} + await asyncio.wait_for(asyncio.gather(*spawned), timeout=5) + return result + + def _discard_periodic_flush(coroutine: Coroutine[object, object, None]) -> None: coroutine.close() @@ -267,17 +278,19 @@ async def test_datadog_logging_http_request(): # Make the completion call for _ in range(5): - response = await litellm.acompletion( - model="gpt-4.1-mini", - messages=[{"role": "user", "content": "what llm are u"}], - max_tokens=10, - temperature=0.2, - mock_response="Accepted", + response = await _await_with_spawned_tasks( + litellm.acompletion( + model="gpt-4.1-mini", + messages=[{"role": "user", "content": "what llm are u"}], + max_tokens=10, + temperature=0.2, + mock_response="Accepted", + ) ) print(response) - # Wait for 5 seconds - await asyncio.sleep(6) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=5) + await dd_logger.flush_queue() # Assert that the mock was called assert mock_post.called, "HTTP request was not made" diff --git a/tests/unit/integrations/generic_api/test_generic_api_callback.py b/tests/unit/integrations/generic_api/test_generic_api_callback.py index e3bb47cad02..7a712786a84 100644 --- a/tests/unit/integrations/generic_api/test_generic_api_callback.py +++ b/tests/unit/integrations/generic_api/test_generic_api_callback.py @@ -4,7 +4,7 @@ import json import logging import os from collections.abc import AsyncIterator, Iterator, Sequence -from typing import Final, TypedDict +from typing import Awaitable, Final, TypedDict, TypeVar from unittest.mock import AsyncMock import httpx @@ -24,6 +24,17 @@ from litellm.types.utils import StandardLoggingPayload from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome +_T = TypeVar("_T") + + +async def _await_with_spawned_tasks(call: Awaitable[_T]) -> _T: + before: Final = asyncio.all_tasks() + result: Final = await call + spawned: Final = asyncio.all_tasks() - before - {asyncio.current_task()} + await asyncio.wait_for(asyncio.gather(*spawned), timeout=5) + return result + + @pytest.fixture(autouse=True) def set_verbose_logger_level() -> Iterator[None]: original_level = verbose_logger.level @@ -51,20 +62,22 @@ async def test_generic_api_callback(): os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint # Initialize the GenericAPILogger and set the mock - generic_logger = GenericAPILogger(endpoint=test_endpoint, headers=test_headers, flush_interval=1) + generic_logger = GenericAPILogger(endpoint=test_endpoint, headers=test_headers, flush_interval=3600) generic_logger.async_httpx_client.post = mock_post litellm.callbacks = [generic_logger] # Make the completion call - response = await litellm.acompletion( - model="gpt-5.5", - messages=[{"role": "user", "content": "Hello, world!"}], - mock_response="hi", - user="test_user", + response = await _await_with_spawned_tasks( + litellm.acompletion( + model="gpt-5.5", + messages=[{"role": "user", "content": "Hello, world!"}], + mock_response="hi", + user="test_user", + ) ) - # Wait for async flush - await asyncio.sleep(3) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS) + await generic_logger.flush_queue() # Assert httpx post was called mock_post.assert_called_once() @@ -130,21 +143,23 @@ async def test_generic_api_callback_multiple_logs(): os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint # Initialize the GenericAPILogger and set the mock - generic_logger = GenericAPILogger(endpoint=test_endpoint, headers=test_headers, flush_interval=5) + generic_logger = GenericAPILogger(endpoint=test_endpoint, headers=test_headers, flush_interval=3600) generic_logger.async_httpx_client.post = mock_post litellm.callbacks = [generic_logger] # Make the completion call for _ in range(10): - response = await litellm.acompletion( - model="gpt-5.5", - messages=[{"role": "user", "content": "Hello, world!"}], - mock_response="hi", - user="test_user", + response = await _await_with_spawned_tasks( + litellm.acompletion( + model="gpt-5.5", + messages=[{"role": "user", "content": "Hello, world!"}], + mock_response="hi", + user="test_user", + ) ) - # Wait for async flush - await asyncio.sleep(6) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS) + await generic_logger.flush_queue() # Assert httpx post was called mock_post.assert_called_once() @@ -208,7 +223,7 @@ async def test_generic_api_callback_ndjson_format(): generic_logger = GenericAPILogger( endpoint=test_endpoint, headers=test_headers, - flush_interval=1, + flush_interval=3600, log_format="ndjson", # Set NDJSON format ) generic_logger.async_httpx_client.post = mock_post @@ -216,15 +231,17 @@ async def test_generic_api_callback_ndjson_format(): # Make multiple completion calls to generate multiple logs for i in range(3): - response = await litellm.acompletion( - model="gpt-5.5", - messages=[{"role": "user", "content": f"Hello, world! {i}"}], - mock_response="hi", - user="test_user", + response = await _await_with_spawned_tasks( + litellm.acompletion( + model="gpt-5.5", + messages=[{"role": "user", "content": f"Hello, world! {i}"}], + mock_response="hi", + user="test_user", + ) ) - # Wait for async flush - await asyncio.sleep(3) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS) + await generic_logger.flush_queue() # Assert httpx post was called mock_post.assert_called_once() @@ -277,7 +294,7 @@ async def test_generic_api_callback_single_format(): generic_logger = GenericAPILogger( endpoint=test_endpoint, headers=test_headers, - flush_interval=1, # Quick flush to trigger batch send + flush_interval=3600, log_format="single", # Set single format ) generic_logger.async_httpx_client.post = mock_post @@ -285,15 +302,17 @@ async def test_generic_api_callback_single_format(): # Make 3 completion calls for i in range(3): - response = await litellm.acompletion( - model="gpt-5.5", - messages=[{"role": "user", "content": f"Hello, world! {i}"}], - mock_response="hi", - user="test_user", + response = await _await_with_spawned_tasks( + litellm.acompletion( + model="gpt-5.5", + messages=[{"role": "user", "content": f"Hello, world! {i}"}], + mock_response="hi", + user="test_user", + ) ) - # Wait for async flush - await asyncio.sleep(3) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS) + await generic_logger.flush_queue() # Assert httpx post was called 3 times (once per log in batch) assert mock_post.call_count == 3, f"Expected 3 calls, got {mock_post.call_count}" @@ -336,7 +355,7 @@ async def test_generic_api_callback_json_array_format_explicit(): generic_logger = GenericAPILogger( endpoint=test_endpoint, headers=test_headers, - flush_interval=1, + flush_interval=3600, log_format="json_array", # Explicitly set json_array ) generic_logger.async_httpx_client.post = mock_post @@ -344,15 +363,17 @@ async def test_generic_api_callback_json_array_format_explicit(): # Make multiple completion calls for i in range(5): - response = await litellm.acompletion( - model="gpt-5.5", - messages=[{"role": "user", "content": f"Hello, world! {i}"}], - mock_response="hi", - user="test_user", + response = await _await_with_spawned_tasks( + litellm.acompletion( + model="gpt-5.5", + messages=[{"role": "user", "content": f"Hello, world! {i}"}], + mock_response="hi", + user="test_user", + ) ) - # Wait for async flush - await asyncio.sleep(3) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS) + await generic_logger.flush_queue() # Assert httpx post was called once (batched) mock_post.assert_called_once() @@ -386,7 +407,7 @@ async def test_generic_api_callback_sumologic_uses_ndjson(): os.environ["SUMOLOGIC_WEBHOOK_URL"] = "https://collectors.sumologic.com/receiver/v1/http/test123" # Initialize using callback_name (loads from JSON config) - generic_logger = GenericAPILogger(callback_name="sumologic", flush_interval=1) + generic_logger = GenericAPILogger(callback_name="sumologic", flush_interval=3600) generic_logger.async_httpx_client.post = mock_post litellm.callbacks = [generic_logger] @@ -395,15 +416,17 @@ async def test_generic_api_callback_sumologic_uses_ndjson(): # Make completion calls for i in range(2): - await litellm.acompletion( - model="gpt-5.5", - messages=[{"role": "user", "content": f"Test {i}"}], - mock_response="response", - user="test_user", + await _await_with_spawned_tasks( + litellm.acompletion( + model="gpt-5.5", + messages=[{"role": "user", "content": f"Test {i}"}], + mock_response="response", + user="test_user", + ) ) - # Wait for async flush - await asyncio.sleep(3) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS) + await generic_logger.flush_queue() # Assert httpx post was called mock_post.assert_called_once() diff --git a/tests/unit/integrations/opik/test_opik.py b/tests/unit/integrations/opik/test_opik.py index 1aad7079a4a..f7f5e252ed4 100644 --- a/tests/unit/integrations/opik/test_opik.py +++ b/tests/unit/integrations/opik/test_opik.py @@ -4,6 +4,7 @@ import logging import os import time from collections.abc import Iterator +from typing import Awaitable, Final, TypeVar from unittest.mock import AsyncMock, Mock import pytest @@ -15,6 +16,17 @@ from litellm.utils import _invalidate_model_cost_lowercase_map from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome +_T = TypeVar("_T") + + +async def _await_with_spawned_tasks(call: Awaitable[_T]) -> _T: + before: Final = asyncio.all_tasks() + result: Final = await call + spawned: Final = asyncio.all_tasks() - before - {asyncio.current_task()} + await asyncio.wait_for(asyncio.gather(*spawned), timeout=5) + return result + + @pytest.fixture(autouse=True) def isolate_opik_logging_state(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: original_level = verbose_logger.level @@ -54,27 +66,31 @@ async def test_opik_logging_http_request(): ] for _ in range(5): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Test message"}], - max_tokens=10, - temperature=0.2, - mock_response="This is a mock response", + await _await_with_spawned_tasks( + litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Test message"}], + max_tokens=10, + temperature=0.2, + mock_response="This is a mock response", + ) ) - await asyncio.sleep(1) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) assert opik_batch_calls() == [], "events below batch_size must stay queued" assert len(test_opik_logger.log_queue) == 10 for _ in range(3): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Test message"}], - max_tokens=10, - temperature=0.2, - mock_response="This is a mock response", + await _await_with_spawned_tasks( + litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Test message"}], + max_tokens=10, + temperature=0.2, + mock_response="This is a mock response", + ) ) - await asyncio.sleep(1) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) assert opik_batch_calls(), "crossing batch_size must flush the queue" events_left_over_after_the_size_triggered_flush = len(test_opik_logger.log_queue) diff --git a/tests/unit/integrations/otel/test_otel_v2_destinations.py b/tests/unit/integrations/otel/test_otel_v2_destinations.py index 068b2aa2e2e..c7a057cca5c 100644 --- a/tests/unit/integrations/otel/test_otel_v2_destinations.py +++ b/tests/unit/integrations/otel/test_otel_v2_destinations.py @@ -2,6 +2,7 @@ import asyncio import contextvars +import threading import time from base64 import b64encode from collections.abc import Mapping @@ -3496,9 +3497,11 @@ class TestEvictionSafety: def __init__(self): super().__init__(InMemorySpanExporter()) self.shutdown_calls = 0 + self.closed = threading.Event() def shutdown(self): self.shutdown_calls += 1 + self.closed.set() def _fan_out(self): built = [] @@ -3540,7 +3543,7 @@ class TestEvictionSafety: assert id(held) in fan_out._retired fan_out._release(held) - self._settle(fan_out, held) + assert held.closed.wait(timeout=5) assert held.shutdown_calls == 1 @@ -3554,7 +3557,7 @@ class TestEvictionSafety: fan_out._release(fan_out._acquire(self._dest(index))) fan_out._release(fan_out._acquire(self._dest(0))) fan_out._release(fan_out._acquire(self._dest(_MAX_CACHED_DESTINATION_PROCESSORS))) - self._settle(fan_out, built[1]) + assert built[1].closed.wait(timeout=5) assert built[1].shutdown_calls == 1 assert built[0].shutdown_calls == 0, "the destination used most recently was the one shed" @@ -3566,7 +3569,7 @@ class TestEvictionSafety: for index in range(_MAX_CACHED_DESTINATION_PROCESSORS + 1): fan_out._acquire(self._dest(index)) fan_out._release(built[-1]) - self._settle(fan_out, built[0]) + assert built[0].closed.wait(timeout=5) assert built[0].shutdown_calls == 1 assert len(fan_out._processors) == _MAX_CACHED_DESTINATION_PROCESSORS @@ -3577,9 +3580,11 @@ class TestEvictionSafety: every other tenant's spans.""" from litellm.integrations.otel.plumbing.providers import _MAX_CACHED_DESTINATION_PROCESSORS + release = threading.Event() + class Slow(self.Recording): def shutdown(self): - time.sleep(3) + release.wait(timeout=10) super().shutdown() built = [] @@ -3589,12 +3594,14 @@ class TestEvictionSafety: return built[-1] fan_out = TenantFanOutSpanProcessor(processor_factory=factory) - started = time.monotonic() - for index in range(_MAX_CACHED_DESTINATION_PROCESSORS + 1): - fan_out._acquire(self._dest(index)) - fan_out._release(built[-1]) + try: + for index in range(_MAX_CACHED_DESTINATION_PROCESSORS + 1): + fan_out._acquire(self._dest(index)) + fan_out._release(built[-1]) - assert time.monotonic() - started < 2 + assert not built[0].closed.is_set(), "the shed processor was closed inline on the export path" + finally: + release.set() def test_shedding_many_processors_does_not_spawn_a_thread_each(self): """A tenant that cycles its destination config sheds a processor per request, @@ -3970,11 +3977,7 @@ class TestEvictionSafety: fan_out.shutdown() fan_out._drain.submit(stray) - for _ in range(500): - if stray.shutdown_calls: - break - time.sleep(0.02) - + assert stray.closed.wait(timeout=5) assert stray.shutdown_calls == 1 def test_releasing_a_straggler_after_shutdown_does_not_block_the_span_thread(self): @@ -4036,7 +4039,7 @@ class TestEvictionSafety: assert held.shutdown_calls == 0 fan_out._release(held) - self._settle(fan_out, held) + assert held.closed.wait(timeout=5) assert held.shutdown_calls == 1 @@ -4153,10 +4156,7 @@ class TestEvictionSafety: close_returned.set() submitter.join(timeout=5) closer.join(timeout=5) - for _ in range(250): - if submitted.shutdown_calls: - break - time.sleep(0.02) + assert submitted.closed.wait(timeout=5), "a processor was queued behind the sentinels and never closed" assert submitted.shutdown_calls == 1, "a processor was queued behind the sentinels and never closed" @@ -4175,7 +4175,7 @@ class TestEvictionSafety: assert held.shutdown_calls == 0 fan_out._release(held) - self._settle(fan_out, held) + assert held.closed.wait(timeout=5) assert held.shutdown_calls == 1 diff --git a/tests/unit/integrations/test_mlflow.py b/tests/unit/integrations/test_mlflow.py index f828c34a9ff..058d0ae885d 100644 --- a/tests/unit/integrations/test_mlflow.py +++ b/tests/unit/integrations/test_mlflow.py @@ -1,5 +1,6 @@ import asyncio import json +import threading from datetime import datetime from unittest.mock import MagicMock, patch @@ -19,7 +20,13 @@ async def test_mlflow_logging_functionality(): mock_span = MagicMock() mock_span.parent_id = None # Simulate root trace mock_span.request_id = "test_trace_id" - mock_client.start_trace.return_value = mock_span + trace_started = threading.Event() + + def start_trace(*args: object, **kwargs: object) -> MagicMock: + trace_started.set() + return mock_span + + mock_client.start_trace.side_effect = start_trace # Mock all MLflow-related imports to avoid requiring MLflow as a dependency mock_mlflow_tracking = MagicMock() @@ -67,8 +74,7 @@ async def test_mlflow_logging_functionality(): }, ) - # Allow time for async processing - await asyncio.sleep(1) + assert await asyncio.to_thread(trace_started.wait, 5) # Verify start_trace was called with tags parameter assert mock_client.start_trace.called, "start_trace should have been called" diff --git a/tests/unit/integrations/test_prometheus_metrics_endpoint.py b/tests/unit/integrations/test_prometheus_metrics_endpoint.py index f0e6495ba22..a41508661d6 100644 --- a/tests/unit/integrations/test_prometheus_metrics_endpoint.py +++ b/tests/unit/integrations/test_prometheus_metrics_endpoint.py @@ -59,22 +59,18 @@ async def _scrape(client: httpx.AsyncClient, headers: Mapping[str, str] | None = @pytest.mark.asyncio async def test_render_does_not_block_the_event_loop(): - ticks: Final[list[float]] = [] # mutable-ok: records loop wakeups while the scrape is in flight - - async def ticker() -> None: - while True: - await asyncio.sleep(0.01) - ticks.append(time.monotonic()) - - ticker_task: Final = asyncio.create_task(ticker()) - async with _client(_registry_with(_SlowCollector(block_seconds=0.5))) as client: + collector: Final = _GatedCollector() + async with _client(_registry_with(collector)) as client: try: - response: Final = await _scrape(client) + scrape: Final = asyncio.create_task(_scrape(client)) + assert await asyncio.to_thread(collector.started.wait, _GATE_TIMEOUT_SECONDS), "render never started" + assert not collector.finished.is_set(), "event loop was blocked while the registry was rendered" + collector.release.set() + response: Final = await scrape finally: - ticker_task.cancel() + collector.release.set() - assert b"slow_metric" in response.content - assert len(ticks) > 5, "event loop was blocked while the registry was rendered" + assert b"gated_metric" in response.content @pytest.mark.asyncio @@ -170,6 +166,7 @@ class _GatedCollector(Collector): def __init__(self) -> None: self.started = threading.Event() self.release = threading.Event() + self.finished = threading.Event() self._lock = threading.Lock() self.collect_calls = 0 @@ -178,6 +175,7 @@ class _GatedCollector(Collector): self.collect_calls += 1 self.started.set() self.release.wait(timeout=_GATE_TIMEOUT_SECONDS) + self.finished.set() family: Final = GaugeMetricFamily("gated_metric", "gated") family.add_metric((), 1.0) yield family diff --git a/tests/unit/integrations/test_s3_v2.py b/tests/unit/integrations/test_s3_v2.py index 8e9a8982cd9..8449eae7bf4 100644 --- a/tests/unit/integrations/test_s3_v2.py +++ b/tests/unit/integrations/test_s3_v2.py @@ -4,6 +4,7 @@ import json import re import sys import textwrap +import threading import time import uuid from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine @@ -5291,12 +5292,15 @@ class _FakeS3Paginator: class _FakeS3Client: def __init__(self): self.objects = defaultdict(dict) + self.object_written = threading.Event() def clear(self): self.objects.clear() + self.object_written.clear() def put_object(self, Bucket, Key, Body, **_kwargs): self.objects[Bucket][Key] = Body + self.object_written.set() return {"ResponseMetadata": {"HTTPStatusCode": 200}} def delete_object(self, Bucket, Key): @@ -5363,7 +5367,7 @@ async def test_basic_s3_logging(sync_mode, streaming): response_id = chunk.id else: response_id = response.id - time.sleep(2) + assert _FAKE_S3_CLIENT.object_written.wait(timeout=5) else: response = await litellm.acompletion( model="gpt-5-mini", @@ -5376,7 +5380,7 @@ async def test_basic_s3_logging(sync_mode, streaming): response_id = chunk.id else: response_id = response.id - await asyncio.sleep(2) + assert await asyncio.to_thread(_FAKE_S3_CLIENT.object_written.wait, 5) total_objects, all_s3_keys = list_all_s3_objects("load-testing-oct") @@ -5411,9 +5415,11 @@ async def test_basic_s3_v2_logging(streaming): litellm.callbacks = [s3_v2_logger] uploaded_keys: list = [] + uploaded = asyncio.Event() async def mock_upload(batch_logging_element): uploaded_keys.append(batch_logging_element.s3_object_key) + uploaded.set() s3_v2_logger.async_upload_data_to_s3 = mock_upload @@ -5431,7 +5437,7 @@ async def test_basic_s3_v2_logging(streaming): else: response_id = response.id - await asyncio.sleep(5) + assert await asyncio.wait_for(uploaded.wait(), timeout=5) assert len(uploaded_keys) > 0, "S3 upload was never called" assert any(response_id in key for key in uploaded_keys), ( @@ -5465,6 +5471,7 @@ async def test_basic_s3_v2_logging_failure(): s3_v2_logger.async_httpx_client.put.return_value = mock_response upload_called = False + uploaded = asyncio.Event() async def mock_upload(batch_logging_element): nonlocal upload_called @@ -5474,6 +5481,7 @@ async def test_basic_s3_v2_logging_failure(): data = '{"model": "gpt-5-mini"}' await s3_v2_logger.async_httpx_client.put(url=url, headers=headers, data=data) + uploaded.set() s3_v2_logger.async_upload_data_to_s3 = mock_upload @@ -5496,7 +5504,7 @@ async def test_basic_s3_v2_logging_failure(): except Exception: pass - await asyncio.sleep(5) + assert await asyncio.wait_for(uploaded.wait(), timeout=5) assert upload_called, "S3 upload method was not called" diff --git a/tests/unit/interactions/test_interactions_streaming_iterator.py b/tests/unit/interactions/test_interactions_streaming_iterator.py index 9f88b2c9611..c52ab36deaa 100644 --- a/tests/unit/interactions/test_interactions_streaming_iterator.py +++ b/tests/unit/interactions/test_interactions_streaming_iterator.py @@ -90,8 +90,10 @@ async def test_custom_logger_only_never_submits_sync_success_handler(monkeypatch ) iterator.completed_response = InteractionsAPIStreamingResponse() + existing_tasks = asyncio.all_tasks() iterator._handle_logging_completed_response() - await asyncio.sleep(0.5) + spawned_tasks = asyncio.all_tasks() - existing_tasks + await asyncio.wait_for(asyncio.gather(*spawned_tasks, return_exceptions=True), timeout=5) assert recorder.async_hook_fired is True assert recording_executor.submitted_for(logging_obj) == [] diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 01fae0958e6..b60544e6606 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -1,5 +1,6 @@ import asyncio import json +import threading import time from typing import Final, NoReturn, Optional from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -548,9 +549,11 @@ async def test_streaming_with_usage_and_logging(sync_mode: bool): mock_callback, "async_log_stream_event" ) as mock_async_log_stream_event, ): + existing_tasks = asyncio.all_tasks() await test_streaming_handler_with_usage( sync_mode=sync_mode, final_usage_block=final_usage_block ) + spawned_tasks = asyncio.all_tasks() - existing_tasks if sync_mode: time.sleep(1) mock_log_success_event.assert_called_once() @@ -560,7 +563,7 @@ async def test_streaming_with_usage_and_logging(sync_mode: bool): == final_usage_block ) else: - await asyncio.sleep(1) + await asyncio.wait_for(asyncio.gather(*spawned_tasks, return_exceptions=True), timeout=5) mock_async_log_success_event.assert_called_once() # mock_async_log_stream_event.assert_called() assert ( @@ -750,10 +753,12 @@ async def test_streaming_completion_start_time(logging_obj: Logging): logging_obj=logging_obj, ) + existing_tasks = asyncio.all_tasks() async for chunk in response: print(chunk) - await asyncio.sleep(2) + spawned_tasks = asyncio.all_tasks() - existing_tasks + await asyncio.wait_for(asyncio.gather(*spawned_tasks, return_exceptions=True), timeout=5) assert logging_obj.model_call_details["completion_start_time"] is not None assert ( @@ -2952,15 +2957,19 @@ async def test_custom_stream_wrapper_anext_does_not_block_event_loop_for_sync_it """ class BlockingIterator: - def __init__(self, chunks, delay_s: float): + def __init__(self, chunks: list[ModelResponseStream]) -> None: self._it = iter(chunks) - self._delay_s = delay_s + self.entered = threading.Event() + self.release = threading.Event() + self.returned = threading.Event() - def __iter__(self): + def __iter__(self) -> "BlockingIterator": return self - def __next__(self): - time.sleep(self._delay_s) # simulate blocking I/O + def __next__(self) -> ModelResponseStream: + self.entered.set() + self.release.wait(timeout=5) + self.returned.set() return next(self._it) test_chunk = ModelResponseStream( @@ -2988,34 +2997,23 @@ async def test_custom_stream_wrapper_anext_does_not_block_event_loop_for_sync_it usage=None, ) - # Delay is intentionally > the wait_for timeout used to detect event loop blocking. + blocking_iterator = BlockingIterator([test_chunk]) wrapper = CustomStreamWrapper( - completion_stream=BlockingIterator([test_chunk], delay_s=0.3), + completion_stream=blocking_iterator, model="test-model", logging_obj=logging_obj, custom_llm_provider="cached_response", ) - tick_event = asyncio.Event() + anext_task = asyncio.create_task(wrapper.__anext__()) + try: + assert await asyncio.to_thread(blocking_iterator.entered.wait, 5) + assert not blocking_iterator.returned.is_set(), "Event loop was blocked by the sync iterator" + finally: + blocking_iterator.release.set() - async def background_tick(): - await asyncio.sleep(0.05) - tick_event.set() - - # Run the two coroutines concurrently and measure wall time. - # If __anext__ blocks the event loop, background_tick can't run and the gather - # takes the full 0.3 s delay; if non-blocking both finish within ~0.35 s total. - start = asyncio.get_event_loop().time() - - out, _ = await asyncio.gather( - wrapper.__anext__(), - background_tick(), - ) - - elapsed = asyncio.get_event_loop().time() - start + out = await anext_task assert isinstance(out, ModelResponseStream) - # background_tick sleeps 0.05 s; total must finish well under 2 × 0.3 s - assert elapsed < 0.5, f"Event loop was likely blocked (elapsed={elapsed:.2f}s)" @pytest.mark.asyncio diff --git a/tests/unit/litellm_core_utils/test_url_utils.py b/tests/unit/litellm_core_utils/test_url_utils.py index fccdc1a2a0a..46bb91959c0 100644 --- a/tests/unit/litellm_core_utils/test_url_utils.py +++ b/tests/unit/litellm_core_utils/test_url_utils.py @@ -1,7 +1,6 @@ import asyncio import socket import threading -import time import httpx import pytest @@ -542,31 +541,33 @@ def test_assert_same_origin_error_message_does_not_leak_hostnames(): async def test_async_safe_get_resolves_dns_off_the_event_loop(monkeypatch): + loop = asyncio.get_running_loop() loop_thread = threading.current_thread() resolver_threads = [] + resolver_blocked = asyncio.Event() + loop_ran_while_resolver_blocked = threading.Event() + released_by_the_loop = [] - def slow_getaddrinfo(host, port, *args, **kwargs): + def blocking_getaddrinfo(host, port, *args, **kwargs): resolver_threads.append(threading.current_thread()) - time.sleep(0.4) + loop.call_soon_threadsafe(resolver_blocked.set) + released_by_the_loop.append(loop_ran_while_resolver_blocked.wait(timeout=5)) return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", port or 443))] - monkeypatch.setattr(url_utils.socket, "getaddrinfo", slow_getaddrinfo) + monkeypatch.setattr(url_utils.socket, "getaddrinfo", blocking_getaddrinfo) class FakeClient: async def get(self, url, **kwargs): return httpx.Response(200, request=httpx.Request("GET", url)) - ticks = [time.perf_counter()] + async def release_resolver_from_the_loop() -> None: + await resolver_blocked.wait() + loop_ran_while_resolver_blocked.set() - async def heartbeat(): - while True: - await asyncio.sleep(0.01) - ticks.append(time.perf_counter()) - - beating = asyncio.create_task(heartbeat()) + releaser = asyncio.create_task(release_resolver_from_the_loop()) response = await url_utils.async_safe_get(FakeClient(), "https://img.example/a.png") - beating.cancel() + await asyncio.wait_for(releaser, timeout=5) assert response.status_code == 200 assert resolver_threads and all(thread is not loop_thread for thread in resolver_threads) - assert max(b - a for a, b in zip(ticks, ticks[1:])) < 0.2 + assert released_by_the_loop and all(released_by_the_loop) diff --git a/tests/unit/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py b/tests/unit/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py index 3906528ae2e..b60da3a7548 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py +++ b/tests/unit/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py @@ -3,7 +3,7 @@ Split from test_vertex.py to satisfy CI per-file size limits. """ import asyncio -import time +import threading from dotenv import load_dotenv @@ -265,8 +265,10 @@ def test_async_transform_request_body_offloads_extensionless_gs_not_plain_text() } ] - def slow_http_get(*args, **kwargs): - time.sleep(0.5) + http_get_released_by_the_loop = threading.Event() + released_while_blocked = [] + + def http_get_response() -> MagicMock: response = MagicMock() response.is_error = False response.status_code = 200 @@ -280,14 +282,22 @@ def test_async_transform_request_body_offloads_extensionless_gs_not_plain_text() mock_v = MagicMock() mock_v.get_access_token.return_value = ("token", "project") mock_http = MagicMock() - mock_http.get.side_effect = slow_http_get - async def run_scenario() -> float: - async def concurrent_sleep() -> float: - start = time.monotonic() - await asyncio.sleep(0.05) - return time.monotonic() - start + async def run_scenario() -> None: + loop = asyncio.get_running_loop() + http_get_blocked = asyncio.Event() + def blocking_http_get(*args: object, **kwargs: object) -> MagicMock: + loop.call_soon_threadsafe(http_get_blocked.set) + released_while_blocked.append(http_get_released_by_the_loop.wait(timeout=5)) + return http_get_response() + + async def release_http_get_from_the_loop() -> None: + await http_get_blocked.wait() + http_get_released_by_the_loop.set() + + mock_http.get.side_effect = blocking_http_get + releaser = asyncio.create_task(release_http_get_from_the_loop()) task = asyncio.create_task( gemini_transformation.async_transform_request_body( gemini_api_key=None, @@ -306,9 +316,8 @@ def test_async_transform_request_body_offloads_extensionless_gs_not_plain_text() vertex_auth_header=None, ) ) - elapsed = await concurrent_sleep() await task - return elapsed + await asyncio.wait_for(releaser, timeout=5) with ( patch.object(gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_v), @@ -322,11 +331,10 @@ def test_async_transform_request_body_offloads_extensionless_gs_not_plain_text() new=fake_check_and_create_cache, ), ): - sleep_elapsed = asyncio.run(run_scenario()) + asyncio.run(run_scenario()) - assert sleep_elapsed < 0.4, ( - f"Event loop blocked for {sleep_elapsed:.3f}s; " - "async_transform_request_body did not offload sync GCS metadata" + assert released_while_blocked and all(released_while_blocked), ( + "Event loop blocked; async_transform_request_body did not offload sync GCS metadata" ) async def fake_cache2(self, **kwargs): diff --git a/tests/unit/llms/vertex_ai/test_vertex_llm_base.py b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py index 514728137cf..b65e0e1d59c 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py @@ -1874,8 +1874,8 @@ class TestVertexBase: # Should return the current (still usable) token immediately assert token == "near-expiry-token" - # Let the background refresh task run - await asyncio.sleep(0.05) + background_refreshes = tuple(vertex_base._background_refresh_tasks.values()) + await asyncio.wait_for(asyncio.gather(*background_refreshes), timeout=5) assert mock_refresh.called, "Background refresh should have been triggered" @@ -1982,8 +1982,8 @@ class TestVertexBase: custom_llm_provider="vertex_ai", ) - # Allow the background task to complete. - await asyncio.sleep(0.1) + background_refreshes = tuple(vertex_base._background_refresh_tasks.values()) + await asyncio.wait_for(asyncio.gather(*background_refreshes), timeout=5) # After completion the entry should have been removed by the done-callback. assert len(vertex_base._background_refresh_tasks) == 0, ( @@ -2002,6 +2002,7 @@ class TestVertexBase: vertex_base = VertexBase() num_keys = 20 + background_refreshes: list[asyncio.Task] = [] for i in range(num_keys): mock_creds = MagicMock() @@ -2032,9 +2033,9 @@ class TestVertexBase: project_id=f"project-{i}", custom_llm_provider="vertex_ai", ) + background_refreshes.extend(vertex_base._background_refresh_tasks.values()) - # Let all background tasks finish. - await asyncio.sleep(0.1) + await asyncio.wait_for(asyncio.gather(*set(background_refreshes)), timeout=5) assert len(vertex_base._background_refresh_tasks) == 0, ( f"Expected 0 tasks after all refreshes completed, " diff --git a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index 34d9741bcc0..ce0e2f77e87 100644 --- a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -60,11 +60,14 @@ class _QueuePubSub: self.queue.put_nowait(message) self.subscribed_channels: list[str] = [] self.closed = False + self.drained = asyncio.Event() async def subscribe(self, *channels: str) -> None: self.subscribed_channels.extend(channels) async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> object | None: + if self.queue.empty(): + self.drained.set() try: return await asyncio.wait_for(self.queue.get(), timeout) except asyncio.TimeoutError: @@ -155,10 +158,7 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None: ) subscriber.start() try: - for _ in range(200): - if cache.in_memory_cache.get_cache("project_id:p-1") is None: - break - await asyncio.sleep(0.01) + assert await asyncio.wait_for(pubsub.drained.wait(), timeout=5) finally: await subscriber.stop() @@ -185,10 +185,7 @@ async def test_subscriber_deletes_key_object_partition_entry_on_message() -> Non ) subscriber.start() try: - for _ in range(200): - if cache.get_cache(hashed_token, model_type=UserAPIKeyAuth) is None: - break - await asyncio.sleep(0.01) + assert await asyncio.wait_for(pubsub.drained.wait(), timeout=5) finally: await subscriber.stop() @@ -216,10 +213,7 @@ async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message() ) subscriber.start() try: - for _ in range(200): - if spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None: - break - await asyncio.sleep(0.01) + assert await asyncio.wait_for(pubsub.drained.wait(), timeout=5) finally: await subscriber.stop() diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_custom_code_bounded_execution.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_custom_code_bounded_execution.py index dc3d3c8cf91..9114d811c57 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_custom_code_bounded_execution.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_custom_code_bounded_execution.py @@ -17,6 +17,12 @@ def _worker_threads() -> list[str]: return [t.name for t in threading.enumerate() if t.name.startswith("guardrail-code:")] +def _join_worker_threads() -> None: + for thread in threading.enumerate(): + if thread.name.startswith("guardrail-code:"): + thread.join(timeout=5) + + def _spin_forever() -> None: n = 0 while True: @@ -82,7 +88,7 @@ def test_call_with_timeout_interrupts_a_busy_loop_and_reclaims_the_thread(fn): assert exc_info.value.timeout == 0.2 assert time.monotonic() - started < 1.5 - time.sleep(0.2) + _join_worker_threads() assert _worker_threads() == [] @@ -105,7 +111,7 @@ async def test_call_off_loop_with_timeout_keeps_the_loop_running_and_stops_the_w ticker.cancel() assert ticks >= 5 - await asyncio.sleep(0.2) + await asyncio.to_thread(_join_worker_threads) assert _worker_threads() == [] @@ -142,22 +148,30 @@ async def test_await_with_timeout_abandons_a_coroutine_that_never_stops_swallowi @pytest.mark.asyncio async def test_call_off_loop_with_timeout_stops_the_worker_when_the_caller_is_cancelled(): - waiting = asyncio.create_task(call_off_loop_with_timeout(_spin_forever, 30.0, label="spin")) - await asyncio.sleep(0.1) + spinning = threading.Event() + + def signal_then_spin() -> None: + spinning.set() + _spin_forever() + + waiting = asyncio.create_task(call_off_loop_with_timeout(signal_then_spin, 30.0, label="spin")) + assert await asyncio.to_thread(spinning.wait, 5) waiting.cancel() with pytest.raises(asyncio.CancelledError): await waiting - await asyncio.sleep(0.3) + await asyncio.to_thread(_join_worker_threads) assert _worker_threads() == [] @pytest.mark.asyncio async def test_await_with_timeout_cancels_the_code_when_the_caller_is_cancelled(): + started = asyncio.Event() interrupted = asyncio.Event() async def sleep_until_cancelled() -> None: + started.set() try: await asyncio.sleep(30) except asyncio.CancelledError: @@ -165,10 +179,10 @@ async def test_await_with_timeout_cancels_the_code_when_the_caller_is_cancelled( raise waiting = asyncio.create_task(await_with_timeout(sleep_until_cancelled(), 30.0, label="sleep")) - await asyncio.sleep(0.1) + assert await asyncio.wait_for(started.wait(), timeout=5) waiting.cancel() with pytest.raises(asyncio.CancelledError): await waiting - await asyncio.wait_for(interrupted.wait(), timeout=1.0) + assert await asyncio.wait_for(interrupted.wait(), timeout=5) diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 920f9d90bb8..b981e8282bf 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -3422,26 +3422,27 @@ async def test_add_proxy_budget_to_db_only_creates_user_no_keys(monkeypatch: pyt litellm_proxy_budget_name = "litellm-proxy-budget" - # Mock generate_key_helper_fn to capture its call arguments - mock_generate_key_helper = AsyncMock( - return_value={ + generate_key_helper_called: Final = asyncio.Event() + + async def _record_generate_key_helper(*args: object, **kwargs: object) -> dict[str, object]: + generate_key_helper_called.set() + return { "user_id": litellm_proxy_budget_name, "max_budget": 100.0, "budget_duration": "30d", "spend": 0, "models": [], } - ) + + # Mock generate_key_helper_fn to capture its call arguments + mock_generate_key_helper = AsyncMock(side_effect=_record_generate_key_helper) # Patch generate_key_helper_fn in proxy_server where it's being called from with patch("litellm.proxy.proxy_server.generate_key_helper_fn", mock_generate_key_helper): # Call the function under test ProxyStartupEvent._add_proxy_budget_to_db() - # Allow async task to complete - import asyncio - - await asyncio.sleep(0.1) + assert await asyncio.wait_for(generate_key_helper_called.wait(), timeout=5) # Verify that generate_key_helper_fn was called mock_generate_key_helper.assert_called_once() @@ -6305,8 +6306,15 @@ async def test_tag_cache_update_called(): "spend": 10.0, } + set_cache_called: Final = asyncio.Event() + + async def _record_set_cache(*args: object, **kwargs: object) -> None: + set_cache_called.set() + with patch.object(cache, "async_batch_get_cache", new=AsyncMock(return_value=[mock_tag_obj])) as mock_get_cache: - with patch.object(cache, "async_set_cache_pipeline", new=AsyncMock()) as mock_set_cache: + with patch.object( + cache, "async_set_cache_pipeline", new=AsyncMock(side_effect=_record_set_cache) + ) as mock_set_cache: await litellm.proxy.proxy_server.update_cache( token=None, user_id=None, @@ -6317,7 +6325,7 @@ async def test_tag_cache_update_called(): tags=["test-tag"], ) - await asyncio.sleep(0.1) + assert await asyncio.wait_for(set_cache_called.wait(), timeout=5) mock_get_cache.assert_awaited_once_with(keys=["tag:test-tag"], parent_otel_span=None, throttle_redis=False) mock_set_cache.assert_awaited_once() @@ -6353,10 +6361,17 @@ async def test_tag_cache_update_multiple_tags(): async def mock_get_cache_side_effect(keys, **kwargs): return [{"tag:tag1": mock_tag1_obj, "tag:tag2": mock_tag2_obj}.get(key) for key in keys] + set_cache_called: Final = asyncio.Event() + + async def _record_set_cache(*args: object, **kwargs: object) -> None: + set_cache_called.set() + with patch.object( cache, "async_batch_get_cache", new=AsyncMock(side_effect=mock_get_cache_side_effect) ) as mock_get_cache: - with patch.object(cache, "async_set_cache_pipeline", new=AsyncMock()) as mock_set_cache: + with patch.object( + cache, "async_set_cache_pipeline", new=AsyncMock(side_effect=_record_set_cache) + ) as mock_set_cache: await litellm.proxy.proxy_server.update_cache( token=None, user_id=None, @@ -6367,7 +6382,7 @@ async def test_tag_cache_update_multiple_tags(): tags=["tag1", "tag2"], ) - await asyncio.sleep(0.1) + assert await asyncio.wait_for(set_cache_called.wait(), timeout=5) mock_get_cache.assert_awaited_once_with(keys=["tag:tag1", "tag:tag2"], parent_otel_span=None, throttle_redis=False) mock_set_cache.assert_awaited_once() @@ -6397,13 +6412,20 @@ async def test_update_cache_pipeline_honors_user_api_key_cache_ttl(): original_cache = litellm.proxy.proxy_server.user_api_key_cache cache = DualCache(default_in_memory_ttl=300) setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache) + set_cache_called: Final = asyncio.Event() + + async def _record_set_cache(*args: object, **kwargs: object) -> None: + set_cache_called.set() + try: with patch.object( cache, "async_batch_get_cache", new=AsyncMock(return_value=[{"tag_name": "active-tag", "spend": 1.0}]), ): - with patch.object(cache, "async_set_cache_pipeline", new=AsyncMock()) as mock_set_cache: + with patch.object( + cache, "async_set_cache_pipeline", new=AsyncMock(side_effect=_record_set_cache) + ) as mock_set_cache: await litellm.proxy.proxy_server.update_cache( token=None, user_id=None, @@ -6414,7 +6436,7 @@ async def test_update_cache_pipeline_honors_user_api_key_cache_ttl(): tags=["active-tag"], ) - await asyncio.sleep(0.1) + assert await asyncio.wait_for(set_cache_called.wait(), timeout=5) mock_set_cache.assert_awaited_once() assert mock_set_cache.call_args.kwargs["ttl"] == 300 diff --git a/tests/unit/responses/test_responses_streaming_iterator.py b/tests/unit/responses/test_responses_streaming_iterator.py index 9ba7dfcaa80..0419fce8ec2 100644 --- a/tests/unit/responses/test_responses_streaming_iterator.py +++ b/tests/unit/responses/test_responses_streaming_iterator.py @@ -119,6 +119,11 @@ def _make_iterator(logging_obj: LitellmLogging) -> ResponsesAPIStreamingIterator return iterator +async def _await_tasks_spawned_since(existing: set["asyncio.Task[object]"]) -> None: + spawned = asyncio.all_tasks() - existing + await asyncio.wait_for(asyncio.gather(*spawned, return_exceptions=True), timeout=5) + + @pytest.mark.asyncio async def test_custom_logger_only_never_submits_sync_success_handler(recording_executor): recorder = RecordingCustomLogger() @@ -128,8 +133,9 @@ async def test_custom_logger_only_never_submits_sync_success_handler(recording_e logging_obj = _make_logging_obj() iterator = _make_iterator(logging_obj) + existing_tasks = asyncio.all_tasks() iterator._log_completed_response(is_async=True) - await asyncio.sleep(0.6) + await _await_tasks_spawned_since(existing_tasks) assert recorder.async_hook_started is not None assert recording_executor.submit_times_for(logging_obj) == [] @@ -149,8 +155,9 @@ async def test_sync_callbacks_run_only_after_async_handler_completes(recording_e logging_obj = _make_logging_obj() iterator = _make_iterator(logging_obj) + existing_tasks = asyncio.all_tasks() iterator._log_completed_response(is_async=True) - await asyncio.sleep(0.8) + await _await_tasks_spawned_since(existing_tasks) assert recorder.async_hook_finished is not None submit_times = recording_executor.submit_times_for(logging_obj) diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index 4d44b8602df..9e3b4e86f6f 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -5,6 +5,7 @@ completion_start_time = end_time.""" import asyncio, importlib import json +import threading from collections.abc import Callable from datetime import datetime from typing import Final, Optional @@ -1350,6 +1351,8 @@ class _FakeLoggingObj: self.start_time = datetime.now() self.completion_start_time = None self.model_call_details = {"litellm_params": {}} + self.success_logged = threading.Event() + self.failure_logged = threading.Event() # Signature alignment with Logging handlers async def dispatch_success_handlers(self, *args, **kwargs): @@ -1360,6 +1363,7 @@ class _FakeLoggingObj: def success_handler(self, *args, **kwargs): self.success_calls += 1 self.last_success_kwargs = kwargs + self.success_logged.set() async def async_success_handler(self, *args, **kwargs): self.async_success_calls += 1 @@ -1367,6 +1371,7 @@ class _FakeLoggingObj: def failure_handler(self, *args, **kwargs): self.failure_calls += 1 + self.failure_logged.set() async def async_failure_handler(self, *args, **kwargs): self.async_failure_calls += 1 @@ -1375,6 +1380,10 @@ class _FakeLoggingObj: self.completion_start_time = completion_start_time self.model_call_details["completion_start_time"] = completion_start_time +async def _await_tasks_spawned_since(existing: set["asyncio.Task[object]"]) -> None: + spawned: Final = asyncio.all_tasks() - existing + await asyncio.wait_for(asyncio.gather(*spawned, return_exceptions=True), timeout=5) + def _make_completed_response(response_id: str = "resp_test") -> ResponseCompletedEvent: return ResponseCompletedEvent( type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, @@ -1525,8 +1534,9 @@ async def test_responses_streaming_triggers_hooks(monkeypatch): type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=SimpleNamespace() ) + existing_tasks = asyncio.all_tasks() iterator._handle_logging_completed_response() - await asyncio.sleep(0.2) # allow async tasks to run + await _await_tasks_spawned_since(existing_tasks) assert logging_obj.success_calls == 1 assert logging_obj.async_success_calls == 1 @@ -1601,8 +1611,7 @@ async def test_responses_streaming_failure_triggers_failure_handlers(): with pytest.raises(ValueError, match="boom"): iterator._process_chunk('{"delta": "chunk"}') - # allow failure callbacks to run - await asyncio.sleep(0.2) + assert await asyncio.to_thread(logging_obj.failure_logged.wait, 5) assert logging_obj.failure_calls >= 1 assert logging_obj.async_failure_calls >= 1 @@ -1769,7 +1778,7 @@ async def test_handle_logging_failed_response_uses_response_error_message(): ) iterator._handle_logging_failed_response() - await asyncio.sleep(0.2) + assert await asyncio.to_thread(logging_obj.failure_logged.wait, 5) assert logging_obj.failure_calls == 1 assert logging_obj.async_failure_calls == 1 @@ -1891,8 +1900,9 @@ async def test_responses_streaming_completed_event_persists_async_cache(): ) iterator.completed_response = _make_completed_response() + existing_tasks = asyncio.all_tasks() iterator._handle_logging_completed_response() - await asyncio.sleep(0.2) + await _await_tasks_spawned_since(existing_tasks) litellm.cache.async_add_cache.assert_called_once() assert litellm.cache.async_add_cache.call_args.kwargs["stream"] is True @@ -1981,7 +1991,7 @@ def test_log_completed_response_sync_direct_path(monkeypatch): iterator.completed_response = _make_completed_response("resp_log_sync") iterator._log_completed_response(is_async=False) - asyncio.run(asyncio.sleep(0.2)) + assert logging_obj.success_logged.wait(timeout=5) assert logging_obj.success_calls == 1 assert logging_obj.async_success_calls == 1 @@ -2012,7 +2022,7 @@ def test_log_completed_response_falls_back_when_model_validate_fails(monkeypatch monkeypatch.setattr(iterator, "_run_post_success_hooks", MagicMock()) iterator._log_completed_response(is_async=False) - asyncio.run(asyncio.sleep(0.2)) + assert logging_obj.success_logged.wait(timeout=5) assert logging_obj.success_calls == 1 assert logging_obj.async_success_calls == 1 @@ -2178,8 +2188,9 @@ async def test_mock_responses_streaming_iterator_async_iteration_logs_completion call_type=CallTypes.responses.value, ) + existing_tasks = asyncio.all_tasks() streamed_events = [event async for event in iterator] - await asyncio.sleep(0.2) + await _await_tasks_spawned_since(existing_tasks) assert streamed_events[0].type == ResponsesAPIStreamEvents.RESPONSE_CREATED assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED @@ -2224,7 +2235,7 @@ def test_mock_responses_streaming_iterator_sync_iteration_logs_completion(monkey ) streamed_events = list(iterator) - asyncio.run(asyncio.sleep(0.2)) + assert logging_obj.success_logged.wait(timeout=5) assert streamed_events[0].type == ResponsesAPIStreamEvents.RESPONSE_CREATED assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED @@ -2277,8 +2288,9 @@ async def test_cached_responses_stream_async_hit_triggers_success_callbacks( call_type=CallTypes.aresponses.value, ) + existing_tasks = asyncio.all_tasks() streamed_events = [event async for event in iterator] - await asyncio.sleep(0.2) + await _await_tasks_spawned_since(existing_tasks) assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED assert logging_obj.success_calls == 1 @@ -2333,7 +2345,7 @@ def test_cached_responses_stream_sync_hit_triggers_success_callbacks(monkeypatch ) streamed_events = list(iterator) - asyncio.run(asyncio.sleep(0.2)) + assert logging_obj.success_logged.wait(timeout=5) assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED assert logging_obj.success_calls == 1 diff --git a/tests/unit/test_litellm_overhead.py b/tests/unit/test_litellm_overhead.py index a944a6da487..c2241b55adb 100644 --- a/tests/unit/test_litellm_overhead.py +++ b/tests/unit/test_litellm_overhead.py @@ -145,6 +145,7 @@ async def test_litellm_overhead_stream(monkeypatch): @pytest.mark.asyncio async def test_litellm_overhead_cache_hit(monkeypatch): from litellm.caching.caching import Cache + from litellm.caching.caching_handler import _PENDING_CACHE_WRITES calls = _mock_openai_completion_transport(monkeypatch, response_id="chatcmpl-cache") litellm.cache = Cache() @@ -157,7 +158,7 @@ async def test_litellm_overhead_cache_hit(monkeypatch): messages=messages, caching=True, ) - await asyncio.sleep(0.5) + await asyncio.wait_for(asyncio.gather(*_PENDING_CACHE_WRITES), timeout=5) response2 = await litellm.acompletion( model="gpt-4o", api_key="test-key", diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 6ad4f75f97f..124369a9528 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -6264,25 +6264,30 @@ async def test_wrapper_async_preserves_original_exception_when_hook_await_is_can the caller - not get replaced by CancelledError/TimeoutError from the hook's own await getting cancelled.""" + hook_entered: Final = asyncio.Event() + hook_released: Final = asyncio.Event() + class SlowLogger(CustomLogger): async def async_post_call_failure_deployment_hook( self, request_data, exception, call_type, fallback_depth=None ): - await asyncio.sleep(5) + hook_entered.set() + await hook_released.wait() monkeypatch.setattr(litellm, "callbacks", [SlowLogger()]) - with pytest.raises(litellm.AuthenticationError): - await asyncio.wait_for( - litellm.acompletion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "hi"}], - mock_response=litellm.AuthenticationError( - message="bad key", llm_provider="openai", model="gpt-4o-mini" - ), - ), - timeout=0.2, + call: Final = asyncio.create_task( + litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + mock_response=litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini"), ) + ) + assert await asyncio.wait_for(hook_entered.wait(), timeout=5) + call.cancel() + + with pytest.raises(litellm.AuthenticationError): + await call @pytest.mark.asyncio From 829bd6a8cc7ccaab24fe7fca1c7fc2fc7b9a38b3 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 13 Oct 2026 12:52:01 -0700 Subject: [PATCH 007/325] fix(model-prices): add TokenHub rows, fix TokenHub DeepSeek-V4-Pro price, Bedrock Claude 3.5 Sonnet and Fireworks context windows, Together dates (#46206) * fix(model-prices): add TokenHub DeepSeek, GLM and Kimi rows, correct TokenHub DeepSeek-V4-Pro price, update Together removal dates Co-authored-by: halewang Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(model-prices): correct bare Bedrock Claude 3.5 Sonnet context window Align bare Sonnet entries with their regional 200K context metadata. Co-authored-by: OrangerXu Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(model-prices): correct Fireworks serverless context windows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: halewang Co-authored-by: OrangerXu --- ...odel_prices_and_context_window_backup.json | 392 +++++++++++++++--- model_prices_and_context_window.json | 392 +++++++++++++++--- tests/unit/test_model_prices_schema.py | 147 +++++++ 3 files changed, 837 insertions(+), 94 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8977c5dd88e..0b96fd3e2b0 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -901,7 +901,7 @@ "anthropic.claude-3-5-sonnet-20240620-v1:0": { "input_cost_per_token": 3e-06, "litellm_provider": "bedrock", - "max_input_tokens": 1000000, + "max_input_tokens": 200000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", @@ -911,10 +911,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "input_cost_per_token_above_200k_tokens": 6e-06, - "output_cost_per_token_above_200k_tokens": 3e-05, - "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, - "cache_read_input_token_cost_above_200k_tokens": 6e-07, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07 }, @@ -923,7 +919,7 @@ "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "litellm_provider": "bedrock", - "max_input_tokens": 1000000, + "max_input_tokens": 200000, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", @@ -935,11 +931,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "input_cost_per_token_above_200k_tokens": 6e-06, - "output_cost_per_token_above_200k_tokens": 3e-05, - "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, - "cache_read_input_token_cost_above_200k_tokens": 6e-07 + "supports_vision": true }, "anthropic.claude-3-7-sonnet-20240620-v1:0": { "cache_creation_input_token_cost": 4.5e-06, @@ -26245,7 +26237,7 @@ "fireworks_ai/accounts/fireworks/models/deepseek-r1": { "input_cost_per_token": 3e-06, "litellm_provider": "fireworks_ai", - "max_input_tokens": 128000, + "max_input_tokens": 163840, "max_output_tokens": 20480, "max_tokens": 20480, "mode": "chat", @@ -26257,7 +26249,7 @@ "fireworks_ai/accounts/fireworks/models/deepseek-r1-0528": { "input_cost_per_token": 3e-06, "litellm_provider": "fireworks_ai", - "max_input_tokens": 160000, + "max_input_tokens": 163840, "max_output_tokens": 160000, "max_tokens": 160000, "mode": "chat", @@ -26269,7 +26261,7 @@ "fireworks_ai/accounts/fireworks/models/deepseek-r1-basic": { "input_cost_per_token": 5.5e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 128000, + "max_input_tokens": 163840, "max_output_tokens": 20480, "max_tokens": 20480, "mode": "chat", @@ -26281,7 +26273,7 @@ "fireworks_ai/accounts/fireworks/models/deepseek-v3": { "input_cost_per_token": 9e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 128000, + "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", @@ -26414,7 +26406,7 @@ "fireworks_ai/accounts/fireworks/models/glm-4p5": { "input_cost_per_token": 5.5e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 128000, + "max_input_tokens": 131072, "max_output_tokens": 96000, "max_tokens": 96000, "mode": "chat", @@ -26428,7 +26420,7 @@ "fireworks_ai/accounts/fireworks/models/glm-4p5-air": { "input_cost_per_token": 2.2e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 128000, + "max_input_tokens": 131072, "max_output_tokens": 96000, "max_tokens": 96000, "mode": "chat", @@ -26443,7 +26435,7 @@ "input_cost_per_token": 5.5e-07, "output_cost_per_token": 2.19e-06, "litellm_provider": "fireworks_ai", - "max_input_tokens": 202800, + "max_input_tokens": 202752, "max_output_tokens": 202800, "max_tokens": 202800, "mode": "chat", @@ -26457,7 +26449,7 @@ "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 6e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 202800, + "max_input_tokens": 202752, "max_output_tokens": 202800, "max_tokens": 202800, "mode": "chat", @@ -26472,7 +26464,7 @@ "cache_read_input_token_cost": 2.6e-07, "input_cost_per_token": 1.4e-06, "litellm_provider": "fireworks_ai", - "max_input_tokens": 202800, + "max_input_tokens": 202752, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", @@ -26636,7 +26628,7 @@ "fireworks_ai/accounts/fireworks/models/llama-v3p1-405b-instruct": { "input_cost_per_token": 3e-06, "litellm_provider": "fireworks_ai", - "max_input_tokens": 128000, + "max_input_tokens": 131072, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", @@ -26649,7 +26641,7 @@ "fireworks_ai/accounts/fireworks/models/llama-v3p1-8b-instruct": { "input_cost_per_token": 1e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 16384, + "max_input_tokens": 131072, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", @@ -26662,7 +26654,7 @@ "fireworks_ai/accounts/fireworks/models/llama-v3p2-11b-vision-instruct": { "input_cost_per_token": 2e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 16384, + "max_input_tokens": 131072, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", @@ -26676,7 +26668,7 @@ "fireworks_ai/accounts/fireworks/models/llama-v3p2-1b-instruct": { "input_cost_per_token": 1e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 16384, + "max_input_tokens": 131072, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", @@ -26689,7 +26681,7 @@ "fireworks_ai/accounts/fireworks/models/llama-v3p2-3b-instruct": { "input_cost_per_token": 1e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 16384, + "max_input_tokens": 131072, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", @@ -26702,7 +26694,7 @@ "fireworks_ai/accounts/fireworks/models/llama-v3p2-90b-vision-instruct": { "input_cost_per_token": 9e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 16384, + "max_input_tokens": 131072, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", @@ -26740,7 +26732,7 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_token": 3e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 204800, + "max_input_tokens": 196608, "max_output_tokens": 204800, "max_tokens": 204800, "mode": "chat", @@ -26841,7 +26833,7 @@ "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 6e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 202800, + "max_input_tokens": 202752, "max_output_tokens": 202800, "max_tokens": 202800, "mode": "chat", @@ -26856,7 +26848,7 @@ "cache_read_input_token_cost": 2.6e-07, "input_cost_per_token": 1.4e-06, "litellm_provider": "fireworks_ai", - "max_input_tokens": 202800, + "max_input_tokens": 202752, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", @@ -27031,7 +27023,7 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_token": 3e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 204800, + "max_input_tokens": 196608, "max_output_tokens": 204800, "max_tokens": 204800, "mode": "chat", @@ -45602,6 +45594,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo": { + "deprecation_date": "2026-10-22", "input_cost_per_token": 1.04e-06, "litellm_provider": "together_ai", "max_input_tokens": 131072, @@ -45758,6 +45751,7 @@ }, "together_ai/deepseek-ai/DeepSeek-V4-Flash-0731": { "cache_read_input_token_cost": 3e-08, + "deprecation_date": "2026-10-22", "input_cost_per_token": 1.4e-07, "litellm_provider": "together_ai", "max_input_tokens": 1048576, @@ -52638,7 +52632,7 @@ }, "fireworks_ai/accounts/fireworks/models/ernie-4p5-21b-a3b-pt": { "max_tokens": 4096, - "max_input_tokens": 4096, + "max_input_tokens": 131072, "max_output_tokens": 4096, "input_cost_per_token": 1e-07, "output_cost_per_token": 1e-07, @@ -52647,7 +52641,7 @@ }, "fireworks_ai/accounts/fireworks/models/ernie-4p5-300b-a47b-pt": { "max_tokens": 4096, - "max_input_tokens": 4096, + "max_input_tokens": 131072, "max_output_tokens": 4096, "input_cost_per_token": 1e-07, "output_cost_per_token": 1e-07, @@ -53107,7 +53101,7 @@ }, "fireworks_ai/accounts/fireworks/models/minimax-m2": { "max_tokens": 4096, - "max_input_tokens": 4096, + "max_input_tokens": 196608, "max_output_tokens": 4096, "input_cost_per_token": 3e-07, "output_cost_per_token": 1.2e-06, @@ -53278,7 +53272,7 @@ }, "fireworks_ai/accounts/fireworks/models/nemotron-nano-v2-12b-vl": { "max_tokens": 4096, - "max_input_tokens": 4096, + "max_input_tokens": 131072, "max_output_tokens": 4096, "input_cost_per_token": 1e-07, "output_cost_per_token": 1e-07, @@ -53341,7 +53335,7 @@ }, "fireworks_ai/accounts/fireworks/models/nvidia-nemotron-nano-12b-v2": { "max_tokens": 131072, - "max_input_tokens": 131072, + "max_input_tokens": 128000, "max_output_tokens": 131072, "input_cost_per_token": 2e-07, "output_cost_per_token": 2e-07, @@ -53350,7 +53344,7 @@ }, "fireworks_ai/accounts/fireworks/models/nvidia-nemotron-nano-9b-v2": { "max_tokens": 131072, - "max_input_tokens": 131072, + "max_input_tokens": 128000, "max_output_tokens": 131072, "input_cost_per_token": 2e-07, "output_cost_per_token": 2e-07, @@ -53946,7 +53940,7 @@ }, "fireworks_ai/accounts/fireworks/models/qwen3-coder-480b-instruct-bf16": { "max_tokens": 4096, - "max_input_tokens": 4096, + "max_input_tokens": 262144, "max_output_tokens": 4096, "input_cost_per_token": 9e-07, "output_cost_per_token": 9e-07, @@ -54064,7 +54058,7 @@ }, "fireworks_ai/accounts/fireworks/models/qwen3-vl-8b-instruct": { "max_tokens": 4096, - "max_input_tokens": 4096, + "max_input_tokens": 262144, "max_output_tokens": 4096, "input_cost_per_token": 2e-07, "output_cost_per_token": 2e-07, @@ -59480,16 +59474,16 @@ }, "tencent/deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 3.625e-09, - "input_cost_per_token": 4.35e-07, - "input_cost_per_token_cache_hit": 3.625e-09, + "cache_read_input_token_cost": 1.45e-07, + "input_cost_per_token": 1.74e-06, + "input_cost_per_token_cache_hit": 1.45e-07, "litellm_provider": "tencent", "max_input_tokens": 1000000, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 8.7e-07, - "source": "https://www.tencentcloud.com/products/tokenhub", + "output_cost_per_token": 3.48e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", "supported_endpoints": [ "/v1/chat/completions" ], @@ -59515,7 +59509,7 @@ "max_tokens": 384000, "mode": "chat", "output_cost_per_token": 2.8e-07, - "source": "https://www.tencentcloud.com/products/tokenhub", + "source": "https://www.tencentcloud.com/document/product/1300/78937", "supported_endpoints": [ "/v1/chat/completions" ], @@ -59530,6 +59524,310 @@ "supports_reasoning": true, "supports_vision": false }, + "tencent/deepseek-v4-pro-0813": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 4.4e-08, + "input_cost_per_token": 1.32e-06, + "input_cost_per_token_cache_hit": 4.4e-08, + "litellm_provider": "tencent", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 2.2e-08, + "input_cost_per_token": 6.6e-07, + "output_cost_per_token": 1.98e-06, + "weekday_timezone": "Asia/Shanghai", + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, + "output_cost_per_token": 3.96e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_vision": false + }, + "tencent/deepseek-v4.1-flash": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, + "litellm_provider": "tencent", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "weekday_timezone": "Asia/Shanghai", + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, + "output_cost_per_token": 1.2e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_vision": true + }, + "tencent/glm-5.3-flashx": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 3.7e-07, + "input_cost_per_token_cache_hit": 7.5e-08, + "litellm_provider": "tencent", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.25e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, + "tencent/glm-5.3-flash": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "input_cost_per_token_cache_hit": 3e-08, + "litellm_provider": "tencent", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, + "tencent/glm-5.3": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "input_cost_per_token_cache_hit": 2.6e-07, + "litellm_provider": "tencent", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "tencent/glm-5.2": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "input_cost_per_token_cache_hit": 2.6e-07, + "litellm_provider": "tencent", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "tencent/kimi-k3": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_cache_hit": 3e-07, + "litellm_provider": "tencent", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "tencent/kimi-k2.8-preview": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_cache_hit": 2.5e-07, + "litellm_provider": "tencent", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "tencent/kimi-k2.7-code-highspeed": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 3.8e-07, + "input_cost_per_token": 1.9e-06, + "input_cost_per_token_cache_hit": 3.8e-07, + "litellm_provider": "tencent", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 8e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "tencent/kimi-k2.7-code": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 9.5e-07, + "input_cost_per_token_cache_hit": 1.9e-07, + "litellm_provider": "tencent", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, "tencent/minimax-m3": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 6e-08, @@ -59539,7 +59837,7 @@ "max_input_tokens": 1000000, "mode": "chat", "output_cost_per_token": 1.2e-06, - "source": "https://www.tencentcloud.com/products/tokenhub", + "source": "https://www.tencentcloud.com/document/product/1300/78937", "supported_endpoints": [ "/v1/chat/completions" ], @@ -66399,7 +66697,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/google/gemma-4-31B-it": { - "deprecation_date": "2026-09-15", + "deprecation_date": "2026-09-14", "input_cost_per_token": 3.9e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, @@ -66542,7 +66840,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/openai/gpt-oss-20b": { - "deprecation_date": "2026-09-15", + "deprecation_date": "2026-09-14", "input_cost_per_token": 5e-08, "litellm_provider": "together_ai", "max_input_tokens": 131072, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8977c5dd88e..0b96fd3e2b0 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -901,7 +901,7 @@ "anthropic.claude-3-5-sonnet-20240620-v1:0": { "input_cost_per_token": 3e-06, "litellm_provider": "bedrock", - "max_input_tokens": 1000000, + "max_input_tokens": 200000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", @@ -911,10 +911,6 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "input_cost_per_token_above_200k_tokens": 6e-06, - "output_cost_per_token_above_200k_tokens": 3e-05, - "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, - "cache_read_input_token_cost_above_200k_tokens": 6e-07, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07 }, @@ -923,7 +919,7 @@ "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "litellm_provider": "bedrock", - "max_input_tokens": 1000000, + "max_input_tokens": 200000, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", @@ -935,11 +931,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "input_cost_per_token_above_200k_tokens": 6e-06, - "output_cost_per_token_above_200k_tokens": 3e-05, - "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, - "cache_read_input_token_cost_above_200k_tokens": 6e-07 + "supports_vision": true }, "anthropic.claude-3-7-sonnet-20240620-v1:0": { "cache_creation_input_token_cost": 4.5e-06, @@ -26245,7 +26237,7 @@ "fireworks_ai/accounts/fireworks/models/deepseek-r1": { "input_cost_per_token": 3e-06, "litellm_provider": "fireworks_ai", - "max_input_tokens": 128000, + "max_input_tokens": 163840, "max_output_tokens": 20480, "max_tokens": 20480, "mode": "chat", @@ -26257,7 +26249,7 @@ "fireworks_ai/accounts/fireworks/models/deepseek-r1-0528": { "input_cost_per_token": 3e-06, "litellm_provider": "fireworks_ai", - "max_input_tokens": 160000, + "max_input_tokens": 163840, "max_output_tokens": 160000, "max_tokens": 160000, "mode": "chat", @@ -26269,7 +26261,7 @@ "fireworks_ai/accounts/fireworks/models/deepseek-r1-basic": { "input_cost_per_token": 5.5e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 128000, + "max_input_tokens": 163840, "max_output_tokens": 20480, "max_tokens": 20480, "mode": "chat", @@ -26281,7 +26273,7 @@ "fireworks_ai/accounts/fireworks/models/deepseek-v3": { "input_cost_per_token": 9e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 128000, + "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", @@ -26414,7 +26406,7 @@ "fireworks_ai/accounts/fireworks/models/glm-4p5": { "input_cost_per_token": 5.5e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 128000, + "max_input_tokens": 131072, "max_output_tokens": 96000, "max_tokens": 96000, "mode": "chat", @@ -26428,7 +26420,7 @@ "fireworks_ai/accounts/fireworks/models/glm-4p5-air": { "input_cost_per_token": 2.2e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 128000, + "max_input_tokens": 131072, "max_output_tokens": 96000, "max_tokens": 96000, "mode": "chat", @@ -26443,7 +26435,7 @@ "input_cost_per_token": 5.5e-07, "output_cost_per_token": 2.19e-06, "litellm_provider": "fireworks_ai", - "max_input_tokens": 202800, + "max_input_tokens": 202752, "max_output_tokens": 202800, "max_tokens": 202800, "mode": "chat", @@ -26457,7 +26449,7 @@ "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 6e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 202800, + "max_input_tokens": 202752, "max_output_tokens": 202800, "max_tokens": 202800, "mode": "chat", @@ -26472,7 +26464,7 @@ "cache_read_input_token_cost": 2.6e-07, "input_cost_per_token": 1.4e-06, "litellm_provider": "fireworks_ai", - "max_input_tokens": 202800, + "max_input_tokens": 202752, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", @@ -26636,7 +26628,7 @@ "fireworks_ai/accounts/fireworks/models/llama-v3p1-405b-instruct": { "input_cost_per_token": 3e-06, "litellm_provider": "fireworks_ai", - "max_input_tokens": 128000, + "max_input_tokens": 131072, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", @@ -26649,7 +26641,7 @@ "fireworks_ai/accounts/fireworks/models/llama-v3p1-8b-instruct": { "input_cost_per_token": 1e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 16384, + "max_input_tokens": 131072, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", @@ -26662,7 +26654,7 @@ "fireworks_ai/accounts/fireworks/models/llama-v3p2-11b-vision-instruct": { "input_cost_per_token": 2e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 16384, + "max_input_tokens": 131072, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", @@ -26676,7 +26668,7 @@ "fireworks_ai/accounts/fireworks/models/llama-v3p2-1b-instruct": { "input_cost_per_token": 1e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 16384, + "max_input_tokens": 131072, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", @@ -26689,7 +26681,7 @@ "fireworks_ai/accounts/fireworks/models/llama-v3p2-3b-instruct": { "input_cost_per_token": 1e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 16384, + "max_input_tokens": 131072, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", @@ -26702,7 +26694,7 @@ "fireworks_ai/accounts/fireworks/models/llama-v3p2-90b-vision-instruct": { "input_cost_per_token": 9e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 16384, + "max_input_tokens": 131072, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", @@ -26740,7 +26732,7 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_token": 3e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 204800, + "max_input_tokens": 196608, "max_output_tokens": 204800, "max_tokens": 204800, "mode": "chat", @@ -26841,7 +26833,7 @@ "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 6e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 202800, + "max_input_tokens": 202752, "max_output_tokens": 202800, "max_tokens": 202800, "mode": "chat", @@ -26856,7 +26848,7 @@ "cache_read_input_token_cost": 2.6e-07, "input_cost_per_token": 1.4e-06, "litellm_provider": "fireworks_ai", - "max_input_tokens": 202800, + "max_input_tokens": 202752, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", @@ -27031,7 +27023,7 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_token": 3e-07, "litellm_provider": "fireworks_ai", - "max_input_tokens": 204800, + "max_input_tokens": 196608, "max_output_tokens": 204800, "max_tokens": 204800, "mode": "chat", @@ -45602,6 +45594,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo": { + "deprecation_date": "2026-10-22", "input_cost_per_token": 1.04e-06, "litellm_provider": "together_ai", "max_input_tokens": 131072, @@ -45758,6 +45751,7 @@ }, "together_ai/deepseek-ai/DeepSeek-V4-Flash-0731": { "cache_read_input_token_cost": 3e-08, + "deprecation_date": "2026-10-22", "input_cost_per_token": 1.4e-07, "litellm_provider": "together_ai", "max_input_tokens": 1048576, @@ -52638,7 +52632,7 @@ }, "fireworks_ai/accounts/fireworks/models/ernie-4p5-21b-a3b-pt": { "max_tokens": 4096, - "max_input_tokens": 4096, + "max_input_tokens": 131072, "max_output_tokens": 4096, "input_cost_per_token": 1e-07, "output_cost_per_token": 1e-07, @@ -52647,7 +52641,7 @@ }, "fireworks_ai/accounts/fireworks/models/ernie-4p5-300b-a47b-pt": { "max_tokens": 4096, - "max_input_tokens": 4096, + "max_input_tokens": 131072, "max_output_tokens": 4096, "input_cost_per_token": 1e-07, "output_cost_per_token": 1e-07, @@ -53107,7 +53101,7 @@ }, "fireworks_ai/accounts/fireworks/models/minimax-m2": { "max_tokens": 4096, - "max_input_tokens": 4096, + "max_input_tokens": 196608, "max_output_tokens": 4096, "input_cost_per_token": 3e-07, "output_cost_per_token": 1.2e-06, @@ -53278,7 +53272,7 @@ }, "fireworks_ai/accounts/fireworks/models/nemotron-nano-v2-12b-vl": { "max_tokens": 4096, - "max_input_tokens": 4096, + "max_input_tokens": 131072, "max_output_tokens": 4096, "input_cost_per_token": 1e-07, "output_cost_per_token": 1e-07, @@ -53341,7 +53335,7 @@ }, "fireworks_ai/accounts/fireworks/models/nvidia-nemotron-nano-12b-v2": { "max_tokens": 131072, - "max_input_tokens": 131072, + "max_input_tokens": 128000, "max_output_tokens": 131072, "input_cost_per_token": 2e-07, "output_cost_per_token": 2e-07, @@ -53350,7 +53344,7 @@ }, "fireworks_ai/accounts/fireworks/models/nvidia-nemotron-nano-9b-v2": { "max_tokens": 131072, - "max_input_tokens": 131072, + "max_input_tokens": 128000, "max_output_tokens": 131072, "input_cost_per_token": 2e-07, "output_cost_per_token": 2e-07, @@ -53946,7 +53940,7 @@ }, "fireworks_ai/accounts/fireworks/models/qwen3-coder-480b-instruct-bf16": { "max_tokens": 4096, - "max_input_tokens": 4096, + "max_input_tokens": 262144, "max_output_tokens": 4096, "input_cost_per_token": 9e-07, "output_cost_per_token": 9e-07, @@ -54064,7 +54058,7 @@ }, "fireworks_ai/accounts/fireworks/models/qwen3-vl-8b-instruct": { "max_tokens": 4096, - "max_input_tokens": 4096, + "max_input_tokens": 262144, "max_output_tokens": 4096, "input_cost_per_token": 2e-07, "output_cost_per_token": 2e-07, @@ -59480,16 +59474,16 @@ }, "tencent/deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 3.625e-09, - "input_cost_per_token": 4.35e-07, - "input_cost_per_token_cache_hit": 3.625e-09, + "cache_read_input_token_cost": 1.45e-07, + "input_cost_per_token": 1.74e-06, + "input_cost_per_token_cache_hit": 1.45e-07, "litellm_provider": "tencent", "max_input_tokens": 1000000, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 8.7e-07, - "source": "https://www.tencentcloud.com/products/tokenhub", + "output_cost_per_token": 3.48e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", "supported_endpoints": [ "/v1/chat/completions" ], @@ -59515,7 +59509,7 @@ "max_tokens": 384000, "mode": "chat", "output_cost_per_token": 2.8e-07, - "source": "https://www.tencentcloud.com/products/tokenhub", + "source": "https://www.tencentcloud.com/document/product/1300/78937", "supported_endpoints": [ "/v1/chat/completions" ], @@ -59530,6 +59524,310 @@ "supports_reasoning": true, "supports_vision": false }, + "tencent/deepseek-v4-pro-0813": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 4.4e-08, + "input_cost_per_token": 1.32e-06, + "input_cost_per_token_cache_hit": 4.4e-08, + "litellm_provider": "tencent", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 2.2e-08, + "input_cost_per_token": 6.6e-07, + "output_cost_per_token": 1.98e-06, + "weekday_timezone": "Asia/Shanghai", + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, + "output_cost_per_token": 3.96e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_vision": false + }, + "tencent/deepseek-v4.1-flash": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "input_cost_per_token_cache_hit": 6e-09, + "litellm_provider": "tencent", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "weekday_timezone": "Asia/Shanghai", + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, + "output_cost_per_token": 1.2e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_vision": true + }, + "tencent/glm-5.3-flashx": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 3.7e-07, + "input_cost_per_token_cache_hit": 7.5e-08, + "litellm_provider": "tencent", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.25e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, + "tencent/glm-5.3-flash": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "input_cost_per_token_cache_hit": 3e-08, + "litellm_provider": "tencent", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, + "tencent/glm-5.3": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "input_cost_per_token_cache_hit": 2.6e-07, + "litellm_provider": "tencent", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "tencent/glm-5.2": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "input_cost_per_token_cache_hit": 2.6e-07, + "litellm_provider": "tencent", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "tencent/kimi-k3": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_cache_hit": 3e-07, + "litellm_provider": "tencent", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "tencent/kimi-k2.8-preview": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_cache_hit": 2.5e-07, + "litellm_provider": "tencent", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "tencent/kimi-k2.7-code-highspeed": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 3.8e-07, + "input_cost_per_token": 1.9e-06, + "input_cost_per_token_cache_hit": 3.8e-07, + "litellm_provider": "tencent", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 8e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "tencent/kimi-k2.7-code": { + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 9.5e-07, + "input_cost_per_token_cache_hit": 1.9e-07, + "litellm_provider": "tencent", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, "tencent/minimax-m3": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 6e-08, @@ -59539,7 +59837,7 @@ "max_input_tokens": 1000000, "mode": "chat", "output_cost_per_token": 1.2e-06, - "source": "https://www.tencentcloud.com/products/tokenhub", + "source": "https://www.tencentcloud.com/document/product/1300/78937", "supported_endpoints": [ "/v1/chat/completions" ], @@ -66399,7 +66697,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/google/gemma-4-31B-it": { - "deprecation_date": "2026-09-15", + "deprecation_date": "2026-09-14", "input_cost_per_token": 3.9e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, @@ -66542,7 +66840,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/openai/gpt-oss-20b": { - "deprecation_date": "2026-09-15", + "deprecation_date": "2026-09-14", "input_cost_per_token": 5e-08, "litellm_provider": "together_ai", "max_input_tokens": 131072, diff --git a/tests/unit/test_model_prices_schema.py b/tests/unit/test_model_prices_schema.py index 50c9643683b..314a401e88e 100644 --- a/tests/unit/test_model_prices_schema.py +++ b/tests/unit/test_model_prices_schema.py @@ -605,6 +605,26 @@ DEEPSEEK_OFF_PEAK_WINDOWS: Final = ( {"hours_utc": "00:00-00:00", "weekdays": [6, 7]}, ) DEEPSEEK_HALVED_RATES: Final = ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost") +TENCENT_TOKENHUB_DEEPSEEK_PRICED_ROWS: Final = ( + "tencent/deepseek-v4-pro-0813", + "tencent/deepseek-v4.1-flash", +) +TENCENT_TOKENHUB_GLM_ROWS: Final = MappingProxyType( + { + "tencent/glm-5.3-flashx": (3.7e-07, 1.25e-06, 7.5e-08, True, True), + "tencent/glm-5.3-flash": (1.5e-07, 5e-07, 3e-08, True, True), + "tencent/glm-5.3": (1.4e-06, 4.4e-06, 2.6e-07, False, False), + "tencent/glm-5.2": (1.4e-06, 4.4e-06, 2.6e-07, False, False), + } +) +TENCENT_TOKENHUB_KIMI_ROWS: Final = MappingProxyType( + { + "tencent/kimi-k3": (3e-06, 1.5e-05, 3e-07, 1000000, False, False), + "tencent/kimi-k2.8-preview": (1e-06, 4e-06, 2.5e-07, 1000000, False, False), + "tencent/kimi-k2.7-code-highspeed": (1.9e-06, 8e-06, 3.8e-07, 256000, False, False), + "tencent/kimi-k2.7-code": (9.5e-07, 4e-06, 1.9e-07, 256000, False, False), + } +) def deepseek_off_peak_drift(entry: Mapping[str, object]) -> str | None: @@ -636,6 +656,133 @@ def test_deepseek_rows_bill_half_rate_outside_weekday_peak_hours(path: Path): assert all(name in rows for name in DEEPSEEK_PRICED_ROWS) +@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) +def test_tencent_tokenhub_deepseek_rows_bill_half_rate_outside_weekday_peak_hours(path: Path): + """Tencent TokenHub applies the DeepSeek weekday peak windows and half-rate off-peak prices + to these models (tencentcloud.com/document/product/1300/78937, read 2026-10-13).""" + rows: Mapping[str, object] = json.loads(path.read_text()) + drifted: Final = { + name: deepseek_off_peak_drift(entry) + for name in TENCENT_TOKENHUB_DEEPSEEK_PRICED_ROWS + if isinstance(entry := rows.get(name), dict) and deepseek_off_peak_drift(entry) is not None + } + invalid_timezones: Final = tuple( + name + for name in TENCENT_TOKENHUB_DEEPSEEK_PRICED_ROWS + if isinstance(entry := rows.get(name), dict) + and isinstance(block := entry.get("off_peak_pricing"), dict) + and block.get("weekday_timezone") != "Asia/Shanghai" + ) + assert drifted == {} + assert all(name in rows for name in TENCENT_TOKENHUB_DEEPSEEK_PRICED_ROWS) + assert invalid_timezones == () + + +@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) +def test_tencent_tokenhub_glm_rows_match_published_metadata(path: Path): + """TokenHub publishes these prices and capabilities in its official model catalog and rate card + (tencentcloud.com/document/product/1300/78934 and /78937, read 2026-10-13).""" + rows: Mapping[str, object] = json.loads(path.read_text()) + actual: Final = MappingProxyType( + { + name: ( + entry.get("input_cost_per_token"), + entry.get("output_cost_per_token"), + entry.get("cache_read_input_token_cost"), + entry.get("supports_vision", False), + entry.get("supports_video_input", False), + ) + for name in TENCENT_TOKENHUB_GLM_ROWS + if isinstance(entry := rows.get(name), dict) + } + ) + invalid_metadata: Final = tuple( + name + for name in TENCENT_TOKENHUB_GLM_ROWS + if not isinstance(entry := rows.get(name), dict) + or entry.get("litellm_provider") != "tencent" + or entry.get("max_input_tokens") != 1000000 + or entry.get("max_output_tokens") != 128000 + or entry.get("max_tokens") != 128000 + or entry.get("input_cost_per_token_cache_hit") != entry.get("cache_read_input_token_cost") + ) + assert actual == TENCENT_TOKENHUB_GLM_ROWS + assert invalid_metadata == () + + +@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) +def test_tencent_tokenhub_kimi_rows_match_published_metadata(path: Path): + """TokenHub publishes these Kimi prices and capabilities in its official model catalog and rate card + (tencentcloud.com/document/product/1300/78934 and /78937, read 2026-10-13).""" + rows: Mapping[str, object] = json.loads(path.read_text()) + actual: Final = MappingProxyType( + { + name: ( + entry.get("input_cost_per_token"), + entry.get("output_cost_per_token"), + entry.get("cache_read_input_token_cost"), + entry.get("max_tokens"), + entry.get("supports_vision", False), + entry.get("supports_video_input", False), + ) + for name in TENCENT_TOKENHUB_KIMI_ROWS + if isinstance(entry := rows.get(name), dict) + } + ) + invalid_metadata: Final = tuple( + name + for name, expected in TENCENT_TOKENHUB_KIMI_ROWS.items() + if not isinstance(entry := rows.get(name), dict) + or entry.get("litellm_provider") != "tencent" + or entry.get("max_input_tokens") != expected[3] + or entry.get("max_output_tokens") != expected[3] + or entry.get("input_cost_per_token_cache_hit") != entry.get("cache_read_input_token_cost") + or not entry.get("supports_reasoning") + or not entry.get("supports_response_schema") + or not entry.get("supports_function_calling") + or not entry.get("supports_prompt_caching") + ) + assert actual == TENCENT_TOKENHUB_KIMI_ROWS + assert invalid_metadata == () + + +class TestBedrockClaude35SonnetContextWindow: + @pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) + def test_bare_model_ids_match_us_and_eu_siblings(self, path: Path): + rows = json.loads(path.read_text()) + for model in ( + "anthropic.claude-3-5-sonnet-20240620-v1:0", + "anthropic.claude-3-5-sonnet-20241022-v2:0", + ): + assert rows[model]["max_input_tokens"] == rows[f"us.{model}"]["max_input_tokens"] + assert rows[model]["max_input_tokens"] == rows[f"eu.{model}"]["max_input_tokens"] + + @pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) + def test_bare_model_ids_have_200k_context_window(self, path: Path): + rows = json.loads(path.read_text()) + # Claude 3.5 Sonnet context window is 200K tokens, https://www.anthropic.com/news/claude-3-5-sonnet (accessed 2026-10-13) + for model in ( + "anthropic.claude-3-5-sonnet-20240620-v1:0", + "anthropic.claude-3-5-sonnet-20241022-v2:0", + ): + assert rows[model]["max_input_tokens"] == 200000 + + @pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) + def test_bare_model_ids_have_no_above_200k_pricing(self, path: Path): + rows = json.loads(path.read_text()) + above_200k_fields = ( + "input_cost_per_token_above_200k_tokens", + "output_cost_per_token_above_200k_tokens", + "cache_creation_input_token_cost_above_200k_tokens", + "cache_read_input_token_cost_above_200k_tokens", + ) + for model in ( + "anthropic.claude-3-5-sonnet-20240620-v1:0", + "anthropic.claude-3-5-sonnet-20241022-v2:0", + ): + assert not set(above_200k_fields).intersection(rows[model]) + + PROVIDER_LABELS_WITHOUT_A_MODEL_SET: Final = frozenset({"sagemaker", "bedrock_converse"}) MODES_SERVED_OUTSIDE_THE_LLM_PROVIDER_REGISTRY: Final = frozenset({"search", "decisions", "evaluation"}) VERTEX_FAMILIES_A_VERTEX_WILDCARD_GRANT_DOES_NOT_LIST: Final = frozenset( From d7b58126a7c1947df1a6b78fea2a5854a4610dd5 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 13 Oct 2026 13:11:59 -0700 Subject: [PATCH 008/325] feat(ui): source Home What's new from day 0 support blog posts (#46260) * feat(ui): source Home What's new from day 0 support blog posts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(ui): Home cards read card_title and card_summary from blog tags with provider logo Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/__init__.py | 1 - litellm/litellm_core_utils/get_blog_posts.py | 100 +++++------ .../public_endpoints/public_endpoints.py | 82 +-------- .../public_endpoints/public_endpoints.py | 19 -- tests/integration/security/_sinks.py | 1 - .../security/test_sweep_sensitivity.py | 2 +- .../test_blog_posts_endpoint.py | 57 +++--- .../public_endpoints/test_public_endpoints.py | 162 +----------------- tests/unit/test_get_blog_posts.py | 121 +++++-------- .../_components/HomePage.integration.test.tsx | 54 ++++-- .../(dashboard)/home/_components/HomePage.tsx | 28 +-- .../home/_components/homeContent.test.ts | 79 ++++++--- .../home/_components/homeContent.ts | 49 ++++-- .../home/_components/useWhatsNew.ts | 10 -- .../hooks/blogPosts/useBlogPosts.ts | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 78 +-------- 16 files changed, 280 insertions(+), 564 deletions(-) delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/home/_components/useWhatsNew.ts diff --git a/litellm/__init__.py b/litellm/__init__.py index 290d15bd07a..504993cb469 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -441,7 +441,6 @@ autorouter_presets_url: str = os.getenv( "LITELLM_AUTOROUTER_PRESETS_URL", "https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/proxy/public_endpoints/autorouter_presets.json", ) -whats_new_url: str | None = os.getenv("LITELLM_WHATS_NEW_URL") suppress_debug_info: bool = False dynamodb_table_name: Optional[str] = None s3_callback_params: Optional[Dict] = None diff --git a/litellm/litellm_core_utils/get_blog_posts.py b/litellm/litellm_core_utils/get_blog_posts.py index 0dd09064dcc..c8bc184f972 100644 --- a/litellm/litellm_core_utils/get_blog_posts.py +++ b/litellm/litellm_core_utils/get_blog_posts.py @@ -8,7 +8,6 @@ Disable remote fetching entirely: export LITELLM_LOCAL_BLOG_POSTS=True """ -import json import os import time import xml.etree.ElementTree as ET @@ -17,23 +16,53 @@ from importlib.resources import files from typing import Final import httpx +from pydantic import ConfigDict from litellm import verbose_logger from litellm.types.llms.base import LiteLLMBaseModel BLOG_POSTS_TTL_SECONDS: Final[int] = 3600 # 1 hour -BLOG_POSTS_MAX: Final[int] = 6 +BLOG_POSTS_MAX: Final[int] = 50 class BlogPost(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + title: str description: str date: str url: str + tags: tuple[str, ...] = () class BlogPostsResponse(LiteLLMBaseModel): - posts: list[BlogPost] + model_config = ConfigDict(frozen=True) + + posts: tuple[BlogPost, ...] + + +def _pub_date(pub_date_el: ET.Element | None) -> str: + if pub_date_el is None or not pub_date_el.text: + return "" + try: + return parsedate_to_datetime(pub_date_el.text).strftime("%Y-%m-%d") + except Exception: + return pub_date_el.text + + +def _parse_item(item: ET.Element) -> BlogPost | None: + title_el: Final = item.find("title") + link_el: Final = item.find("link") + if title_el is None or link_el is None: + return None + desc_el: Final = item.find("description") + return BlogPost( + title=title_el.text or "", + description=(desc_el.text or "") if desc_el is not None else "", + date=_pub_date(item.find("pubDate")), + url=link_el.text or "", + tags=tuple(category.text for category in item.findall("category") if category.text), + ) class GetBlogPosts: @@ -41,19 +70,20 @@ class GetBlogPosts: Fetches, validates, and caches LiteLLM blog posts. - Fetches RSS feed from docs site with a 5-second timeout - - Parses the XML and extracts the latest BLOG_POSTS_MAX blog posts + - Parses the XML and extracts up to BLOG_POSTS_MAX blog posts, with their RSS categories as tags - Caches the result in-process for BLOG_POSTS_TTL_SECONDS (1 hour) - Falls back to the bundled local backup on any failure """ - _cached_posts: list[dict[str, str]] | None = None + _cached_posts: tuple[BlogPost, ...] | None = None _last_fetch_time: float = 0.0 @staticmethod - def load_local_blog_posts() -> list[dict[str, str]]: + def load_local_blog_posts() -> tuple[BlogPost, ...]: """Load the bundled local backup blog posts.""" - content: Final = json.loads(files("litellm").joinpath("blog_posts.json").read_text(encoding="utf-8")) - return content.get("posts", []) + return BlogPostsResponse.model_validate_json( + files("litellm").joinpath("blog_posts.json").read_text(encoding="utf-8") + ).posts @staticmethod def fetch_rss_feed(url: str, timeout: int = 5) -> str: @@ -67,54 +97,24 @@ class GetBlogPosts: return response.text @staticmethod - def parse_rss_to_posts(xml_text: str, max_posts: int = BLOG_POSTS_MAX) -> list[dict[str, str]]: + def parse_rss_to_posts(xml_text: str, max_posts: int = BLOG_POSTS_MAX) -> tuple[BlogPost, ...]: """ - Parse RSS XML and return a list of blog post dicts. + Parse RSS XML into blog posts. - Extracts title, description, date (YYYY-MM-DD), and url from each . + Extracts title, description, date (YYYY-MM-DD), url and the tags from each . """ root: Final = ET.fromstring(xml_text) channel: Final = root.find("channel") if channel is None: raise ValueError("RSS feed missing element") - posts: Final[list[dict[str, str]]] = [] - for item in channel.findall("item"): - if len(posts) >= max_posts: - break - - title_el = item.find("title") - link_el = item.find("link") - desc_el = item.find("description") - pub_date_el = item.find("pubDate") - - if title_el is None or link_el is None: - continue - - # Parse RFC 2822 date to YYYY-MM-DD - date_str = "" - if pub_date_el is not None and pub_date_el.text: - try: - dt = parsedate_to_datetime(pub_date_el.text) - date_str = dt.strftime("%Y-%m-%d") - except Exception: - date_str = pub_date_el.text - - posts.append( - { - "title": title_el.text or "", - "description": desc_el.text or "" if desc_el is not None else "", - "date": date_str, - "url": link_el.text or "", - } - ) - - return posts + parsed: Final = (_parse_item(item) for item in channel.findall("item")) + return tuple(post for post in parsed if post is not None)[:max_posts] @staticmethod - def validate_blog_posts(posts: list[dict[str, str]]) -> bool: - """Return True if posts is a non-empty list.""" - if not isinstance(posts, list) or len(posts) == 0: + def validate_blog_posts(posts: tuple[BlogPost, ...]) -> bool: + """Return True if there is at least one post.""" + if len(posts) == 0: verbose_logger.warning( "LiteLLM: Parsed RSS feed has no valid posts. Falling back to local backup.", ) @@ -122,9 +122,9 @@ class GetBlogPosts: return True @classmethod - def get_blog_posts(cls, url: str) -> list[dict[str, str]]: + def get_blog_posts(cls, url: str) -> tuple[BlogPost, ...]: """ - Return the blog posts list. + Return the blog posts. Uses the in-process cache if within BLOG_POSTS_TTL_SECONDS. Fetches from ``url`` otherwise, falling back to local backup on failure. @@ -156,6 +156,6 @@ class GetBlogPosts: return posts -def get_blog_posts(url: str) -> list[dict[str, str]]: - """Public entry point — returns the blog posts list.""" +def get_blog_posts(url: str) -> tuple[BlogPost, ...]: + """Public entry point, returns the blog posts.""" return GetBlogPosts.get_blog_posts(url=url) diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 26bef3b3a61..1d3516b9ebd 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -6,15 +6,13 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from importlib.resources import files from typing import TYPE_CHECKING, Final, Protocol -import httpx from fastapi import APIRouter, HTTPException, Request -from pydantic import TypeAdapter, ValidationError +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_blog_posts import ( - BlogPost, BlogPostsResponse, GetBlogPosts, get_blog_posts, @@ -38,7 +36,6 @@ from litellm.types.proxy.public_endpoints.public_endpoints import ( ProviderCreateInfo, PublicModelHubInfo, SupportedEndpointsResponse, - WhatsNewResponse, ) from litellm.types.utils import LlmProviders @@ -479,21 +476,18 @@ async def get_litellm_model_cost_map(catalog_only: bool = False): tags=["public"], response_model=BlogPostsResponse, ) -async def get_litellm_blog_posts(): +async def get_litellm_blog_posts() -> BlogPostsResponse: """ - Public endpoint to get the latest LiteLLM blog posts. + Public endpoint to get the latest LiteLLM blog posts, each with its RSS categories as ``tags``. - Fetches from GitHub with a 1-hour in-process cache. + Fetches the docs RSS feed with a 1-hour in-process cache. Falls back to the bundled local backup on any failure. """ try: - posts_data = get_blog_posts(url=litellm.blog_posts_url) + return BlogPostsResponse(posts=get_blog_posts(url=litellm.blog_posts_url)) except Exception as e: verbose_logger.warning("LiteLLM: get_litellm_blog_posts endpoint fallback triggered: %s", str(e)) - posts_data = GetBlogPosts.load_local_blog_posts() - - posts: Final = [BlogPost(**p) for p in posts_data] - return BlogPostsResponse(posts=posts) + return BlogPostsResponse(posts=GetBlogPosts.load_local_blog_posts()) _AUTOROUTER_PRESETS_ADAPTER: Final = TypeAdapter(dict[str, AutoRouterPresetRecord]) @@ -581,70 +575,6 @@ async def get_public_autorouter_presets() -> Mapping[str, AutoRouterPresetRecord return await get_autorouter_presets(url=litellm.autorouter_presets_url) -def _load_bundled_whats_new() -> WhatsNewResponse: - return WhatsNewResponse.model_validate_json( - files("litellm.proxy.public_endpoints").joinpath("whats_new.json").read_text(encoding="utf-8") - ) - - -async def _fetch_remote_whats_new(url: str) -> WhatsNewResponse: - return WhatsNewResponse.model_validate_json(await _fetch_remote_bytes(url)) - - -async def _resolve_whats_new(url: str | None, fetch: Callable[[str], Awaitable[WhatsNewResponse]]) -> WhatsNewResponse: - if url is None or os.getenv("LITELLM_LOCAL_WHATS_NEW", "").lower() == "true": - return _load_bundled_whats_new() - try: - return await fetch(url) - except (httpx.HTTPError, ValidationError, OSError) as e: - verbose_logger.warning( - "LiteLLM: failed to fetch What's new launches from %s: %s. Serving the bundled list for the life of this process.", - url, - str(e), - ) - return _load_bundled_whats_new() - - -class _WhatsNewCache: - launches: WhatsNewResponse | None = None - lock: asyncio.Lock | None = None - - -async def get_whats_new( - url: str | None, - fetch: Callable[[str], Awaitable[WhatsNewResponse]] = _fetch_remote_whats_new, -) -> WhatsNewResponse: - cached: Final = _WhatsNewCache.launches - if cached is not None: - return cached - if _WhatsNewCache.lock is None: - _WhatsNewCache.lock = asyncio.Lock() - async with _WhatsNewCache.lock: - held: Final = _WhatsNewCache.launches - if held is not None: - return held - resolved: Final = await _resolve_whats_new(url=url, fetch=fetch) - _WhatsNewCache.launches = resolved - return resolved - - -@router.get( - "/public/whats_new", - tags=["public"], - response_model=WhatsNewResponse, -) -async def get_public_whats_new() -> WhatsNewResponse: - """ - Return the launches the dashboard Home page shows under What's new. - - Serves the list bundled with this release. Set ``LITELLM_WHATS_NEW_URL`` to fetch a remote list instead, - resolved once per process on the first request and falling back to the bundled list on any failure, so an - airgapped proxy serves the bundled list without retrying. ``LITELLM_LOCAL_WHATS_NEW=True`` forces the - bundled list even when a URL is set. A restart picks up a newly published list. - """ - return await get_whats_new(url=litellm.whats_new_url) - - @router.get( "/public/endpoints", tags=["public"], diff --git a/litellm/types/proxy/public_endpoints/public_endpoints.py b/litellm/types/proxy/public_endpoints/public_endpoints.py index 6e4adb4bbe7..c487ce8b6c5 100644 --- a/litellm/types/proxy/public_endpoints/public_endpoints.py +++ b/litellm/types/proxy/public_endpoints/public_endpoints.py @@ -1,5 +1,4 @@ from collections.abc import Mapping, Sequence -from datetime import date from typing import Any, Literal from pydantic import ConfigDict @@ -114,24 +113,6 @@ class AutoRouterPresetRecord(LiteLLMBaseModel): complexity_router_config: AutoRouterPresetConfig -class WhatsNewLaunch(LiteLLMBaseModel): - """One launch card in the dashboard Home page's What's new section.""" - - model_config = ConfigDict(frozen=True) - - icon: str - title: str - description: str - href: str - published_on: date - - -class WhatsNewResponse(LiteLLMBaseModel): - model_config = ConfigDict(frozen=True) - - launches: tuple[WhatsNewLaunch, ...] - - class ComplexityScorerDefaults(LiteLLMBaseModel): """The complexity router's shipped heuristic scorer defaults. diff --git a/tests/integration/security/_sinks.py b/tests/integration/security/_sinks.py index 1aebaba1621..906384e8760 100644 --- a/tests/integration/security/_sinks.py +++ b/tests/integration/security/_sinks.py @@ -80,7 +80,6 @@ LOCAL_CATALOGS: Final = MappingProxyType( "LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "LITELLM_LOCAL_AUTOROUTER_PRESETS", "LITELLM_LOCAL_POLICY_TEMPLATES", - "LITELLM_LOCAL_WHATS_NEW", ) } ) diff --git a/tests/integration/security/test_sweep_sensitivity.py b/tests/integration/security/test_sweep_sensitivity.py index 46c004ad45d..d3377093965 100644 --- a/tests/integration/security/test_sweep_sensitivity.py +++ b/tests/integration/security/test_sweep_sensitivity.py @@ -85,7 +85,7 @@ def test_rig_with_an_overridden_master_key_resolves_the_config_deployment(tmp_pa assert overridden.proxy.key == master_key assert overridden.model_id assert overridden.proxy.request("GET", "/model/info").status_code == 200 - assert overridden.proxy.request("GET", "/public/whats_new").status_code == 200 + assert overridden.proxy.request("GET", "/public/litellm_blog_posts").status_code == 200 assert overridden.egress() == () diff --git a/tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py b/tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py index dbbc31279c3..a771b1ccdb0 100644 --- a/tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py +++ b/tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py @@ -5,14 +5,17 @@ from unittest.mock import patch import pytest from fastapi.testclient import TestClient -SAMPLE_POSTS = [ - { - "title": "Test Post", - "description": "A test post.", - "date": "2026-01-01", - "url": "https://www.litellm.ai/blog/test", - } -] +from litellm.litellm_core_utils.get_blog_posts import BlogPost, GetBlogPosts + +SAMPLE_POSTS = ( + BlogPost( + title="Test Post", + description="A test post.", + date="2026-01-01", + url="https://www.litellm.ai/blog/test", + tags=("anthropic", "day 0 support"), + ), +) @pytest.fixture @@ -35,27 +38,25 @@ def test_get_blog_posts_returns_response_shape(client): response = client.get("/public/litellm_blog_posts") assert response.status_code == 200 - data = response.json() - assert "posts" in data - assert len(data["posts"]) == 1 - post = data["posts"][0] - assert post["title"] == "Test Post" - assert post["description"] == "A test post." - assert post["date"] == "2026-01-01" - assert post["url"] == "https://www.litellm.ai/blog/test" + assert response.json() == { + "posts": [ + { + "title": "Test Post", + "description": "A test post.", + "date": "2026-01-01", + "url": "https://www.litellm.ai/blog/test", + "tags": ["anthropic", "day 0 support"], + } + ] + } def test_get_blog_posts_returns_every_fetched_post(client): """The parser owns the cap (BLOG_POSTS_MAX), so the endpoint passes every fetched post through.""" - many_posts = [ - { - "title": f"Post {i}", - "description": "desc", - "date": "2026-01-01", - "url": f"https://www.litellm.ai/blog/{i}", - } + many_posts = tuple( + BlogPost(title=f"Post {i}", description="desc", date="2026-01-01", url=f"https://www.litellm.ai/blog/{i}") for i in range(10) - ] + ) with patch( "litellm.proxy.public_endpoints.public_endpoints.get_blog_posts", @@ -68,14 +69,14 @@ def test_get_blog_posts_returns_every_fetched_post(client): def test_get_blog_posts_returns_local_backup_on_failure(client): - """Endpoint returns local backup (non-empty list) when fetcher fails.""" + """Endpoint returns the local backup, with empty tags, when the fetcher raises.""" with patch( "litellm.proxy.public_endpoints.public_endpoints.get_blog_posts", side_effect=Exception("fetch failed"), ): response = client.get("/public/litellm_blog_posts") - # Should not 500 — returns local backup assert response.status_code == 200 - assert "posts" in response.json() - assert len(response.json()["posts"]) > 0 + backup = [post.model_dump(mode="json") for post in GetBlogPosts.load_local_blog_posts()] + assert response.json()["posts"] == backup + assert all(post["tags"] == [] for post in response.json()["posts"]) diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index 4f1cf2d8c28..129dbaf9d5e 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -1,9 +1,6 @@ import json -import os import re -from collections.abc import Iterator from datetime import datetime, timezone -from importlib.resources import files from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -19,7 +16,7 @@ from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_pres from litellm.types.proxy.management_endpoints.model_management_endpoints import ( ModelGroupInfoProxy, ) -from litellm.types.proxy.public_endpoints.public_endpoints import ProviderCreateInfo, WhatsNewResponse +from litellm.types.proxy.public_endpoints.public_endpoints import ProviderCreateInfo from litellm.types.utils import LlmProviders @@ -78,7 +75,6 @@ def test_get_provider_create_fields(): def test_get_litellm_model_cost_map_catalog_only_excludes_runtime_registered_entries( monkeypatch: pytest.MonkeyPatch, ) -> None: - import litellm from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) @@ -1538,159 +1534,3 @@ async def test_fetch_remote_autorouter_presets_parses_and_rejects_empty(monkeypa response.content = b"{}" with pytest.raises(ValueError, match="empty"): await _fetch_remote_autorouter_presets("https://example.test/presets.json") - - -@pytest.fixture -def reset_whats_new_cache() -> Iterator[None]: - from litellm.proxy.public_endpoints.public_endpoints import _WhatsNewCache - - _WhatsNewCache.launches = None - _WhatsNewCache.lock = None - yield - _WhatsNewCache.launches = None - _WhatsNewCache.lock = None - - -def _remote_whats_new() -> WhatsNewResponse: - return WhatsNewResponse.model_validate( - { - "launches": [ - { - "icon": "sparkles", - "title": "Remote Launch", - "description": "from the remote list", - "href": "https://docs.litellm.ai/blog/remote", - "published_on": "2026-10-11", - } - ] - } - ) - - -def test_get_whats_new_local_mode_serves_bundled_list( - monkeypatch: pytest.MonkeyPatch, reset_whats_new_cache: None -) -> None: - monkeypatch.setenv("LITELLM_LOCAL_WHATS_NEW", "True") - app = FastAPI() - app.include_router(router) - client = TestClient(app) - - response = client.get("/public/whats_new") - - assert response.status_code == 200 - served: Final = WhatsNewResponse.model_validate_json(response.content) - bundled: Final = WhatsNewResponse.model_validate_json( - files("litellm.proxy.public_endpoints").joinpath("whats_new.json").read_text(encoding="utf-8") - ) - assert served == bundled - assert len(served.launches) > 0 - - -@pytest.mark.asyncio -async def test_get_whats_new_serves_bundled_list_without_fetching_when_no_url_is_configured( - monkeypatch: pytest.MonkeyPatch, reset_whats_new_cache: None -) -> None: - from litellm.proxy.public_endpoints.public_endpoints import _load_bundled_whats_new, get_whats_new - - monkeypatch.delenv("LITELLM_LOCAL_WHATS_NEW", raising=False) - - async def forbidden_fetch(url: str) -> WhatsNewResponse: - raise AssertionError(f"no URL is configured, so nothing may be fetched, got {url}") - - served = await get_whats_new(url=None, fetch=forbidden_fetch) - - assert served == _load_bundled_whats_new() - assert len(served.launches) > 0 - - -def test_whats_new_url_default_mirrors_the_env_var_so_no_remote_url_is_baked_in() -> None: - assert litellm.whats_new_url == os.environ.get("LITELLM_WHATS_NEW_URL") - - -@pytest.mark.asyncio -async def test_public_whats_new_route_serves_the_bundled_list_without_fetching_when_no_url_is_configured( - monkeypatch: pytest.MonkeyPatch, reset_whats_new_cache: None -) -> None: - from litellm.proxy.public_endpoints import public_endpoints - from litellm.proxy.public_endpoints.public_endpoints import _load_bundled_whats_new, get_public_whats_new - - monkeypatch.delenv("LITELLM_LOCAL_WHATS_NEW", raising=False) - - async def forbidden_fetch(url: str) -> bytes: - raise AssertionError(f"no URL is configured, so the route may not fetch, got {url}") - - monkeypatch.setattr(public_endpoints, "_fetch_remote_bytes", forbidden_fetch) - - served = await get_public_whats_new() - - assert served == _load_bundled_whats_new() - assert len(served.launches) > 0 - - -@pytest.mark.asyncio -async def test_get_whats_new_fetches_once_per_process( - monkeypatch: pytest.MonkeyPatch, reset_whats_new_cache: None -) -> None: - from litellm.proxy.public_endpoints.public_endpoints import get_whats_new - - monkeypatch.delenv("LITELLM_LOCAL_WHATS_NEW", raising=False) - remote = _remote_whats_new() - calls: Final[list[str]] = [] - - async def fake_fetch(url: str) -> WhatsNewResponse: - calls.append(url) - return remote - - first = await get_whats_new(url="https://example.test/whats_new.json", fetch=fake_fetch) - second = await get_whats_new(url="https://example.test/whats_new.json", fetch=fake_fetch) - - assert first == remote - assert second == remote - assert calls == ["https://example.test/whats_new.json"] - - -@pytest.mark.asyncio -async def test_get_whats_new_caches_bundled_fallback_when_remote_is_unreachable( - monkeypatch: pytest.MonkeyPatch, reset_whats_new_cache: None -) -> None: - from litellm.proxy.public_endpoints.public_endpoints import _load_bundled_whats_new, get_whats_new - - monkeypatch.delenv("LITELLM_LOCAL_WHATS_NEW", raising=False) - calls: Final[list[str]] = [] - - async def unreachable_fetch(url: str) -> WhatsNewResponse: - calls.append(url) - raise OSError("network is unreachable") - - first = await get_whats_new(url="https://example.test/whats_new.json", fetch=unreachable_fetch) - second = await get_whats_new(url="https://example.test/whats_new.json", fetch=unreachable_fetch) - - assert first == _load_bundled_whats_new() - assert second == first - assert len(calls) == 1 - - -@pytest.mark.asyncio -async def test_fetch_remote_whats_new_rejects_a_malformed_list_so_the_bundled_one_is_served( - monkeypatch: pytest.MonkeyPatch, reset_whats_new_cache: None -) -> None: - import httpx - - from litellm.proxy.public_endpoints import public_endpoints - - monkeypatch.delenv("LITELLM_LOCAL_WHATS_NEW", raising=False) - request = httpx.Request("GET", "https://example.test/whats_new.json") - client = MagicMock() - client.get = AsyncMock( - return_value=httpx.Response( - 200, - request=request, - content=b'{"launches": [{"icon": "box", "title": "No date", "description": "d", "href": "https://x"}]}', - ) - ) - - with patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=client): - served = await public_endpoints.get_whats_new(url="https://example.test/whats_new.json") - - assert served == public_endpoints._load_bundled_whats_new() - client.get.assert_awaited_once_with("https://example.test/whats_new.json", timeout=5.0) diff --git a/tests/unit/test_get_blog_posts.py b/tests/unit/test_get_blog_posts.py index 5686f45efd3..499dd9bcf42 100644 --- a/tests/unit/test_get_blog_posts.py +++ b/tests/unit/test_get_blog_posts.py @@ -25,6 +25,8 @@ SAMPLE_RSS = """\ https://docs.litellm.ai/blog/test A test post. Wed, 01 Jan 2026 10:00:00 GMT + anthropic + day 0 support Second Post @@ -36,6 +38,8 @@ SAMPLE_RSS = """\ """ +CACHED_POST = BlogPost(title="Cached", description="D", date="2026-01-01", url="https://x.com") + @pytest.fixture(autouse=True) def reset_blog_posts_cache(): @@ -46,30 +50,32 @@ def reset_blog_posts_cache(): GetBlogPosts._last_fetch_time = 0.0 -def test_load_local_blog_posts_returns_list(): +def test_load_local_blog_posts_returns_posts_without_tags(): posts = GetBlogPosts.load_local_blog_posts() - assert isinstance(posts, list) assert len(posts) > 0 - first = posts[0] - assert "title" in first - assert "description" in first - assert "date" in first - assert "url" in first + assert all(isinstance(post, BlogPost) for post in posts) + assert posts[0].tags == () def test_parse_rss_to_posts(): posts = GetBlogPosts.parse_rss_to_posts(SAMPLE_RSS, max_posts=1) assert len(posts) == 1 - assert posts[0]["title"] == "Test Post" - assert posts[0]["url"] == "https://docs.litellm.ai/blog/test" - assert posts[0]["description"] == "A test post." - assert posts[0]["date"] == "2026-01-01" + assert posts[0].title == "Test Post" + assert posts[0].url == "https://docs.litellm.ai/blog/test" + assert posts[0].description == "A test post." + assert posts[0].date == "2026-01-01" + + +def test_parse_rss_to_posts_reads_categories_as_tags(): + posts = GetBlogPosts.parse_rss_to_posts(SAMPLE_RSS) + assert posts[0].tags == ("anthropic", "day 0 support") + assert posts[1].tags == () def test_parse_rss_to_posts_multiple(): posts = GetBlogPosts.parse_rss_to_posts(SAMPLE_RSS, max_posts=5) assert len(posts) == 2 - assert posts[1]["title"] == "Second Post" + assert posts[1].title == "Second Post" def test_parse_rss_to_posts_default_cap(): @@ -80,8 +86,16 @@ def test_parse_rss_to_posts_default_cap(): ) posts = GetBlogPosts.parse_rss_to_posts(f"{items}") assert len(posts) == BLOG_POSTS_MAX - assert posts[0]["title"] == "Post 0" - assert posts[-1]["title"] == f"Post {BLOG_POSTS_MAX - 1}" + assert posts[0].title == "Post 0" + assert posts[-1].title == f"Post {BLOG_POSTS_MAX - 1}" + + +def test_parse_rss_to_posts_skips_items_without_title_or_link_before_capping(): + items = "no title" + "".join( + f"Post {i}https://docs.litellm.ai/blog/{i}" for i in range(2) + ) + posts = GetBlogPosts.parse_rss_to_posts(f"{items}", max_posts=2) + assert [post.title for post in posts] == ["Post 0", "Post 1"] def test_parse_rss_to_posts_invalid_xml(): @@ -95,16 +109,11 @@ def test_parse_rss_to_posts_missing_channel(): def test_validate_blog_posts_valid(): - posts = [{"title": "T", "description": "D", "date": "2026-01-01", "url": "https://x.com"}] - assert GetBlogPosts.validate_blog_posts(posts) is True + assert GetBlogPosts.validate_blog_posts((CACHED_POST,)) is True -def test_validate_blog_posts_empty_list(): - assert GetBlogPosts.validate_blog_posts([]) is False - - -def test_validate_blog_posts_not_list(): - assert GetBlogPosts.validate_blog_posts("not a list") is False +def test_validate_blog_posts_empty(): + assert GetBlogPosts.validate_blog_posts(()) is False def test_get_blog_posts_success(): @@ -120,7 +129,8 @@ def test_get_blog_posts_success(): posts = get_blog_posts(url=litellm.blog_posts_url) assert len(posts) == 2 - assert posts[0]["title"] == "Test Post" + assert posts[0].title == "Test Post" + assert posts[0].tags == ("anthropic", "day 0 support") def test_get_blog_posts_network_error_falls_back_to_local(): @@ -131,8 +141,7 @@ def test_get_blog_posts_network_error_falls_back_to_local(): ): posts = get_blog_posts(url=litellm.blog_posts_url) - assert isinstance(posts, list) - assert len(posts) > 0 + assert posts == GetBlogPosts.load_local_blog_posts() def test_get_blog_posts_invalid_xml_falls_back_to_local(): @@ -147,52 +156,25 @@ def test_get_blog_posts_invalid_xml_falls_back_to_local(): ): posts = get_blog_posts(url=litellm.blog_posts_url) - assert isinstance(posts, list) - assert len(posts) > 0 + assert posts == GetBlogPosts.load_local_blog_posts() def test_get_blog_posts_ttl_cache_not_refetched(): """Within TTL window, does not re-fetch.""" - cached = [ - { - "title": "Cached", - "description": "D", - "date": "2026-01-01", - "url": "https://x.com", - } - ] - GetBlogPosts._cached_posts = cached - GetBlogPosts._last_fetch_time = time.time() # just now + GetBlogPosts._cached_posts = (CACHED_POST,) + GetBlogPosts._last_fetch_time = time.time() - call_count = 0 - - def mock_get(*args, **kwargs): - nonlocal call_count - call_count += 1 - m = MagicMock() - m.text = SAMPLE_RSS - m.raise_for_status = MagicMock() - return m - - with patch("litellm.litellm_core_utils.get_blog_posts.httpx.get", side_effect=mock_get): + with patch("litellm.litellm_core_utils.get_blog_posts.httpx.get") as mock_get: posts = get_blog_posts(url=litellm.blog_posts_url) - assert call_count == 0 # cache hit, no fetch - assert len(posts) == 1 + mock_get.assert_not_called() + assert posts == (CACHED_POST,) def test_get_blog_posts_ttl_expired_refetches(): """After TTL window, re-fetches from remote.""" - cached = [ - { - "title": "Cached", - "description": "D", - "date": "2026-01-01", - "url": "https://x.com", - } - ] - GetBlogPosts._cached_posts = cached - GetBlogPosts._last_fetch_time = time.time() - 7200 # 2 hours ago + GetBlogPosts._cached_posts = (CACHED_POST,) + GetBlogPosts._last_fetch_time = time.time() - 7200 mock_response = MagicMock() mock_response.text = SAMPLE_RSS @@ -213,20 +195,9 @@ def test_get_blog_posts_local_env_var_skips_remote(monkeypatch): with patch("litellm.litellm_core_utils.get_blog_posts.httpx.get") as mock_get: posts = get_blog_posts(url=litellm.blog_posts_url) mock_get.assert_not_called() - assert isinstance(posts, list) - assert len(posts) > 0 + assert posts == GetBlogPosts.load_local_blog_posts() -def test_blog_post_pydantic_model(): - post = BlogPost( - title="T", - description="D", - date="2026-01-01", - url="https://example.com", - ) - assert post.title == "T" - - -def test_blog_posts_response_pydantic_model(): - resp = BlogPostsResponse(posts=[BlogPost(title="T", description="D", date="2026-01-01", url="https://x.com")]) - assert len(resp.posts) == 1 +def test_blog_posts_response_serializes_tags_as_a_list(): + resp = BlogPostsResponse(posts=(CACHED_POST.model_copy(update={"tags": ("day 0 support",)}),)) + assert resp.model_dump(mode="json")["posts"][0]["tags"] == ["day 0 support"] diff --git a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx index 4e87cd06f7a..598fbf74efc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx @@ -33,17 +33,15 @@ vi.mock("@/components/networking", async (importOriginal) => ({ getProxyBaseUrl: () => "http://proxy.test", })); -const blogPosts = { posts: [{ title: "Post one", description: "d", date: "2026-10-09", url: "https://x/1" }] }; - -const launch = (title: string, icon: string, published_on: string) => ({ - icon, +const post = (title: string, date: string, tags: string[] = []) => ({ title, description: `${title} description`, - href: `https://docs.litellm.ai/blog/${title}`, - published_on, + date, + url: title === "Post one" ? "https://x/1" : `https://docs.litellm.ai/blog/${title}`, + tags, }); -const whatsNew = vi.hoisted(() => ({ launches: [] as unknown[] })); +const blogPosts = vi.hoisted(() => ({ posts: [] as unknown[] })); const requestUrl = (input: RequestInfo | URL): string => { if (typeof input === "string") return input; @@ -57,9 +55,8 @@ const fakeFetch = (input: RequestInfo | URL) => { if (requestUrl(input).endsWith("/sso/get/ui_settings")) { return Promise.resolve(Response.json({ PROXY_BASE_URL: "http://proxy.test" })); } - const body = requestUrl(input).endsWith("/public/whats_new") ? whatsNew : blogPosts; return Promise.resolve( - new Response(JSON.stringify(body), { status: 200, headers: { "Content-Type": "application/json" } }), + new Response(JSON.stringify(blogPosts), { status: 200, headers: { "Content-Type": "application/json" } }), ); }; @@ -118,7 +115,7 @@ describe("HomePage", () => { network.aggregated.mockReset().mockResolvedValue(activity); network.gateway.mockReset().mockResolvedValue({ total_successful_requests: 0, total_failed_requests: 0 }); network.userInfo.mockReset().mockResolvedValue({ user_info: { max_budget: null } }); - whatsNew.launches = [launch("mid", "box", "2026-10-06"), launch("new", "zap", "2026-10-11")]; + blogPosts.posts = [post("Post one", "2026-10-09"), post("mid", "2026-10-06", ["day 0 support"])]; vi.stubGlobal("fetch", vi.fn(fakeFetch)); }); afterEach(() => { @@ -168,11 +165,13 @@ describe("HomePage", () => { expect(network.aggregated).toHaveBeenCalledTimes(2); }); - it("renders the launches served by /public/whats_new newest first, whatever order the JSON lists them in", async () => { - whatsNew.launches = [ - launch("mid", "box", "2026-10-06"), - launch("new", "an-icon-from-a-newer-list", "2026-10-11"), - launch("old", "scale", "2026-09-30"), + it("renders the three newest day 0 support blog posts as What's new, skipping untagged posts", async () => { + blogPosts.posts = [ + post("mid", "2026-10-06", ["openai", "day 0 support"]), + post("townhall", "2026-10-12", ["announcement"]), + post("new", "2026-10-11", ["day 0 support"]), + post("old", "2026-09-30", ["day 0 support"]), + post("oldest", "2026-09-01", ["day 0 support"]), ]; renderHome(); const section = await screen.findByTestId("home-whats-new"); @@ -187,10 +186,31 @@ describe("HomePage", () => { "https://docs.litellm.ai/blog/old", ]); expect(within(section).getByText("Oct 11")).toBeInTheDocument(); + expect(within(section).queryByText("townhall")).not.toBeInTheDocument(); }); - it("hides What's new when /public/whats_new has no launches", async () => { - whatsNew.launches = []; + it("renders card_title, card_summary and the provider logo, falling back to the title and first sentence", async () => { + blogPosts.posts = [ + post("Gemini 4 Argon API Pricing and Day 0 Support on LiteLLM", "2026-10-13", [ + "gemini", + "card_title:Gemini 4 Argon", + "card_summary:Gemini 4 Argon costs $2.00 per 1M input tokens.", + "day 0 support", + ]), + post("Day 0 Support: Grok 4.7", "2026-09-21", ["nano banana", "day 0 support"]), + ]; + renderHome(); + const section = await screen.findByTestId("home-whats-new"); + expect(await within(section).findByText("Gemini 4 Argon")).toBeInTheDocument(); + expect(within(section).getByText("Gemini 4 Argon costs $2.00 per 1M input tokens.")).toBeInTheDocument(); + expect(within(section).getByRole("img", { name: "gemini logo" })).toBeInTheDocument(); + expect(within(section).getByText("Day 0 Support: Grok 4.7")).toBeInTheDocument(); + expect(within(section).getByText("Day 0 Support: Grok 4.7 description")).toBeInTheDocument(); + expect(within(section).getAllByRole("img")).toHaveLength(1); + }); + + it("hides What's new when no blog post is tagged day 0 support", async () => { + blogPosts.posts = [post("Post one", "2026-10-09", ["announcement"])]; renderHome(); expect(await screen.findByRole("link", { name: /post one/i })).toBeInTheDocument(); expect(screen.queryByTestId("home-whats-new")).not.toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.tsx index dcb0cd771b6..b0c63bb5dc0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.tsx @@ -1,6 +1,6 @@ "use client"; -import { ArrowRight, Plus, Sparkles } from "lucide-react"; +import { ArrowRight, Plus, Rocket, Sparkles } from "lucide-react"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { Button } from "@/components/ui/button"; import { Card } from "@/components/ui/card"; @@ -12,17 +12,19 @@ import { LiteAdminHome } from "@/components/liteadmin/LiteAdmin"; import { Leaderboard, useBreakdown } from "@/app/(dashboard)/usage/_components/components/overview/BreakdownChart"; import { Panel } from "@/app/(dashboard)/usage/_components/components/overview/Primitives"; import UsageStatStrip from "@/app/(dashboard)/usage/_components/components/overview/UsageStatStrip"; -import { formatPublishedOn, launchIcon } from "./homeContent"; +import { useBlogPosts } from "@/app/(dashboard)/hooks/blogPosts/useBlogPosts"; +import { cardProvider, cardSummary, cardTitle, formatPublishedOn, latestLaunches } from "./homeContent"; +import { Logo } from "@/components/molecules/logo/Logo"; import UpdatesFeed from "./UpdatesFeed"; import { HOME_USAGE_DAYS, useHomeUsage } from "./useHomeUsage"; -import { useWhatsNew } from "./useWhatsNew"; const LEADERBOARD = { metric: "tokens", dimension: "model_groups" } as const; export default function HomePage() { const { isViewOnly } = useAuthorized(); const usage = useHomeUsage(); - const whatsNew = useWhatsNew(); + const blog = useBlogPosts(); + const launches = latestLaunches(blog.data?.posts ?? []); const { series, ranking } = useBreakdown(usage.results, LEADERBOARD, 8); return ( @@ -62,7 +64,7 @@ export default function HomePage() {
- {(whatsNew.isLoading || (whatsNew.data?.length ?? 0) > 0) && ( + {(blog.isLoading || launches.length > 0) && (

@@ -71,22 +73,22 @@ export default function HomePage() {

- {whatsNew.isLoading && + {blog.isLoading && [0, 1, 2].map((i) => ( ))} - {whatsNew.data?.map((item) => { - const Icon = launchIcon(item.icon); + {launches.map((post) => { + const provider = cardProvider(post); return ( - +
- + {provider ? : }
-
{item.title}
-

{item.description}

+
{cardTitle(post)}
+

{cardSummary(post)}

- {formatPublishedOn(item.published_on)} + {formatPublishedOn(post.date)}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/homeContent.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/homeContent.test.ts index 7fad5767d7f..5757b26a6cf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/homeContent.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/homeContent.test.ts @@ -1,40 +1,75 @@ -import { Box, Sparkles } from "lucide-react"; import { describe, expect, it } from "vitest"; -import { latestLaunches, launchIcon, type WhatsNewLaunch } from "./homeContent"; +import type { BlogPost } from "@/app/(dashboard)/hooks/blogPosts/useBlogPosts"; +import { DAY_ZERO_TAG, cardProvider, cardSummary, cardTitle, firstSentence, latestLaunches } from "./homeContent"; -const launch = (title: string, published_on: string): WhatsNewLaunch => ({ - icon: "box", +const post = (title: string, date: string, tags: string[] = [DAY_ZERO_TAG]): BlogPost => ({ title, description: "d", - href: `https://docs.litellm.ai/${title}`, - published_on, + date, + url: `https://docs.litellm.ai/blog/${title}`, + tags, }); describe("latestLaunches", () => { - it("orders launches by published_on, newest first, without mutating the input", () => { - const launches = [launch("old", "2026-09-01"), launch("new", "2026-10-11"), launch("mid", "2026-10-01")]; - expect(latestLaunches(launches).map((l) => l.title)).toEqual(["new", "mid", "old"]); - expect(launches.map((l) => l.title)).toEqual(["old", "new", "mid"]); + it("keeps only posts tagged day 0 support", () => { + const posts = [ + post("townhall", "2026-10-12", ["announcement"]), + post("argon", "2026-10-13"), + post("untagged", "2026-10-11", []), + ]; + expect(latestLaunches(posts).map((p) => p.title)).toEqual(["argon"]); + }); + + it("orders launches by date, newest first, without mutating the input", () => { + const posts = [post("old", "2026-09-01"), post("new", "2026-10-11"), post("mid", "2026-10-01")]; + expect(latestLaunches(posts).map((p) => p.title)).toEqual(["new", "mid", "old"]); + expect(posts.map((p) => p.title)).toEqual(["old", "new", "mid"]); }); it("keeps only the three newest so the cards fill one row", () => { - const launches = [ - launch("oldest", "2026-08-01"), - launch("a", "2026-10-03"), - launch("b", "2026-10-02"), - launch("c", "2026-10-01"), + const posts = [ + post("oldest", "2026-08-01"), + post("a", "2026-10-03"), + post("b", "2026-10-02"), + post("c", "2026-10-01"), ]; - expect(latestLaunches(launches).map((l) => l.title)).toEqual(["a", "b", "c"]); + expect(latestLaunches(posts).map((p) => p.title)).toEqual(["a", "b", "c"]); }); }); -describe("launchIcon", () => { - it("maps a known icon name to its lucide icon", () => { - expect(launchIcon("box")).toBe(Box); +describe("card fields", () => { + const argon = post("Gemini 4 Argon API Pricing and Day 0 Support on LiteLLM", "2026-10-13", [ + "gemini", + "gemini 4", + "card_title:Gemini 4 Argon", + "card_summary:Gemini 4 Argon costs $2.00 per 1M input tokens.", + DAY_ZERO_TAG, + ]); + const legacy = { + ...post("Day 0 Support: Grok 4.7", "2026-09-21", ["xai", DAY_ZERO_TAG]), + description: "Day 0 support for Grok 4.7 on LiteLLM, at the same price as Grok 4.6. Use it anywhere.", + }; + + it("reads card_title and card_summary from the tags", () => { + expect(cardTitle(argon)).toBe("Gemini 4 Argon"); + expect(cardSummary(argon)).toBe("Gemini 4 Argon costs $2.00 per 1M input tokens."); }); - it("falls back to Sparkles for an icon name this dashboard version does not know", () => { - expect(launchIcon("hologram")).toBe(Sparkles); - expect(launchIcon("constructor")).toBe(Sparkles); + it("falls back to the full title and the first sentence of the description", () => { + expect(cardTitle(legacy)).toBe("Day 0 Support: Grok 4.7"); + expect(cardSummary(legacy)).toBe("Day 0 support for Grok 4.7 on LiteLLM, at the same price as Grok 4.6."); + expect(firstSentence("No terminal punctuation here")).toBe("No terminal punctuation here"); + expect(firstSentence("Costs $0.10 per 1M tokens. Twice.")).toBe("Costs $0.10 per 1M tokens."); + }); + + it("treats an empty card tag as absent", () => { + const blank = post("Full title", "2026-10-01", ["card_title: ", DAY_ZERO_TAG]); + expect(cardTitle(blank)).toBe("Full title"); + }); + + it("picks the provider logo from the first tag that is a litellm provider", () => { + expect(cardProvider(argon)).toBe("gemini"); + expect(cardProvider(legacy)).toBe("xai"); + expect(cardProvider(post("x", "2026-10-01", ["nano banana", DAY_ZERO_TAG]))).toBeUndefined(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/homeContent.ts b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/homeContent.ts index c3eb4b2c838..37549b48588 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/homeContent.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/homeContent.ts @@ -1,26 +1,37 @@ -import { Bot, Box, Gauge, Plug, Rocket, Scale, Shield, Sparkles, Zap, type LucideIcon } from "lucide-react"; -import type { components } from "@/lib/http/schema"; +import type { BlogPost } from "@/app/(dashboard)/hooks/blogPosts/useBlogPosts"; +import { getProviderLogoAndName } from "@/components/provider_info_helpers"; -export type WhatsNewLaunch = components["schemas"]["WhatsNewLaunch"]; - -const LAUNCH_ICONS: ReadonlyMap = new Map([ - ["bot", Bot], - ["box", Box], - ["gauge", Gauge], - ["plug", Plug], - ["rocket", Rocket], - ["scale", Scale], - ["shield", Shield], - ["sparkles", Sparkles], - ["zap", Zap], -]); - -export const launchIcon = (name: string): LucideIcon => LAUNCH_ICONS.get(name) ?? Sparkles; +export const DAY_ZERO_TAG = "day 0 support"; export const WHATS_NEW_MAX = 3; -export const latestLaunches = (launches: readonly WhatsNewLaunch[]): readonly WhatsNewLaunch[] => - [...launches].sort((a, b) => b.published_on.localeCompare(a.published_on)).slice(0, WHATS_NEW_MAX); +const CARD_TITLE_TAG = "card_title:"; + +const CARD_SUMMARY_TAG = "card_summary:"; + +export const latestLaunches = (posts: readonly BlogPost[]): readonly BlogPost[] => + posts + .filter((post) => post.tags.includes(DAY_ZERO_TAG)) + .sort((a, b) => b.date.localeCompare(a.date)) + .slice(0, WHATS_NEW_MAX); + +const taggedValue = (tags: readonly string[], prefix: string): string | undefined => { + const value = tags + .find((tag) => tag.startsWith(prefix)) + ?.slice(prefix.length) + .trim(); + return value ? value : undefined; +}; + +export const firstSentence = (text: string): string => /^.*?[.!?](?=\s|$)/.exec(text)?.[0] ?? text; + +export const cardTitle = (post: BlogPost): string => taggedValue(post.tags, CARD_TITLE_TAG) ?? post.title; + +export const cardSummary = (post: BlogPost): string => + taggedValue(post.tags, CARD_SUMMARY_TAG) ?? firstSentence(post.description); + +export const cardProvider = (post: BlogPost): string | undefined => + post.tags.find((tag) => getProviderLogoAndName(tag).logo !== ""); export const formatPublishedOn = (isoDate: string): string => new Date(`${isoDate}T00:00:00`).toLocaleDateString("en-US", { month: "short", day: "numeric" }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/useWhatsNew.ts b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/useWhatsNew.ts deleted file mode 100644 index c2171f3349a..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/useWhatsNew.ts +++ /dev/null @@ -1,10 +0,0 @@ -import { $api } from "@/lib/http/api"; -import { latestLaunches } from "./homeContent"; - -export const useWhatsNew = () => - $api.useQuery( - "get", - "/public/whats_new", - {}, - { staleTime: 60 * 60 * 1000, select: (data) => latestLaunches(data.launches) }, - ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/blogPosts/useBlogPosts.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/blogPosts/useBlogPosts.ts index 81d55e87650..dee3274966f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/blogPosts/useBlogPosts.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/blogPosts/useBlogPosts.ts @@ -6,6 +6,7 @@ export interface BlogPost { description: string; date: string; url: string; + tags: string[]; } export interface BlogPostsResponse { diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index adf020ca96a..f598ad012db 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -13454,9 +13454,9 @@ export interface paths { }; /** * Get Litellm Blog Posts - * @description Public endpoint to get the latest LiteLLM blog posts. + * @description Public endpoint to get the latest LiteLLM blog posts, each with its RSS categories as ``tags``. * - * Fetches from GitHub with a 1-hour in-process cache. + * Fetches the docs RSS feed with a 1-hour in-process cache. * Falls back to the bundled local backup on any failure. */ get: operations["get_litellm_blog_posts_public_litellm_blog_posts_get"]; @@ -13661,31 +13661,6 @@ export interface paths { patch?: never; trace?: never; }; - "/public/whats_new": { - parameters: { - query?: never; - header?: never; - path?: never; - cookie?: never; - }; - /** - * Get Public Whats New - * @description Return the launches the dashboard Home page shows under What's new. - * - * Serves the list bundled with this release. Set ``LITELLM_WHATS_NEW_URL`` to fetch a remote list instead, - * resolved once per process on the first request and falling back to the bundled list on any failure, so an - * airgapped proxy serves the bundled list without retrying. ``LITELLM_LOCAL_WHATS_NEW=True`` forces the - * bundled list even when a URL is set. A restart picks up a newly published list. - */ - get: operations["get_public_whats_new_public_whats_new_get"]; - put?: never; - post?: never; - delete?: never; - options?: never; - head?: never; - patch?: never; - trace?: never; - }; "/queue/chat/completions": { parameters: { query?: never; @@ -27056,6 +27031,11 @@ export interface components { date: string; /** Description */ description: string; + /** + * Tags + * @default [] + */ + tags: string[]; /** Title */ title: string; /** Url */ @@ -49281,30 +49261,6 @@ export interface components { type: "web_search" | "web_search_2025_08_26"; user_location?: components["schemas"]["openai__types__responses__web_search_tool_param__UserLocation"] | null; }; - /** - * WhatsNewLaunch - * @description One launch card in the dashboard Home page's What's new section. - */ - WhatsNewLaunch: { - /** Description */ - description: string; - /** Href */ - href: string; - /** Icon */ - icon: string; - /** - * Published On - * Format: date - */ - published_on: string; - /** Title */ - title: string; - }; - /** WhatsNewResponse */ - WhatsNewResponse: { - /** Launches */ - launches: components["schemas"]["WhatsNewLaunch"][]; - }; /** WorkerRegistryEntry */ WorkerRegistryEntry: { /** Name */ @@ -68725,26 +68681,6 @@ export interface operations { }; }; }; - get_public_whats_new_public_whats_new_get: { - parameters: { - query?: never; - header?: never; - path?: never; - cookie?: never; - }; - requestBody?: never; - responses: { - /** @description Successful Response */ - 200: { - headers: { - [name: string]: unknown; - }; - content: { - "application/json": components["schemas"]["WhatsNewResponse"]; - }; - }; - }; - }; async_queue_request_queue_chat_completions_post: { parameters: { query?: { From 74e5dc9747a692c6ef54efbbab526928e4155c59 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:17:40 -0700 Subject: [PATCH 009/325] fix(proxy): keep provider refusal, logprobs and system_fingerprint in chat completions (#44788) * fix(proxy): keep provider refusal, logprobs and system_fingerprint in chat completions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): handle missing chat completion roles Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(types): restore throwaway Usage skip in ModelResponse Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): handle explicit fields on custom response dumps Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: update refusal expectations for LIT-9235 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(lint): suppress type discipline violations with reasons Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style: ruff format proxy utils Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(types): keep basedpyright counts at merge-base for azure dump and ModelResponse Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(types): keep an explicit null system_fingerprint passed to ModelResponse Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../convert_dict_to_response.py | 10 +- litellm/llms/azure/azure.py | 14 +- .../llms/databricks/chat/transformation.py | 4 + litellm/llms/openai/common_utils.py | 13 +- litellm/llms/openai/openai.py | 9 +- litellm/proxy/utils.py | 40 ++++- litellm/types/llms/databricks.py | 2 + litellm/types/utils.py | 32 ++-- .../providers/test_dashscope_chat_wire.py | 2 +- .../providers/test_openai_chat_wire.py | 1 - .../providers/test_tencent_chat_wire.py | 1 - .../providers/test_xiaomi_mimo_wire.py | 2 - .../test_chat_completions_basic_azure.py | 3 - .../test_chat_completions_basic_openai.py | 3 - .../test_convert_dict_to_response.py | 6 +- tests/unit/llms/openai/test_openai.py | 44 ++++++ .../test_model_dump_with_preserved_fields.py | 149 +++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 18 files changed, 290 insertions(+), 47 deletions(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 7b2860b9f97..6f16e1e3fb4 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -709,7 +709,7 @@ def convert_to_model_response_object( message = Message( content=content, - role=choice["message"]["role"] or "assistant", + role=choice["message"].get("role") or "assistant", function_call=choice["message"].get("function_call", None), tool_calls=tool_calls, audio=choice["message"].get("audio", None), @@ -720,6 +720,8 @@ def convert_to_model_response_object( images=_normalize_images_for_message(choice["message"].get("images", None)), ) finish_reason = choice.get("finish_reason", None) + if "refusal" in choice["message"]: + message.refusal = choice["message"]["refusal"] if finish_reason is None: # gpt-4 vision can return 'finish_reason' or 'finish_details' finish_reason = choice.get("finish_details") or "stop" @@ -728,9 +730,10 @@ def convert_to_model_response_object( ## PROVIDER SPECIFIC FIELDS ## provider_specific_fields = {f: choice[f] for f in choice.keys() - _CHOICES_FIELDS} - logprobs = choice.get("logprobs", None) + provider_sent_logprobs = "logprobs" in choice enhancements = choice.get("enhancements", None) + choice = Choices( finish_reason=finish_reason, index=idx, @@ -739,6 +742,9 @@ def convert_to_model_response_object( enhancements=enhancements, provider_specific_fields=provider_specific_fields, ) + if provider_sent_logprobs and logprobs is None: + choice.logprobs = None + choice_list.append(choice) model_response_object.choices = choice_list diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index d396e1f4867..67b25bf7d12 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -13,6 +13,7 @@ from openai import ( BadRequestError, OpenAI, ) +from pydantic import BaseModel import litellm from litellm.constants import AZURE_OPERATION_POLLING_TIMEOUT, DEFAULT_MAX_RETRIES @@ -41,6 +42,7 @@ from ..base import BaseLLM from ..openai.common_utils import ( build_output_token_limit_response, is_output_token_limit_error, + model_dump_provider_response, ) from .common_utils import ( AzureOpenAIError, @@ -352,12 +354,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): headers, response = self.make_sync_azure_openai_chat_completion_request( azure_client=azure_client, data=data, timeout=timeout ) - if isinstance(response, str): + if not isinstance(response, BaseModel): raise AzureOpenAIError( status_code=500, - message=f"Unexpected string response from Azure: {response[:500]}", + message=f"Unexpected response from Azure: {response!r:.500}", ) - stringified_response: Final = response.model_dump() + stringified_response: Final = model_dump_provider_response(response) ## LOGGING logging_obj.post_call( input=messages, @@ -447,12 +449,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): ) logging_obj.model_call_details["response_headers"] = headers - if isinstance(response, str): + if not isinstance(response, BaseModel): raise AzureOpenAIError( status_code=500, - message=f"Unexpected string response from Azure: {response[:500]}", + message=f"Unexpected response from Azure: {response!r:.500}", ) - stringified_response: Final = response.model_dump() + stringified_response: Final = model_dump_provider_response(response) logging_obj.post_call( input=data["messages"], api_key=api_key, diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index dceec658824..94a74d703df 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -697,6 +697,10 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): model_response.model = f"{_custom_llm_provider}/{_response_model}" model_response.id = completion_response["id"] model_response.created = completion_response["created"] + if "system_fingerprint" in completion_response: + model_response.system_fingerprint = ( # rebind-ok: fills the caller's model_response like id above + completion_response["system_fingerprint"] + ) setattr(model_response, "usage", Usage(**completion_response["usage"])) served_tier: Final = completion_response.get("service_tier") if isinstance(served_tier, str) and served_tier: diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index 839ca04e082..4ab91392ad4 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -12,7 +12,7 @@ import time import uuid from collections.abc import AsyncIterator, Iterator, Mapping from contextlib import suppress -from typing import TYPE_CHECKING, Final, Literal, NamedTuple, Optional +from typing import TYPE_CHECKING, Final, Literal, NamedTuple, Optional, Protocol from urllib.parse import urlsplit import httpx @@ -23,6 +23,7 @@ from openai.types.chat.chat_completion import Choice from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice from openai.types.chat.chat_completion_chunk import ChoiceDelta from openai.types.completion_usage import CompletionUsage +from pydantic import BaseModel if TYPE_CHECKING: from aiohttp import ClientSession @@ -51,6 +52,16 @@ _AZURE_OPENAI_INIT_PARAMS: Final[tuple[str, ...]] = _get_client_init_params(Azur _OPENAI_API_HOST: Final[str] = "api.openai.com" +class _ModelDumpable(Protocol): + def model_dump(self) -> dict[str, object]: ... # mutable-ok: pydantic model_dump returns a dict + + +def model_dump_provider_response(response: _ModelDumpable) -> dict[str, object]: # mutable-ok: callers need a dict + if isinstance(response, BaseModel): + return response.model_dump(exclude_unset=True) + return response.model_dump() + + _OPENAI_HTTPX_DEFAULT_TIMEOUT: Final = CompletionTimeout.normalize(openai.DEFAULT_TIMEOUT) _OPENAI_HTTPX_CONNECTION_LIMITS: Final = httpx.Limits( max_connections=openai.DEFAULT_CONNECTION_LIMITS.max_connections, diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 3a12326d433..58eadeadc1e 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -60,6 +60,7 @@ from .common_utils import ( drop_params_from_unprocessable_entity_error, is_openai_backed_api_base, is_output_token_limit_error, + model_dump_provider_response, ) from .workload_identity import resolve_openai_workload_identity_config @@ -852,7 +853,9 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ) logging_obj.model_call_details["response_headers"] = headers - stringified_response = provider_config.transform_parsed_response_dict(response.model_dump()) + stringified_response = provider_config.transform_parsed_response_dict( + model_dump_provider_response(response) + ) logging_obj.post_call( input=messages, api_key=api_key, @@ -990,7 +993,9 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): timeout=timeout, logging_obj=logging_obj, ) - stringified_response = provider_config.transform_parsed_response_dict(response.model_dump()) + stringified_response = provider_config.transform_parsed_response_dict( + model_dump_provider_response(response) + ) logging_obj.post_call( input=data["messages"], api_key=api_key, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5fa78a44d09..267c87694b9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -111,7 +111,7 @@ except ImportError: raise ImportError("backoff is not installed. Please install it via 'pip install backoff'") from fastapi import HTTPException, status -from pydantic import TypeAdapter, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError import litellm import litellm.litellm_core_utils @@ -9308,12 +9308,15 @@ def validate_model_access( _PRESERVED_NONE_FIELDS: Final[list[tuple[str, str]]] = [ ("message", "content"), # null when tool_calls present (issue #6677) ("message", "role"), # always required by OpenAI spec + ("message", "refusal"), # explicit null from provider ("delta", "content"), # null in streaming chunks ] +_PRESERVED_NONE_CHOICE_FIELDS: Final = ("logprobs",) # explicit null from provider +_PRESERVED_NONE_RESPONSE_FIELDS: Final = ("system_fingerprint",) # explicit null from provider def model_dump_with_preserved_fields( - obj: Any, + obj: BaseModel, preserve_fields: list[str] | None = None, exclude_unset: bool = True, ) -> dict[str, object]: @@ -9321,7 +9324,7 @@ def model_dump_with_preserved_fields( Serialize a Pydantic model to a dictionary while preserving specific fields even if they are None. - Fields listed in _PRESERVED_NONE_FIELDS are restored after + Explicit null values in the preserved field lists are restored after model_dump(exclude_none=True) strips them. Args: @@ -9332,21 +9335,42 @@ def model_dump_with_preserved_fields( Returns: Dictionary representation with None values excluded except for preserved fields """ - result: Final = obj.model_dump(exclude_none=True, exclude_unset=exclude_unset) + result: Final = cast( # cast-ok: Pydantic model_dump returns serialized model fields + dict[str, object], + obj.model_dump(exclude_none=True, exclude_unset=exclude_unset), + ) - choices: Final = result.get("choices") + for field_name in _PRESERVED_NONE_RESPONSE_FIELDS: + if field_name not in result and field_name in getattr(obj, "model_fields_set", ()): + result[field_name] = getattr(obj, field_name) + + choices: Final = cast( # cast-ok: model_dump serializes choices as a list of dicts + Sequence[dict[str, object]] | None, + result.get("choices"), + ) if not choices: return result - obj_choices: Final = obj.choices + obj_choices: Final = cast( # cast-ok: obj is duck-typed, choices is a list on ModelResponse + Sequence[object], + getattr(obj, "choices", ()), + ) for choice_obj, choice_dict in zip(obj_choices, choices): for sub_object, field_name in _PRESERVED_NONE_FIELDS: - sub_dict = choice_dict.get(sub_object) + sub_dict = cast( # cast-ok: message and delta dump as dicts + dict[str, object] | None, choice_dict.get(sub_object) + ) if sub_dict is None: continue if field_name not in sub_dict: - sub_obj = getattr(choice_obj, sub_object, None) + sub_obj = cast( # cast-ok: getattr on a duck-typed choice + object | None, getattr(choice_obj, sub_object, None) + ) if sub_obj is not None and hasattr(sub_obj, field_name): sub_dict[field_name] = getattr(sub_obj, field_name) + for field_name in _PRESERVED_NONE_CHOICE_FIELDS: + if field_name not in choice_dict and hasattr(choice_obj, field_name): + choice_dict[field_name] = getattr(choice_obj, field_name) + return result diff --git a/litellm/types/llms/databricks.py b/litellm/types/llms/databricks.py index a9c027bd2de..c7023e35777 100644 --- a/litellm/types/llms/databricks.py +++ b/litellm/types/llms/databricks.py @@ -2,6 +2,7 @@ from typing import Any, Literal from pydantic import BaseModel from typing_extensions import ( + NotRequired, ReadOnly, Required, TypedDict, @@ -82,3 +83,4 @@ class DatabricksResponse(TypedDict): model: str choices: list[DatabricksChoice] usage: ChatCompletionUsageBlock + system_fingerprint: NotRequired[ReadOnly[str | None]] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 06af506076a..2d98bf0486f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1416,6 +1416,7 @@ def add_provider_specific_fields(object: BaseModel, provider_specific_fields: di class Message(SafeAttributeModel, OpenAIObject): content: str | None role: Literal["assistant", "user", "system", "tool", "function"] + refusal: str | None = None tool_calls: list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] | None function_call: FunctionCall | None audio: ChatCompletionAudioResponse | None = None @@ -1493,6 +1494,9 @@ class Message(SafeAttributeModel, OpenAIObject): if hasattr(self, "annotations"): del self.annotations + if "refusal" not in params and hasattr(self, "refusal"): + del self.refusal + if reasoning_content is None: # ensure default response matches OpenAI spec if hasattr(self, "reasoning_content"): @@ -2212,6 +2216,20 @@ class ModelResponseStream(ModelResponseBase): return self.dict() +def _normalize_model_response_usage(usage: Mapping[str, object] | BaseModel) -> Usage: + if isinstance(usage, dict): + return Usage(**usage) + dump: Final = usage.model_dump() if hasattr(usage, "model_dump") else usage.dict() + return Usage(**dump) + + +class _UnsetSystemFingerprint: + __slots__ = () + + +_UNSET_SYSTEM_FINGERPRINT: Final = _UnsetSystemFingerprint() + + class ModelResponse(ModelResponseBase): choices: list[Choices] """The list of completion choices the model generated for the input prompt.""" @@ -2223,7 +2241,7 @@ class ModelResponse(ModelResponseBase): created=None, model=None, object=None, - system_fingerprint=None, + system_fingerprint=_UNSET_SYSTEM_FINGERPRINT, usage=None, stream=None, stream_options=None, @@ -2259,13 +2277,8 @@ class ModelResponse(ModelResponseBase): created = created model = model if usage is not None: - if isinstance(usage, dict): - usage = Usage(**usage) - elif isinstance(usage, BaseModel): - dump = usage.model_dump() if hasattr(usage, "model_dump") else usage.dict() - usage = Usage(**dump) - else: - usage = usage + if isinstance(usage, dict | BaseModel): + usage = _normalize_model_response_usage(usage) elif stream is None or stream is False: usage = None # avoid constructing throwaway Usage; set by convert_to_model_response_object if hidden_params: @@ -2280,11 +2293,12 @@ class ModelResponse(ModelResponseBase): "created": created, "model": model, "object": object, - "system_fingerprint": system_fingerprint, } if usage is not None: init_values["usage"] = usage + if system_fingerprint is not _UNSET_SYSTEM_FINGERPRINT: + init_values["system_fingerprint"] = system_fingerprint super().__init__( **init_values, diff --git a/tests/integration/providers/test_dashscope_chat_wire.py b/tests/integration/providers/test_dashscope_chat_wire.py index a2b3a36d6e3..87933d1eaf9 100644 --- a/tests/integration/providers/test_dashscope_chat_wire.py +++ b/tests/integration/providers/test_dashscope_chat_wire.py @@ -55,7 +55,7 @@ def test_dashscope_chat_forwards_reasoning_effort_none_to_the_provider(gateway: { "finish_reason": "stop", "index": 0, - "message": {"role": "assistant", "content": "27", "provider_specific_fields": {"refusal": None}}, + "message": {"role": "assistant", "content": "27"}, "provider_specific_fields": {}, } ] diff --git a/tests/integration/providers/test_openai_chat_wire.py b/tests/integration/providers/test_openai_chat_wire.py index 8113a37442a..5bb8a1f2c96 100644 --- a/tests/integration/providers/test_openai_chat_wire.py +++ b/tests/integration/providers/test_openai_chat_wire.py @@ -61,7 +61,6 @@ def test_openai_chat_tool_choice_without_tools_is_not_forwarded(gateway: Gateway "message": { "role": "assistant", "content": "One sentence.", - "provider_specific_fields": {"refusal": None}, }, "provider_specific_fields": {}, } diff --git a/tests/integration/providers/test_tencent_chat_wire.py b/tests/integration/providers/test_tencent_chat_wire.py index 84e9eb8ffea..25e5cac0fb6 100644 --- a/tests/integration/providers/test_tencent_chat_wire.py +++ b/tests/integration/providers/test_tencent_chat_wire.py @@ -77,7 +77,6 @@ def test_tencent_thinking_is_sent_in_provider_body_instead_of_failing_the_reques "role": "assistant", "content": "43", "reasoning_content": "17 plus 26 is 43.", - "provider_specific_fields": {"refusal": None}, }, "provider_specific_fields": {}, } diff --git a/tests/integration/providers/test_xiaomi_mimo_wire.py b/tests/integration/providers/test_xiaomi_mimo_wire.py index 96b9dc17bdf..2722e501c47 100644 --- a/tests/integration/providers/test_xiaomi_mimo_wire.py +++ b/tests/integration/providers/test_xiaomi_mimo_wire.py @@ -136,7 +136,6 @@ def test_xiaomi_mimo_nonstream_surfaces_reasoning_and_charges_registry_price(gat "role": "assistant", "content": "43", "reasoning_content": "17 plus 26 is 43.", - "provider_specific_fields": {"refusal": None}, }, "provider_specific_fields": {}, } @@ -250,7 +249,6 @@ def test_xiaomi_mimo_tool_call_is_forwarded_and_returned(gateway: Gateway) -> No "content": None, "reasoning_content": "Need the tool.", "tool_calls": [tool_call], - "provider_specific_fields": {"refusal": None}, }, "provider_specific_fields": {}, } diff --git a/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_azure.py b/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_azure.py index e19b5223804..b388d303228 100644 --- a/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_azure.py +++ b/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_azure.py @@ -24,7 +24,4 @@ from integration.translation.runner import assert_translation ids=lambda case: case.id, ) def test_chat_completions_basic_azure(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: - pytest.skip( - "BUG: LIT-9235 chat completions moves message.refusal into provider_specific_fields and drops null system_fingerprint and logprobs" - ) assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_openai.py b/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_openai.py index f8d46ee28ac..daa54612208 100644 --- a/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_openai.py +++ b/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_openai.py @@ -24,7 +24,4 @@ from integration.translation.runner import assert_translation ids=lambda case: case.id, ) def test_chat_completions_basic_openai(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: - pytest.skip( - "BUG: LIT-9235 chat completions moves message.refusal into provider_specific_fields and drops null system_fingerprint and logprobs" - ) assert_translation(case, gateway, provider) diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index 44d823fa82f..8fdc6c8ba47 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -1580,10 +1580,8 @@ def test_convert_to_model_response_object_no_provider_specific_fields_still_work ) assert isinstance(result, ModelResponse) - psf = result.choices[0].message.provider_specific_fields - # refusal is not a Message model field, so it should be in provider_specific_fields - assert psf is not None - assert "refusal" in psf + assert result.choices[0].message.refusal is None + assert result.choices[0].message.provider_specific_fields is None def test_convert_to_model_response_object_with_error_code_only(): """ diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index 788269dd043..713b331f8a5 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -38,6 +38,7 @@ from litellm.utils import _invalidate_model_cost_lowercase_map from openai.types.beta.assistant import Assistant from openai.types.beta.assistant_deleted import AssistantDeleted from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome +from litellm.proxy.utils import model_dump_with_preserved_fields @pytest.mark.parametrize( @@ -220,6 +221,49 @@ async def test_acompletion_returns_prompt_cache_details_over_injected_transport( assert prompt_tokens_details.cached_tokens == 7 +@respx.mock +def test_completion_defaults_missing_message_role_and_preserves_omitted_fields() -> None: + route: Final = respx.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-missing-role", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.4", + "choices": [{"index": 0, "message": {"content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + ) + + response: Final = litellm.completion( + model="openai/gpt-5.4", + api_base="https://api.openai.com/v1", + api_key="sk-test", + messages=[{"role": "user", "content": "Hi"}], + num_retries=0, + ) + + assert route.called + assert response.choices[0].message.role == "assistant" + assert model_dump_with_preserved_fields(response, exclude_unset=True) == { + "id": "chatcmpl-missing-role", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": {"content": "hi", "role": "assistant"}, + "provider_specific_fields": {}, + } + ], + "created": 1, + "model": "gpt-5.4", + "object": "chat.completion", + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + + @pytest.mark.asyncio @pytest.mark.parametrize( ("status", "expected_error"), diff --git a/tests/unit/proxy/test_model_dump_with_preserved_fields.py b/tests/unit/proxy/test_model_dump_with_preserved_fields.py index 316c1e879cc..31531dccd16 100644 --- a/tests/unit/proxy/test_model_dump_with_preserved_fields.py +++ b/tests/unit/proxy/test_model_dump_with_preserved_fields.py @@ -2,12 +2,20 @@ Regression tests for model_dump_with_preserved_fields. This function serializes ModelResponse / ModelResponseStream objects to dicts -while preserving 3 specific None fields for OpenAI API compatibility: +while preserving specific fields for OpenAI API compatibility: - choices[*].message.content (null when tool_calls present) - choices[*].message.role (always present) - choices[*].delta.content (null in streaming chunks) + - choices[*].message.refusal (explicit null from provider) + - choices[*].logprobs (explicit null from provider) + - system_fingerprint (explicit null from provider) """ +from typing import Final + +from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_model_response_object, +) from litellm.proxy.utils import model_dump_with_preserved_fields from litellm.types.utils import ( Choices, @@ -118,8 +126,23 @@ def test_none_fields_stripped_from_message(): assert "audio" not in msg -def test_none_fields_stripped_from_top_level(): - """system_fingerprint=None should be omitted from top-level.""" +def test_unset_system_fingerprint_stripped_from_top_level(): + """A system_fingerprint the provider never sent should be omitted from top-level.""" + response = ModelResponse( + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="Hello", role="assistant"), + ) + ], + ) + result = model_dump_with_preserved_fields(response, exclude_unset=True) + assert "system_fingerprint" not in result + + +def test_explicit_none_system_fingerprint_kept_at_top_level(): + """A provider that sent system_fingerprint: null keeps it through ModelResponse(**payload).""" response = ModelResponse( choices=[ Choices( @@ -131,7 +154,7 @@ def test_none_fields_stripped_from_top_level(): system_fingerprint=None, ) result = model_dump_with_preserved_fields(response, exclude_unset=True) - assert "system_fingerprint" not in result + assert result["system_fingerprint"] is None def test_multiple_choices_independent(): @@ -404,3 +427,121 @@ def test_preserve_fields_param_backward_compat(): assert result_default == result_explicit assert result_default["choices"][0]["message"]["content"] is None assert result_default["choices"][0]["message"]["role"] == "assistant" + + +def test_provider_explicit_null_fields_are_preserved() -> None: + response_object: Final = { + "id": "chatcmpl-provider-null-fields", + "created": 1, + "model": "gpt-5.4", + "object": "chat.completion", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": {"role": "assistant", "content": "Hello", "refusal": None}, + "logprobs": None, + } + ], + "system_fingerprint": None, + } + model_response: Final = convert_to_model_response_object( + response_object=response_object, + model_response_object=ModelResponse(), + ) + + assert model_dump_with_preserved_fields(model_response, exclude_unset=True) == { + "id": "chatcmpl-provider-null-fields", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": {"content": "Hello", "role": "assistant", "refusal": None}, + "provider_specific_fields": {}, + "logprobs": None, + } + ], + "created": 1, + "model": "gpt-5.4", + "object": "chat.completion", + "system_fingerprint": None, + } + + +def test_provider_omitted_null_fields_stay_absent() -> None: + response_object: Final = { + "id": "chatcmpl-provider-omitted-fields", + "created": 2, + "model": "gpt-5.4", + "object": "chat.completion", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": {"role": "assistant", "content": "Hello"}, + } + ], + } + model_response: Final = convert_to_model_response_object( + response_object=response_object, + model_response_object=ModelResponse(), + ) + + assert model_dump_with_preserved_fields(model_response, exclude_unset=True) == { + "id": "chatcmpl-provider-omitted-fields", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": {"content": "Hello", "role": "assistant"}, + "provider_specific_fields": {}, + } + ], + "created": 2, + "model": "gpt-5.4", + "object": "chat.completion", + } + + +def test_provider_refusal_string_is_preserved_on_message() -> None: + response_object: Final = { + "id": "chatcmpl-provider-refusal", + "created": 3, + "model": "gpt-5.4", + "object": "chat.completion", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "role": "assistant", + "content": "I can't help with that.", + "refusal": "I can't help with that.", + }, + } + ], + } + model_response: Final = convert_to_model_response_object( + response_object=response_object, + model_response_object=ModelResponse(), + ) + + assert model_dump_with_preserved_fields(model_response, exclude_unset=True) == { + "id": "chatcmpl-provider-refusal", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "I can't help with that.", + "role": "assistant", + "refusal": "I can't help with that.", + }, + "provider_specific_fields": {}, + } + ], + "created": 3, + "model": "gpt-5.4", + "object": "chat.completion", + } + assert model_response.choices[0].message.refusal == "I can't help with that." diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f598ad012db..f44c425c2c0 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -37616,6 +37616,8 @@ export interface components { reasoning_content?: string | null; /** Reasoning Items */ reasoning_items?: components["schemas"]["ChatCompletionReasoningItem"][] | null; + /** Refusal */ + refusal?: string | null; /** * Role * @enum {string} From 93b29c3977173a12cadf72c25fb3835dee18b136 Mon Sep 17 00:00:00 2001 From: mrinal-berri Date: Tue, 13 Oct 2026 16:25:22 -0400 Subject: [PATCH 010/325] fix(proxy): broadcast team cache invalidation on team writes (#46113) * fix(proxy): broadcast team cache invalidation on team writes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover cross-instance team cache invalidation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 17 +- .../test_team_blocking_multi_instance.py | 4 +- .../test_team_update_multi_instance.py | 216 ++++++++++++++++++ .../test_team_endpoints.py | 146 ++++++++++++ 4 files changed, 379 insertions(+), 4 deletions(-) create mode 100644 tests/integration/management/test_team_update_multi_instance.py diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 6a4c8ac2c9e..9cc1afd7beb 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -40,6 +40,7 @@ import litellm from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.caching.redis_batch import active_request_redis_batches from litellm.integrations.prometheus import PrometheusLogger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( @@ -115,7 +116,10 @@ from litellm.proxy.auth.auth_utils import ( enforce_output_token_estimates_are_admin_only, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + evict_and_broadcast, + publish_auth_cache_invalidation, +) from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -473,7 +477,7 @@ async def refresh_cached_team( proxy_logging_obj: ProxyLogging, ) -> None: """ - Refresh the in-memory cached team object after a DB write. + Refresh the cached team object after a DB write and broadcast invalidations to other instances. Every endpoint that mutates `litellm_teamtable` must call this so the cached `LiteLLM_TeamTableCachedObj` used by `common_checks` stays in @@ -486,12 +490,19 @@ async def refresh_cached_team( via `model_dump()` to match the cache shape `_cache_team_object` expects. """ + team_table: Final = LiteLLM_TeamTableCachedObj.model_validate(team_row.model_dump()) await cache_team_object( team_id=team_row.team_id, - team_table=LiteLLM_TeamTableCachedObj.model_validate(team_row.model_dump()), + team_table=team_table, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) + batches: Final = active_request_redis_batches() + if batches is not None: + await batches.flush_all() + await publish_auth_cache_invalidation(cache_key=f"team_id:{team_row.team_id}") + if team_table.team_alias is not None: + await publish_auth_cache_invalidation(cache_key=f"team_alias:{team_table.team_alias}") _refresh_cached_team: Final = refresh_cached_team diff --git a/tests/integration/management/test_team_blocking_multi_instance.py b/tests/integration/management/test_team_blocking_multi_instance.py index 1173fccaeaa..2d66d614d35 100644 --- a/tests/integration/management/test_team_blocking_multi_instance.py +++ b/tests/integration/management/test_team_blocking_multi_instance.py @@ -12,7 +12,7 @@ def _chat(proxy: Gateway, model: str, key: str) -> httpx.Response: def _rejected(proxy: Gateway, model: str, key: str) -> httpx.Response: - return eventually(lambda: _chat(proxy, model, key), lambda response: response.status_code == 401, seconds=70) + return eventually(lambda: _chat(proxy, model, key), lambda response: response.status_code == 401, seconds=10) def _blocked_flag(proxy: Gateway, team_id: str) -> object: @@ -27,6 +27,8 @@ def test_a_team_blocked_on_one_instance_is_rejected_on_both(gateway: Gateway, pe assert _blocked_flag(peer, team) is False allowed: Final = _chat(gateway, model, key) assert allowed.status_code == 200, allowed.text + peer_allowed: Final = _chat(peer, model, key) + assert peer_allowed.status_code == 200, peer_allowed.text peer.post("/team/update", {"team_id": team, "blocked": True}) diff --git a/tests/integration/management/test_team_update_multi_instance.py b/tests/integration/management/test_team_update_multi_instance.py new file mode 100644 index 00000000000..e5b9d31cb6c --- /dev/null +++ b/tests/integration/management/test_team_update_multi_instance.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +import json +from collections.abc import Mapping +from typing import Final +from uuid import uuid4 + +import httpx +import pytest +from pydantic import JsonValue + +from tests.integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value +from tests.integration._support.wire import Reply, Request, wire_server + +_UPSTREAM_REPLIES: Final[Mapping[str, Mapping[str, JsonValue]]] = { + "/v1/chat/completions": { + "id": "chatcmpl_team_update", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.6", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + "/v1/responses": { + "id": "resp_team_update", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_team_update", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "parallel_tool_calls": False, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, +} + + +def _respond(request: Request) -> Reply: + assert request.method == "POST", request.method + assert request.target in _UPSTREAM_REPLIES, request.target + return Reply(body=json.dumps(_UPSTREAM_REPLIES[request.target]).encode()) + + +def _request(proxy: Gateway, route: str, model: str, key: str) -> httpx.Response: + marker: Final = f"team update cache {uuid4().hex}" + if route == "/v1/messages": + return proxy.request( + "POST", + route, + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, + key=key, + headers={"anthropic-version": "2023-06-01"}, + ) + if route == "/v1/responses": + return proxy.request( + "POST", + route, + {"model": model, "input": marker, "max_output_tokens": 16, "store": False}, + key=key, + ) + return proxy.request( + "POST", route, {"model": model, "messages": [{"role": "user", "content": marker}]}, key=key + ) + + +def _chat(proxy: Gateway, model: str, key: str) -> httpx.Response: + return proxy.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"team limit {uuid4().hex}"}]}, + key=key, + ) + + +def _team_limit_error(response: httpx.Response) -> None: + assert response.status_code == 429, response.text + body: Final = JSON_OBJECT.validate_json(response.content) + error: Final = object_value(body["error"]) + message: Final = string_value(error["message"]) + assert "Current limit: 1" in message, response.text + + +def _scenario_model(scenario: Scenario) -> str: + return scenario.model() + + +def _wire_model(scenario: Scenario, upstream_url: str) -> str: + return scenario.model( + model="openai/gpt-5.6", + api_base=upstream_url + "/v1", + api_key="synthetic-openai-key", + ) + + +@pytest.mark.parametrize( + "route", + ( + pytest.param("/v1/chat/completions", id="chat-completions"), + pytest.param("/v1/messages", id="messages"), + pytest.param("/v1/responses", id="responses"), + ), +) +def test_team_update_adds_model_on_peer_for_inference_routes(gateway: Gateway, peer: Gateway, route: str) -> None: + with wire_server(_respond) as upstream, gateway.scenario() as scenario: + model_a: Final = _wire_model(scenario, upstream.url) + model_b: Final = _wire_model(scenario, upstream.url) + team: Final = scenario.team(models=[model_a]) + key: Final = scenario.key(team_id=team) + + peer_a: Final = _request(peer, route, model_a, key) + assert peer_a.status_code == 200, peer_a.text + peer_b_before: Final = _request(peer, route, model_b, key) + assert peer_b_before.status_code == 403, peer_b_before.text + + update: Final = gateway.request("POST", "/team/update", {"team_id": team, "models": [model_a, model_b]}) + assert update.status_code == 200, update.text + + peer_b_after: Final = eventually( + lambda: _request(peer, route, model_b, key), + lambda response: response.status_code == 200, + seconds=10, + return_last_on_timeout=True, + ) + requests: Final = upstream.drain() + assert peer_b_after.status_code == 200, peer_b_after.text + expected_target: Final = "/v1/chat/completions" if route == "/v1/chat/completions" else "/v1/responses" + assert tuple(request.target for request in requests) == (expected_target, expected_target), peer_b_after.text + + +def test_team_update_raising_rpm_limit_is_enforced_by_peer(gateway: Gateway, peer: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = _scenario_model(scenario) + team: Final = scenario.team(models=[model], rpm_limit=1) + key: Final = scenario.key(team_id=team) + + peer_first: Final = _chat(peer, model, key) + assert peer_first.status_code == 200, peer_first.text + peer_limited: Final = _chat(peer, model, key) + _team_limit_error(peer_limited) + + update: Final = gateway.request("POST", "/team/update", {"team_id": team, "rpm_limit": 1000}) + assert update.status_code == 200, update.text + + peer_after: Final = eventually( + lambda: _chat(peer, model, key), + lambda response: response.status_code == 200, + seconds=10, + return_last_on_timeout=True, + ) + assert peer_after.status_code == 200, peer_after.text + + +def test_team_update_lowering_rpm_limit_is_enforced_by_peer(gateway: Gateway, peer: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = _scenario_model(scenario) + team: Final = scenario.team(models=[model], rpm_limit=1000) + key: Final = scenario.key(team_id=team) + + peer_warm: Final = _chat(peer, model, key) + assert peer_warm.status_code == 200, peer_warm.text + + update: Final = gateway.request("POST", "/team/update", {"team_id": team, "rpm_limit": 1}) + assert update.status_code == 200, update.text + + peer_limited: Final = eventually( + lambda: _chat(peer, model, key), + lambda response: response.status_code == 429, + seconds=10, + return_last_on_timeout=True, + ) + _team_limit_error(peer_limited) + + +@pytest.mark.parametrize( + ("route", "initial_models", "expected_before", "expected_after"), + ( + pytest.param("/team/model/add", "one", 403, 200, id="add"), + pytest.param("/team/model/delete", "both", 200, 403, id="delete"), + ), +) +def test_team_model_add_delete_is_enforced_by_peer( + gateway: Gateway, + peer: Gateway, + route: str, + initial_models: str, + expected_before: int, + expected_after: int, +) -> None: + with gateway.scenario() as scenario: + model_a: Final = _scenario_model(scenario) + model_b: Final = _scenario_model(scenario) + team: Final = scenario.team(models=[model_a] if initial_models == "one" else [model_a, model_b]) + key: Final = scenario.key(team_id=team) + + peer_before: Final = _chat(peer, model_b, key) + assert peer_before.status_code == expected_before, peer_before.text + + update: Final = gateway.request("POST", route, {"team_id": team, "models": [model_b]}) + assert update.status_code == 200, update.text + + peer_after: Final = eventually( + lambda: _chat(peer, model_b, key), + lambda response: response.status_code == expected_after, + seconds=10, + return_last_on_timeout=True, + ) + assert peer_after.status_code == expected_after, peer_after.text diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index e921e4df5c1..bbb83ec1248 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -2384,6 +2384,152 @@ async def test_team_model_add_delete_refresh_team_cache(endpoint_name): assert update_call_kwargs.get("include", {}).get("object_permission") is True +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_alias", "expected_keys"), + [ + ("team-alias", ("team_id:team-refresh", "team_alias:team-alias")), + (None, ("team_id:team-refresh",)), + ], +) +async def test_refresh_cached_team_broadcasts_team_and_alias_invalidations( + monkeypatch, team_alias: str | None, expected_keys: tuple[str, ...] +) -> None: + from litellm.proxy.management_endpoints import team_endpoints + + published: Final = MagicMock() + + async def record_cache_team_object(**_kwargs: object) -> None: + return None + + async def record_publish(cache_key: str) -> None: + published(cache_key) + + monkeypatch.setattr(team_endpoints, "cache_team_object", record_cache_team_object) + monkeypatch.setattr(team_endpoints, "publish_auth_cache_invalidation", record_publish) + + await team_endpoints.refresh_cached_team( + team_row=LiteLLM_TeamTable(team_id="team-refresh", team_alias=team_alias), + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert tuple(call.args[0] for call in published.call_args_list) == expected_keys + + +@pytest.mark.asyncio +async def test_refresh_cached_team_flushes_redis_set_before_broadcast(monkeypatch) -> None: + from datetime import timedelta + from types import SimpleNamespace + from unittest.mock import AsyncMock + + from litellm.caching.in_memory_cache import InMemoryCache + from litellm.caching.redis_batch import request_redis_batch_scope + from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints import team_endpoints + + events: Final = MagicMock() + + class RecordingPipeline: + def __init__(self, calls: MagicMock) -> None: + self.calls: Final = MagicMock() + self.events: Final = calls + + async def __aenter__(self) -> "RecordingPipeline": + return self + + async def __aexit__(self, *exc: object) -> None: + return None + + def set(self, key: str, _value: str, ex: timedelta | None = None) -> None: + self.calls("SET", key) + + def delete(self, *keys: str) -> None: + for key in keys: + self.calls("DEL", key) + + async def execute(self, raise_on_error: bool = True) -> list[object]: + assert raise_on_error is False + for command in self.calls.call_args_list: + self.events(*command.args) + return [True for _ in self.calls.call_args_list] + + class RecordingRedisClient: + def __init__(self, calls: MagicMock) -> None: + self.events: Final = calls + + def pipeline(self, transaction: bool = True) -> RecordingPipeline: + assert transaction is False + return RecordingPipeline(self.events) + + class RecordingRedisCache(RedisCache): + def __init__(self) -> None: + self.redis_kwargs: Final[dict[str, object]] = {} + self.namespace: Final[str | None] = None + self.client: Final = RecordingRedisClient(events) + self._circuit_breaker: Final = RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30) + self.service_logger_obj: Final = SimpleNamespace( + async_service_success_hook=AsyncMock(), async_service_failure_hook=AsyncMock() + ) + self.default_ttl: Final[int | None] = None + + def init_async_client( + self, + ) -> RecordingRedisClient: # pyright: ignore[reportIncompatibleMethodOverride] # in-memory Redis client + return self.client + + async def record_publish(cache_key: str) -> None: + events("PUBLISH", cache_key) + + redis_cache: Final = RecordingRedisCache() + cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.internal_usage_cache.dual_cache = None + monkeypatch.setattr(team_endpoints, "publish_auth_cache_invalidation", record_publish) + + with request_redis_batch_scope(): + await team_endpoints.refresh_cached_team( + team_row=LiteLLM_TeamTable(team_id="team-order", team_alias=None), + user_api_key_cache=cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert [(call.args[0], call.args[1]) for call in events.call_args_list] == [ + ("SET", "team_id:team-order"), + ("PUBLISH", "team_id:team-order"), + ] + + +@pytest.mark.asyncio +async def test_get_team_object_cache_miss_does_not_broadcast_invalidation(monkeypatch) -> None: + from unittest.mock import AsyncMock + + from litellm.caching.in_memory_cache import InMemoryCache + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + team_id: Final = "team-read-cache-miss" + prisma_client: Final = MagicMock() + prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=LiteLLM_TeamTable(team_id=team_id, models=["model-a"]) + ) + cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache()) + publish: Final = AsyncMock() + monkeypatch.setattr(auth_checks, "publish_auth_cache_invalidation", publish) + + team: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=cache, + ) + + assert team.team_id == team_id + assert await cache.async_get_cache(key=f"team_id:{team_id}", model_type=LiteLLM_TeamTableCachedObj) is not None + publish.assert_not_awaited() + + @pytest.mark.asyncio @pytest.mark.parametrize("endpoint_name", ["team_model_add", "team_model_delete"]) async def test_team_model_add_delete_keep_model_aliases_in_team_cache(endpoint_name, monkeypatch): From eacc143dca815fed52d0aa7b08130d1e485a1479 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:25:31 -0700 Subject: [PATCH 011/325] feat(pricing): add github_copilot/gemini-4-argon (#46280) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- ...model_prices_and_context_window_backup.json | 18 ++++++++++++++++++ model_prices_and_context_window.json | 18 ++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0b96fd3e2b0..dcbb31863c5 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -73045,5 +73045,23 @@ "supports_vision": true, "supports_web_search": true, "web_search_billing_unit": "per_query" + }, + "github_copilot/gemini-4-argon": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "github_copilot", + "max_input_tokens": 200000, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_vision": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0b96fd3e2b0..dcbb31863c5 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -73045,5 +73045,23 @@ "supports_vision": true, "supports_web_search": true, "web_search_billing_unit": "per_query" + }, + "github_copilot/gemini-4-argon": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "github_copilot", + "max_input_tokens": 200000, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_vision": true } } From 581449b4d883080ed4bcaf4c26061e31116f985a Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 13 Oct 2026 13:26:19 -0700 Subject: [PATCH 012/325] fix(proxy): warn when --max_requests_before_restart runs with a single worker (#46185) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_cli.py | 39 ++++- tests/unit/proxy/test_proxy_cli.py | 252 +++++++++++++++++++++++++++++ 2 files changed, 288 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index c7371f294f0..7764d877b98 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -368,6 +368,26 @@ class ProxyInitializationHelpers: _apply_uvicorn_max_requests_jitter = apply_uvicorn_max_requests_jitter + @staticmethod + def single_worker_max_requests_warning( + num_workers: int, + max_requests_before_restart: int | None, + run_gunicorn: bool, + ) -> str | None: + if max_requests_before_restart is None or num_workers > 1: + return None + outage: Final = ( + "gunicorn accepts no connections until the replacement worker has booted" + if run_gunicorn + else "uvicorn exits the process, so the container restarts" + ) + return ( + f"LiteLLM Proxy: --max_requests_before_restart {max_requests_before_restart} with --num_workers " + f"{num_workers} does not recycle a worker. After {max_requests_before_restart} requests, health " + f"probes included, {outage}. Use --num_workers 2 or more with --max_requests_before_restart_jitter, " + "or rely on pod replacement instead." + ) + @staticmethod def _get_reload_options(config_path: str | None) -> dict: """Build uvicorn reload kwargs so --reload also reacts to .env and YAML edits.""" @@ -968,7 +988,12 @@ class ProxyInitializationHelpers: "--max_requests_before_restart", default=None, type=int, - help="Restart worker after this many requests (uvicorn: limit_max_requests, gunicorn: max_requests)", + help=( + "Recycle each worker after it has served this many HTTP requests, health probes included " + "(uvicorn: limit_max_requests, gunicorn: max_requests). Only safe with --num_workers 2 or more. " + "With a single worker the whole proxy stops serving when the limit is hit, uvicorn exits the " + "process and gunicorn accepts no connections until the replacement worker has booted." + ), envvar="MAX_REQUESTS_BEFORE_RESTART", ) @click.option( @@ -1070,7 +1095,7 @@ def run_server( max_budget, test, local, - num_workers, + num_workers: int, granian_threads, test_async, iam_token_db_auth, @@ -1092,7 +1117,7 @@ def run_server( validate_config: bool, keepalive_timeout, timeout_worker_healthcheck, - max_requests_before_restart, + max_requests_before_restart: int | None, max_requests_before_restart_jitter: int | None, limit_concurrency: int | None, enforce_prisma_migration_check: bool, @@ -1571,6 +1596,14 @@ def run_server( # Optional: recycle uvicorn workers after N requests if max_requests_before_restart is not None: uvicorn_args["limit_max_requests"] = max_requests_before_restart + if run_hypercorn is False and run_granian is False: + single_worker_warning: Final = ProxyInitializationHelpers.single_worker_max_requests_warning( + num_workers=num_workers, + max_requests_before_restart=max_requests_before_restart, + run_gunicorn=run_gunicorn is True, + ) + if single_worker_warning is not None: + print(f"\033[1;33m{single_worker_warning}\033[0m\n") if run_gunicorn is False and run_hypercorn is False and run_granian is False: if limit_concurrency is not None: uvicorn_args["limit_concurrency"] = limit_concurrency diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index c73b60f3c26..33f0216ff85 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -1768,6 +1768,258 @@ class TestProxyInitializationHelpers: assert g_kwargs["max_requests_before_restart"] == 900 assert g_kwargs["max_requests_before_restart_jitter"] == 75 + @pytest.mark.parametrize("run_gunicorn", [False, True]) + @pytest.mark.parametrize( + "num_workers,max_requests_before_restart", + [(1, None), (2, 5)], + ) + def test_single_worker_max_requests_warning_returns_none( + self, num_workers, max_requests_before_restart, run_gunicorn + ): + """no warning when the flag is unset or more than one worker serves""" + warning = ProxyInitializationHelpers.single_worker_max_requests_warning( + num_workers=num_workers, + max_requests_before_restart=max_requests_before_restart, + run_gunicorn=run_gunicorn, + ) + assert warning is None + + def test_single_worker_max_requests_warning_uvicorn(self): + warning = ProxyInitializationHelpers.single_worker_max_requests_warning( + num_workers=1, + max_requests_before_restart=5, + run_gunicorn=False, + ) + assert warning is not None + assert "--max_requests_before_restart 5" in warning + assert "--num_workers 1" in warning + assert "uvicorn exits the process" in warning + + def test_single_worker_max_requests_warning_gunicorn(self): + warning = ProxyInitializationHelpers.single_worker_max_requests_warning( + num_workers=1, + max_requests_before_restart=5, + run_gunicorn=True, + ) + assert warning is not None + assert "gunicorn accepts no connections" in warning + assert "uvicorn exits" not in warning + + @patch("uvicorn.run") + @patch("builtins.print") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + def test_single_worker_max_requests_warning_printed_on_uvicorn( + self, mock_setup_db, mock_print, mock_uvicorn_run + ): + """uvicorn single-worker startup prints the does-not-recycle warning""" + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + }, + ), + patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" + ) as mock_get_args, + ): + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, ["--local", "--max_requests_before_restart", "123"] + ) + + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + assert any( + "--max_requests_before_restart 123 with --num_workers 1" in str(c) + and "uvicorn exits the process" in str(c) + for c in mock_print.call_args_list + ) + + @patch("litellm.proxy.proxy_cli.start_query_engine_reaper") + @patch("uvicorn.run") + @patch("builtins.print") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + def test_single_worker_max_requests_no_warning_with_two_workers( + self, mock_setup_db, mock_print, mock_uvicorn_run, mock_start_reaper + ): + """num_workers 2 with --max_requests_before_restart prints no warning""" + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + }, + ), + patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" + ) as mock_get_args, + ): + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, + ["--local", "--num_workers", "2", "--max_requests_before_restart", "123"], + ) + + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + assert not any( + "does not recycle a worker" in str(c) for c in mock_print.call_args_list + ) + + @patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.run_gunicorn_server") + @patch("uvicorn.run") + @patch("builtins.print") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + def test_single_worker_max_requests_warning_printed_on_gunicorn( + self, mock_setup_db, mock_print, mock_uvicorn_run, mock_run_gunicorn + ): + """--run_gunicorn single-worker startup prints the gunicorn warning variant""" + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + }, + ), + patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" + ) as mock_get_args, + ): + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, + ["--local", "--run_gunicorn", "--max_requests_before_restart", "123"], + ) + + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + assert any( + "--max_requests_before_restart 123 with --num_workers 1" in str(c) + and "gunicorn accepts no connections" in str(c) + for c in mock_print.call_args_list + ) + + @patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.init_hypercorn_server") + @patch("uvicorn.run") + @patch("builtins.print") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + def test_single_worker_max_requests_no_warning_on_hypercorn( + self, mock_setup_db, mock_print, mock_uvicorn_run, mock_init_hypercorn + ): + """hypercorn ignores --max_requests_before_restart, so no warning prints""" + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + }, + ), + patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" + ) as mock_get_args, + ): + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, + ["--local", "--run_hypercorn", "--max_requests_before_restart", "5"], + ) + + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + mock_init_hypercorn.assert_called_once() + assert not any( + "does not recycle a worker" in str(c) for c in mock_print.call_args_list + ) + @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") def test_gunicorn_options_include_max_requests_jitter(self): """_run_gunicorn_server puts max_requests_jitter into the gunicorn options""" From 426ab45f2dac1db821f98c6b53aadfa0cd53dbdc Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:31:37 -0700 Subject: [PATCH 013/325] fix(cost-map): set together_ai Qwen3.7-Max to published price and restore 2026-09-15 removal dates (#46281) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 12 ++++++------ model_prices_and_context_window.json | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index dcbb31863c5..b3002fafc45 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -45710,14 +45710,14 @@ "supports_reasoning": true }, "together_ai/Qwen/Qwen3.7-Max": { - "cache_read_input_token_cost": 5e-07, - "input_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 1.5e-06, "litellm_provider": "together_ai", "max_input_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 7.5e-06, - "source": "https://api.together.ai/v1/models", + "output_cost_per_token": 4.5e-06, + "source": "https://www.together.ai/pricing", "supports_prompt_caching": true }, "together_ai/Qwen/Qwen3.7-Plus": { @@ -66697,7 +66697,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/google/gemma-4-31B-it": { - "deprecation_date": "2026-09-14", + "deprecation_date": "2026-09-15", "input_cost_per_token": 3.9e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, @@ -66840,7 +66840,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/openai/gpt-oss-20b": { - "deprecation_date": "2026-09-14", + "deprecation_date": "2026-09-15", "input_cost_per_token": 5e-08, "litellm_provider": "together_ai", "max_input_tokens": 131072, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index dcbb31863c5..b3002fafc45 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -45710,14 +45710,14 @@ "supports_reasoning": true }, "together_ai/Qwen/Qwen3.7-Max": { - "cache_read_input_token_cost": 5e-07, - "input_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 1.5e-06, "litellm_provider": "together_ai", "max_input_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 7.5e-06, - "source": "https://api.together.ai/v1/models", + "output_cost_per_token": 4.5e-06, + "source": "https://www.together.ai/pricing", "supports_prompt_caching": true }, "together_ai/Qwen/Qwen3.7-Plus": { @@ -66697,7 +66697,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/google/gemma-4-31B-it": { - "deprecation_date": "2026-09-14", + "deprecation_date": "2026-09-15", "input_cost_per_token": 3.9e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, @@ -66840,7 +66840,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/openai/gpt-oss-20b": { - "deprecation_date": "2026-09-14", + "deprecation_date": "2026-09-15", "input_cost_per_token": 5e-08, "litellm_provider": "together_ai", "max_input_tokens": 131072, From d9123f459561754f41fd598a73b405b091d3bf84 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:59:27 -0700 Subject: [PATCH 014/325] feat(realtime): proxy OpenAI Live sessions on /v1/live/sessions (#43621) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- gateway/routes/allowlist.py | 2 + litellm/constants.py | 5 + litellm/cost_calculator.py | 214 ++- litellm/litellm_core_utils/core_helpers.py | 4 + litellm/litellm_core_utils/logging_worker.py | 2 +- litellm/llms/openai/live/__init__.py | 3 + litellm/llms/openai/live/handler.py | 555 ++++++++ litellm/llms/openai/realtime/handler.py | 4 +- litellm/proxy/_types.py | 3 + .../auth/managed_authorization.py | 5 +- litellm/proxy/auth/auth_checks.py | 55 +- litellm/proxy/auth/user_api_key_auth.py | 20 +- litellm/proxy/proxy_server.py | 213 ++- litellm/realtime_api/main.py | 103 +- litellm/types/utils.py | 3 + tests/unit/llms/openai/live/__init__.py | 0 tests/unit/llms/openai/live/test_handler.py | 1208 +++++++++++++++++ .../auth/test_managed_authorization.py | 14 + .../unit/proxy/auth/test_user_api_key_auth.py | 49 + tests/unit/proxy/test_proxy_server.py | 499 ++++++- tests/unit/realtime_api/test_main.py | 191 +++ tests/unit/test_cost_calculator.py | 261 ++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 114 ++ 23 files changed, 3463 insertions(+), 64 deletions(-) create mode 100644 litellm/llms/openai/live/__init__.py create mode 100644 litellm/llms/openai/live/handler.py create mode 100644 tests/unit/llms/openai/live/__init__.py create mode 100644 tests/unit/llms/openai/live/test_handler.py diff --git a/gateway/routes/allowlist.py b/gateway/routes/allowlist.py index 3cfe9569b34..cced49f788a 100644 --- a/gateway/routes/allowlist.py +++ b/gateway/routes/allowlist.py @@ -121,6 +121,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = ( # Realtime / streaming "/v1/realtime", "/realtime", + "/v1/live/sessions", + "/live/sessions", # Health & ops "/health", "/metrics", diff --git a/litellm/constants.py b/litellm/constants.py index 04ed307330c..6a6816eb4fa 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -391,6 +391,11 @@ MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY: Final = "litellm.model_group_alias_resolve MCP_PEEKED_BODY_SCOPE_KEY: Final = "litellm_mcp_peeked_body" MCP_ADMISSION_BODY_PEEK_TIMEOUT_SECONDS: Final = 5.0 REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" +OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS: Final = 30 +OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS: Final = 5 +OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES: Final = frozenset( + {"response.completed", "response.incomplete", "response.failed"} +) REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged" # SSL/TLS cipher configuration for faster handshakes diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 263d0bfcb81..96d687b44fb 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1,14 +1,16 @@ # What is this? ## File for 'response_cost' calculation in Logging import logging +import math import time from collections.abc import Mapping, Sequence +from datetime import datetime from functools import lru_cache from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast from httpx import Response -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, ValidationError from typing_extensions import ReadOnly, TypedDict, assert_never import litellm @@ -17,7 +19,9 @@ from litellm import verbose_logger from litellm.constants import ( DEFAULT_MAX_LRU_CACHE_SIZE, DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND, + OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.cost_discount import resolve_cost_discount from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, is_decisions_call_type from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR @@ -110,7 +114,6 @@ from litellm.types.llms.openai import ( OpenAIModerationResponse, OpenAIRealtimeStreamList, OpenAIRealtimeStreamResponseBaseObject, - OpenAIRealtimeStreamSessionEvents, ResponseAPIUsage, ResponsesAPIResponse, ) @@ -2979,6 +2982,141 @@ class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor): _TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed" +class _LiveResponsePayload(BaseModel): + model_config = ConfigDict(frozen=True) + + id: str | None = None + model: str | None = None + usage: Mapping[str, object] | None = None + + +class _LiveResponseEvent(BaseModel): + model_config = ConfigDict(frozen=True) + + type: str = "" + response: _LiveResponsePayload | None = None + + +class _LiveResponseEnvelope(BaseModel): + model_config = ConfigDict(frozen=True) + + type: str = "" + event: _LiveResponseEvent | None = None + + +def _live_response_payload(result: Mapping[str, object]) -> _LiveResponsePayload | None: + try: + envelope: Final = _LiveResponseEnvelope.model_validate(result) + except ValidationError: + return None + if ( + envelope.type != "response.event" + or envelope.event is None + or envelope.event.type not in OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES + or envelope.event.response is None + or envelope.event.response.usage is None + ): + return None + return envelope.event.response + + +def _unique_live_response_payloads(results: Sequence[Mapping[str, object]]) -> tuple[_LiveResponsePayload, ...]: + response_payloads: Final = tuple( + response for result in results if (response := _live_response_payload(result)) is not None + ) + first_payload_by_id: Final = { + response.id: response for response in reversed(response_payloads) if response.id is not None + } + return tuple(reversed(first_payload_by_id.values())) + + +def _live_session_model(result: Mapping[str, object]) -> str | None: + session: Final = as_str_mapping(result.get("session")) + if session is None: + return None + model: Final = session.get("model") + return model if isinstance(model, str) else None + + +def _live_usage_seconds(result: Mapping[str, object]) -> float | None: + usage: Final = as_str_mapping(result.get("usage")) + if usage is None: + return None + seconds: Final = usage.get("seconds") + if isinstance(seconds, (int, float)) and not isinstance(seconds, bool) and math.isfinite(seconds): + return float(seconds) + return None + + +def _last_live_usage_seconds(results: Sequence[Mapping[str, object]], event_type: str) -> float | None: + return next( + ( + seconds + for result in reversed(results) + if result.get("type") == event_type and (seconds := _live_usage_seconds(result)) is not None + ), + None, + ) + + +def _live_session_wall_clock_seconds(litellm_logging_obj: LitellmLoggingObject | None) -> float: + if litellm_logging_obj is None: + return 0.0 + model_call_details: Final = litellm_logging_obj.model_call_details + start_time: Final = model_call_details.get("start_time") + end_time: Final = model_call_details.get("end_time") + if isinstance(start_time, datetime) and isinstance(end_time, datetime): + return max((end_time - start_time).total_seconds(), 0.0) + return 0.0 + + +def _live_input_cost_per_second(model_name: str, custom_llm_provider: str) -> float | None: + model_info: Final = _lookup_model_info_or_none(model_name, custom_llm_provider) + if model_info is None: + return None + rate: Final = model_info.get("input_cost_per_second") + if isinstance(rate, (int, float)) and not isinstance(rate, bool) and math.isfinite(rate): + return float(rate) + return None + + +def _first_live_input_cost_per_second( + potential_model_names: Sequence[str | None], + custom_llm_provider: str, +) -> float: + return next( + ( + rate + for model_name in potential_model_names + if model_name is not None + and (rate := _live_input_cost_per_second(model_name, custom_llm_provider)) is not None + ), + 0.0, + ) + + +def _live_response_token_costs( + response: _LiveResponsePayload, + custom_llm_provider: str, + data_residency: str | None, +) -> tuple[float, float]: + if response.id is None or response.model is None or response.usage is None: + return 0.0, 0.0 + try: + usage: Final = get_usage_object(response.model_dump()) + except (ValidationError, TypeError, ValueError, KeyError, AttributeError): + return 0.0, 0.0 + if usage is None: + return 0.0, 0.0 + costs: Final = _candidate_realtime_token_costs( + model_name=response.model, + combined_usage_object=usage, + custom_llm_provider=custom_llm_provider, + data_residency=data_residency, + ) + return (0.0, 0.0) if costs is None else costs + + def _candidate_realtime_token_costs( model_name: str, combined_usage_object: Usage, @@ -3065,14 +3203,21 @@ def handle_realtime_stream_cost_calculation( base_pricing_model: the deployment's resolved base_model, tried ahead of the session-reported model but after custom rates """ - received_model = None - potential_model_names: Final = [custom_pricing_model, base_pricing_model] - for result in results: - if result["type"] == "session.created": - received_model = cast(OpenAIRealtimeStreamSessionEvents, result)["session"].get("model", None) - potential_model_names.append(received_model) - - potential_model_names.append(litellm_model_name) + received_model: Final = next( + ( + model_name + for result in results + if result["type"] in ("session.created", "session.started") + and (model_name := _live_session_model(result)) is not None + ), + None, + ) + potential_model_names: Final = ( + custom_pricing_model, + base_pricing_model, + received_model, + litellm_model_name, + ) input_cost_per_token, output_cost_per_token = _first_priced_realtime_token_costs( potential_model_names=potential_model_names, combined_usage_object=combined_usage_object, @@ -3089,15 +3234,56 @@ def handle_realtime_stream_cost_calculation( if any(r.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE for r in results) else 0.0 ) - total_cost: Final = input_cost_per_token + output_cost_per_token + transcription_cost + live_lifecycle_event_types: Final = frozenset({"session.started", "session.usage.updated", "session.closed"}) + live_session_started: Final = any(result["type"] == "session.started" for result in results) + live_session_closed: Final = any(result["type"] == "session.closed" for result in results) + live_event_seen: Final = any(result["type"] in live_lifecycle_event_types for result in results) + closed_seconds: Final = _last_live_usage_seconds(results, "session.closed") + reported_seconds: Final = _last_live_usage_seconds(results, "session.usage.updated") + latest_reported_seconds: Final = 0.0 if reported_seconds is None else reported_seconds + voice_seconds: Final = ( + closed_seconds + if closed_seconds is not None + else max( + latest_reported_seconds, + _live_session_wall_clock_seconds(litellm_logging_obj), + ) + if live_session_started and not live_session_closed + else latest_reported_seconds + ) + live_model_names: Final = ( + (custom_pricing_model, base_pricing_model, received_model, litellm_model_name) if live_event_seen else () + ) + live_voice_cost: Final = voice_seconds * _first_live_input_cost_per_second( + live_model_names, + custom_llm_provider, + ) + live_response_costs: Final = tuple( + _live_response_token_costs( + response=response, + custom_llm_provider=custom_llm_provider, + data_residency=data_residency, + ) + for response in _unique_live_response_payloads(results) + ) + live_response_input_cost: Final = sum(costs[0] for costs in live_response_costs) + live_response_output_cost: Final = sum(costs[1] for costs in live_response_costs) + input_cost: Final = input_cost_per_token + live_response_input_cost + output_cost: Final = output_cost_per_token + live_response_output_cost + total_cost: Final = input_cost + output_cost + transcription_cost + live_voice_cost + additional_costs: Final = tuple( + (cost_name, cost) + for cost_name, cost in (("transcription_cost", transcription_cost), ("live_voice_cost", live_voice_cost)) + if cost > 0 + ) _store_cost_breakdown_in_logging_obj( litellm_logging_obj=litellm_logging_obj, - prompt_tokens_cost_usd_dollar=input_cost_per_token, - completion_tokens_cost_usd_dollar=output_cost_per_token, + prompt_tokens_cost_usd_dollar=input_cost, + completion_tokens_cost_usd_dollar=output_cost, cost_for_built_in_tools_cost_usd_dollar=0.0, total_cost_usd_dollar=total_cost, - additional_costs={"transcription_cost": transcription_cost} if transcription_cost > 0 else None, + additional_costs=dict(additional_costs) if additional_costs else None, data_residency=data_residency, ) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index e857461785c..c12eb4d458f 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -52,6 +52,10 @@ def safe_divide_seconds(seconds: float, denominator: float, default: float | Non _DROP_PARAMS_BOOL: Final = TypeAdapter(bool) +def as_str_mapping(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, Mapping) else None + + def normalize_drop_params(value: object) -> bool | None: if value is None or isinstance(value, bool): return value diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index 57da3f8dabe..ff649953e69 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -466,7 +466,7 @@ class LoggingWorker: # Process all tasks concurrently for maximum speed await asyncio.gather(*[self._process_single_task(task) for task in tasks]) - def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine): + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, None, None]) -> None: """ Ensure the logging worker is initialized and enqueue the coroutine. """ diff --git a/litellm/llms/openai/live/__init__.py b/litellm/llms/openai/live/__init__.py new file mode 100644 index 00000000000..03d40156213 --- /dev/null +++ b/litellm/llms/openai/live/__init__.py @@ -0,0 +1,3 @@ +from .handler import LiveSessionStart, OpenAILiveSessions, parse_live_session_start + +__all__ = ("LiveSessionStart", "OpenAILiveSessions", "parse_live_session_start") diff --git a/litellm/llms/openai/live/handler.py b/litellm/llms/openai/live/handler.py new file mode 100644 index 00000000000..297a147a288 --- /dev/null +++ b/litellm/llms/openai/live/handler.py @@ -0,0 +1,555 @@ +import asyncio +import json +import math +import ssl as ssl_module +from collections.abc import AsyncIterator, Callable, Mapping +from contextlib import AbstractAsyncContextManager, suppress +from dataclasses import dataclass +from datetime import datetime +from itertools import count +from types import MappingProxyType +from typing import Final, Protocol, TypeAlias +from urllib.parse import urlsplit, urlunsplit + +from pydantic import TypeAdapter, ValidationError +from typing_extensions import NotRequired, ReadOnly, TypedDict + +from litellm._logging import redact_secrets +from litellm.constants import ( + OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS, + OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES, + REALTIME_SESSION_SUCCESS_LOGGED_KEY, + REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, +) +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker +from litellm.litellm_core_utils.realtime_errors import ( + close_after_upstream_handshake_refusal, + realtime_error_event, + websocket_close_reason, +) +from litellm.litellm_core_utils.realtime_streaming import backend_close_from +from litellm.llms.openai.realtime.handler import OpenAIRealtime +from litellm.types.utils import LiteLLMRealtimeStreamLoggingObject, Usage + +_LIVE_EVENT_FRAME_ADAPTER: Final = TypeAdapter(dict[str, object]) +_LIVE_USAGE_EVENT_TYPES: Final = frozenset({"session.usage.updated", "session.closed"}) +_LIVE_WEBSOCKET_SCHEMES: Final = MappingProxyType({"https": "wss", "http": "ws"}) +_LIVE_SESSION_CLOSE_FRAME: Final = '{"type": "session.close"}' + + +class _LiveSessionModel(TypedDict): + model: ReadOnly[str] + + +class _LiveSessionStartedEvent(TypedDict): + type: ReadOnly[str] + session: NotRequired[ReadOnly[_LiveSessionModel]] + + +class _LiveUsageSeconds(TypedDict): + seconds: ReadOnly[object] + + +class _LiveUsageEvent(TypedDict): + type: ReadOnly[str] + usage: NotRequired[ReadOnly[_LiveUsageSeconds]] + + +class _LiveResponseSummary(TypedDict): + id: ReadOnly[str] + model: ReadOnly[object] + usage: ReadOnly[object] + + +class _LiveTerminalResponseEvent(TypedDict): + type: ReadOnly[str] + response: ReadOnly[_LiveResponseSummary] + + +class _LiveResponseEvent(TypedDict): + type: ReadOnly[str] + event: ReadOnly[_LiveTerminalResponseEvent] + + +class _LivePreCallInput(TypedDict): + session_start: ReadOnly[Mapping[str, object]] + + +class _LivePreCallArgs(TypedDict): + api_base: ReadOnly[str] + headers: ReadOnly[Mapping[str, str]] + complete_input_dict: ReadOnly[_LivePreCallInput] + + +class LiveClientWebSocket(Protocol): + async def receive_text(self) -> str: ... + + async def send_text(self, data: str) -> None: ... + + async def close(self, code: int = 1000, reason: str | None = None) -> None: ... + + +class LiveBackendWebSocket(Protocol): + async def send(self, message: str, /) -> None: ... + + async def recv(self) -> str | bytes: ... + + async def close(self) -> None: ... + + +class LiveWebSocketConnector(Protocol): + def __call__( + self, + url: str, + *, + additional_headers: Mapping[str, str], + max_size: int | None, + ssl: bool | str | ssl_module.SSLContext | None, + ) -> AbstractAsyncContextManager[LiveBackendWebSocket]: ... + + +def _openai_live_websocket_connect( + url: str, + *, + additional_headers: Mapping[str, str], + max_size: int | None, + ssl: bool | str | ssl_module.SSLContext | None, +) -> AbstractAsyncContextManager[LiveBackendWebSocket]: + import websockets + + return websockets.connect( + url, + additional_headers=additional_headers, + max_size=max_size, + ssl=ssl, + ) + + +def _is_finite_usage_seconds(value: object) -> bool: + return ( + isinstance(value, (int, float)) + and not isinstance(value, bool) + and (not isinstance(value, float) or math.isfinite(value)) + ) + + +def _session_started_event(event: Mapping[str, object]) -> _LiveSessionStartedEvent: + session: Final = as_str_mapping(event.get("session")) + session_model: Final = session.get("model") if session is not None else None + if isinstance(session_model, str): + started_with_model: Final[_LiveSessionStartedEvent] = { + "type": "session.started", + "session": {"model": session_model}, + } + return started_with_model + started: Final[_LiveSessionStartedEvent] = {"type": "session.started"} + return started + + +def _usage_event(event: Mapping[str, object], event_type: str) -> _LiveUsageEvent: + usage: Final = as_str_mapping(event.get("usage")) + seconds: Final = usage.get("seconds") if usage is not None else None + if _is_finite_usage_seconds(seconds): + usage_with_seconds: Final[_LiveUsageEvent] = {"type": event_type, "usage": {"seconds": seconds}} + return usage_with_seconds + usage_without_seconds: Final[_LiveUsageEvent] = {"type": event_type} + return usage_without_seconds + + +def _terminal_response_event(event: Mapping[str, object]) -> _LiveResponseEvent | None: + nested_event: Final = as_str_mapping(event.get("event")) + if nested_event is None: + return None + nested_event_type: Final = nested_event.get("type") + if not isinstance(nested_event_type, str) or nested_event_type not in OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES: + return None + response: Final = as_str_mapping(nested_event.get("response")) + response_id: Final = response.get("id") if response is not None else None + if not isinstance(response_id, str) or response is None: + return None + response_event: Final[_LiveResponseEvent] = { + "type": "response.event", + "event": { + "type": nested_event_type, + "response": { + "id": response_id, + "model": response.get("model"), + "usage": response.get("usage"), + }, + }, + } + return response_event + + +def _metering_event(message: str) -> Mapping[str, object] | None: + try: + event: Final = _LIVE_EVENT_FRAME_ADAPTER.validate_json(message) + except ValidationError: + return None + event_type: Final = event.get("type") + if event_type == "session.started": + return _session_started_event(event) + if isinstance(event_type, str) and event_type in _LIVE_USAGE_EVENT_TYPES: + return _usage_event(event, event_type) + if event_type == "response.event": + return _terminal_response_event(event) + return None + + +@dataclass(frozen=True, slots=True) +class LiveSessionStart: + frame: Mapping[str, object] + model: str + + +_LIVE_SESSION_START_ADAPTER: Final = TypeAdapter(dict[str, object]) + + +def parse_live_session_start(first_message: str) -> LiveSessionStart | None: + try: + frame: Final = _LIVE_SESSION_START_ADAPTER.validate_json(first_message.replace("\r", "").replace("\n", "")) + session: Final = _LIVE_SESSION_START_ADAPTER.validate_python(frame.get("session")) + except ValidationError: + return None + model: Final = session.get("model") + if frame.get("type") != "session.start" or not isinstance(model, str) or not model.strip(): + return None + return LiveSessionStart(frame=MappingProxyType(frame), model=model) + + +def _is_session_close_message(message: str) -> bool: + return _live_event_type(message) == "session.close" + + +def _is_session_closed_message(message: str) -> bool: + return _live_event_type(message) == "session.closed" + + +def _live_event_type(message: str) -> object: + try: + event: Final = _LIVE_EVENT_FRAME_ADAPTER.validate_json(message) + except ValidationError: + return None + return event.get("type") + + +_RetainedEvent: TypeAlias = tuple[int, Mapping[str, object]] + + +def _response_event_has_usage(event: Mapping[str, object]) -> bool: + nested_event: Final = as_str_mapping(event.get("event")) + response: Final = as_str_mapping(nested_event.get("response")) if nested_event is not None else None + return response is not None and as_str_mapping(response.get("usage")) is not None + + +class _LiveEventRetention: + def __init__(self) -> None: + self._event_sequence = count() + self._session_started: _RetainedEvent | None = None + self._latest_usage: _RetainedEvent | None = None + self._session_closed: _RetainedEvent | None = None + self._response_events: dict[str, _RetainedEvent] = {} # mutable-ok: one entry per response id, avoids copies + + def retain(self, event: Mapping[str, object]) -> None: + retained_item: Final = (next(self._event_sequence), event) + event_type: Final = event.get("type") + if event_type == "session.started": + if self._session_started is None: + self._session_started = retained_item + return + if event_type == "session.usage.updated": + usage: Final = as_str_mapping(event.get("usage")) + seconds: Final = usage.get("seconds") if usage is not None else None + if _is_finite_usage_seconds(seconds): + self._latest_usage = retained_item + return + if event_type == "session.closed": + self._session_closed = retained_item + return + nested_event: Final = as_str_mapping(event.get("event")) + response: Final = as_str_mapping(nested_event.get("response")) if nested_event is not None else None + response_id: Final = response.get("id") if response is not None else None + if not isinstance(response_id, str): + return + retained_response: Final = self._response_events.get(response_id) + if retained_response is None or not _response_event_has_usage(retained_response[1]): + self._response_events[response_id] = retained_item + + def results(self) -> tuple[Mapping[str, object], ...]: + retained_items: Final = tuple( + item for item in (self._session_started, self._latest_usage, self._session_closed) if item is not None + ) + tuple(self._response_events.values()) + return tuple(event for _, event in sorted(retained_items, key=lambda item: item[0])) + + +async def _forward_live_client_frames( + *, + websocket: LiveClientWebSocket, + backend_websocket: LiveBackendWebSocket, + client_gone: asyncio.Event, + client_close_sent: asyncio.Event, +) -> None: + async for message in _live_client_messages(websocket, client_gone): + await backend_websocket.send(message) + if _is_session_close_message(message): + client_close_sent.set() + + +async def _live_client_messages( + websocket: LiveClientWebSocket, + client_gone: asyncio.Event, +) -> AsyncIterator[str]: + from starlette.websockets import WebSocketDisconnect + + try: + while True: + yield await websocket.receive_text() + except (WebSocketDisconnect, RuntimeError): + client_gone.set() + + +async def _forward_live_backend_frames( + *, + websocket: LiveClientWebSocket, + backend_websocket: LiveBackendWebSocket, + client_gone: asyncio.Event, + event_retention: _LiveEventRetention, +) -> None: + from starlette.websockets import WebSocketDisconnect + from websockets.exceptions import ConnectionClosed + + event_loop: Final = asyncio.get_running_loop() + async for message, close_deadline in _live_backend_messages(websocket, backend_websocket): + if (event := _metering_event(message)) is not None: + event_retention.retain(event) + if client_gone.is_set(): + continue + try: + if close_deadline is None: + await websocket.send_text(message) + else: + await asyncio.wait_for( + websocket.send_text(message), + timeout=max(close_deadline - event_loop.time(), 0), + ) + except (TimeoutError, asyncio.TimeoutError, WebSocketDisconnect, RuntimeError, ConnectionClosed, OSError): + client_gone.set() + + +async def _live_backend_messages( + websocket: LiveClientWebSocket, + backend_websocket: LiveBackendWebSocket, +) -> AsyncIterator[tuple[str, float | None]]: + from starlette.websockets import WebSocketDisconnect + from websockets.exceptions import ConnectionClosed + + try: + while not _is_session_closed_message(message := _decode_live_websocket_frame(await backend_websocket.recv())): + yield message, None + except ConnectionClosed as error: + upstream_close: Final = backend_close_from(error) + with suppress(WebSocketDisconnect, RuntimeError, ConnectionClosed, OSError): + await websocket.close(code=upstream_close.code, reason=upstream_close.reason) + return + event_loop: Final = asyncio.get_running_loop() + close_deadline: Final = event_loop.time() + OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS + try: + yield message, close_deadline + while True: + yield ( + _decode_live_websocket_frame( + await asyncio.wait_for( + backend_websocket.recv(), + timeout=max(close_deadline - event_loop.time(), 0), + ) + ), + close_deadline, + ) + except (TimeoutError, asyncio.TimeoutError): + with suppress(WebSocketDisconnect, RuntimeError, ConnectionClosed, OSError): + await websocket.close(code=1000) + except ConnectionClosed as error: + with suppress(WebSocketDisconnect, RuntimeError, ConnectionClosed, OSError): + if error.rcvd is None: + await websocket.close(code=1000) + else: + await websocket.close(code=error.rcvd.code, reason=error.rcvd.reason) + + +def _decode_live_websocket_frame(frame: str | bytes) -> str: + return frame.decode("utf-8") if isinstance(frame, bytes) else frame + + +class OpenAILiveSessions(OpenAIRealtime): + def __init__( + self, + logging_worker: LoggingWorker = GLOBAL_LOGGING_WORKER, + websocket_connector: LiveWebSocketConnector | None = None, + clock: Callable[[], datetime] = datetime.now, + ) -> None: + super().__init__() + self._logging_worker = logging_worker + self._clock = clock + self._websocket_connector = ( + websocket_connector if websocket_connector is not None else _openai_live_websocket_connect + ) + + @staticmethod + def _construct_live_url(api_base: str) -> str: + parsed_url: Final = urlsplit(api_base) + websocket_scheme: Final = _LIVE_WEBSOCKET_SCHEMES.get(parsed_url.scheme, parsed_url.scheme) + return urlunsplit((websocket_scheme, parsed_url.netloc, "/v1/live/sessions", "", "")) + + async def async_live_session( + self, + *, + model: str, + websocket: LiveClientWebSocket, + logging_obj: LiteLLMLogging, + session_start: Mapping[str, object], + api_base: str | None, + api_key: str | None, + ) -> None: + if api_key is None: + raise ValueError("api_key is required for OpenAI Live session calls") + from websockets.exceptions import InvalidStatus + + resolved_api_base: Final = api_base or self._get_default_api_base() + url: Final = self._construct_live_url(resolved_api_base) + headers: Final = self._get_additional_headers(api_key) + ssl_config: Final = self._get_ssl_config(url) + session: Final = as_str_mapping(session_start.get("session")) + if session is None: + raise TypeError("session_start.session must be a mapping") + rewritten_start: Final = { + **session_start, + "session": {**session, "model": model}, + } + upstream_start: Final = json.dumps(rewritten_start) + pre_call_args: Final[_LivePreCallArgs] = { + "api_base": url, + "headers": headers, + "complete_input_dict": {"session_start": rewritten_start}, + } + event_retention: Final = _LiveEventRetention() + try: + logging_obj.pre_call(input=None, api_key=api_key, additional_args=pre_call_args) + async with self._websocket_connector( + url, + additional_headers=headers, + max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, + ssl=ssl_config, + ) as backend_websocket: + await backend_websocket.send(upstream_start) + try: + await self._bidirectional_forward( + websocket=websocket, + backend_websocket=backend_websocket, + event_retention=event_retention, + ) + finally: + session_end_time: Final = self._clock() + retained_events: Final = event_retention.results() + logging_result: Final = LiteLLMRealtimeStreamLoggingObject.model_validate( + {"usage": Usage(), "results": list(retained_events)} + ) + self._logging_worker.ensure_initialized_and_enqueue( + logging_obj.dispatch_success_handlers( + logging_result, end_time=session_end_time, prefer_async_handlers=True + ) + ) + logging_obj.model_call_details[ # rebind-ok: flags the shared logging object as already logged + REALTIME_SESSION_SUCCESS_LOGGED_KEY + ] = True + except InvalidStatus as error: + await close_after_upstream_handshake_refusal(websocket, error.response.status_code) + except Exception as error: # noqa: BLE001 # any relay failure must still close the client socket with a redacted reason + redacted_error: Final = redact_secrets(str(error)) + with suppress(Exception): + await websocket.send_text(realtime_error_event(redacted_error, error_type="server_error")) + try: + await websocket.close( + code=1011, + reason=websocket_close_reason(redacted_error, fallback="Internal server error"), + ) + except RuntimeError as close_error: + if "already completed" not in str(close_error) and "websocket.close" not in str(close_error): + raise RuntimeError(f"Unexpected error while closing WebSocket: {close_error}") from close_error + + async def _bidirectional_forward( + self, + *, + websocket: LiveClientWebSocket, + backend_websocket: LiveBackendWebSocket, + event_retention: _LiveEventRetention, + ) -> None: + from websockets.exceptions import ConnectionClosed + + client_gone: Final = asyncio.Event() + client_close_sent: Final = asyncio.Event() + client_forward_task: Final = asyncio.create_task( + _forward_live_client_frames( + websocket=websocket, + backend_websocket=backend_websocket, + client_gone=client_gone, + client_close_sent=client_close_sent, + ) + ) + backend_forward_task: Final = asyncio.create_task( + _forward_live_backend_frames( + websocket=websocket, + backend_websocket=backend_websocket, + client_gone=client_gone, + event_retention=event_retention, + ) + ) + client_gone_task: Final = asyncio.create_task(client_gone.wait()) + forwarding_tasks: Final = (client_forward_task, backend_forward_task, client_gone_task) + try: + completed_tasks: Final = await asyncio.wait( + forwarding_tasks, + return_when=asyncio.FIRST_COMPLETED, + ) + done_tasks: Final = completed_tasks[0] + if backend_forward_task in done_tasks: + backend_forward_task.result() + return + client_forward_error: Final = client_forward_task.exception() if client_forward_task in done_tasks else None + if client_forward_task not in done_tasks: + client_forward_task.cancel() + await asyncio.gather(client_forward_task, return_exceptions=True) + event_loop: Final = asyncio.get_running_loop() + close_deadline: Final = event_loop.time() + OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS + if not client_close_sent.is_set(): + remaining_send_timeout: Final = max(close_deadline - event_loop.time(), 0) + with suppress(ConnectionClosed, TimeoutError, asyncio.TimeoutError): + await asyncio.wait_for( + backend_websocket.send(_LIVE_SESSION_CLOSE_FRAME), + timeout=remaining_send_timeout, + ) + remaining_wait_timeout: Final = max(close_deadline - event_loop.time(), 0) + close_wait_result: Final = await asyncio.wait( + (backend_forward_task,), + timeout=remaining_wait_timeout, + ) + remaining_close_timeout: Final = max(close_deadline - event_loop.time(), 0) + with suppress(TimeoutError, asyncio.TimeoutError): + await asyncio.wait_for( + backend_websocket.close(), + timeout=remaining_close_timeout, + ) + if backend_forward_task in close_wait_result[0]: + backend_forward_task.result() + if client_forward_error is not None and not isinstance(client_forward_error, ConnectionClosed): + raise client_forward_error + return + backend_forward_task.cancel() + if client_forward_error is not None: + raise client_forward_error + finally: + for task in forwarding_tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*forwarding_tasks, return_exceptions=True) diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index 325169cfe15..853f707bc91 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -45,7 +45,7 @@ class OpenAIRealtime(OpenAIChatCompletion): api_key: str, *, openai_beta_realtime: bool = False, - ) -> dict: + ) -> dict[str, str]: """ Headers for the upstream OpenAI Realtime WebSocket. @@ -54,7 +54,7 @@ class OpenAIRealtime(OpenAIChatCompletion): so the legacy beta API is used. GA clients omit that header on the client connection and must send GA-shaped ``session.update`` payloads. """ - headers: Final[dict] = {"Authorization": f"Bearer {api_key}"} + headers: Final[dict[str, str]] = {"Authorization": f"Bearer {api_key}"} if openai_beta_realtime: headers["OpenAI-Beta"] = "realtime=v1" return headers diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 642d932a712..08bd4bcba2e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -318,6 +318,9 @@ class LiteLLMRoutes(enum.Enum): "/realtime", "/v1/realtime", "/openai/v1/realtime", + "/live/sessions", + "/v1/live/sessions", + "/openai/v1/live/sessions", "/realtime?{model}", "/v1/realtime?{model}", "/openai/v1/realtime?{model}", diff --git a/litellm/proxy/agent_endpoints/auth/managed_authorization.py b/litellm/proxy/agent_endpoints/auth/managed_authorization.py index 64e50cd74e7..cbbd1548301 100644 --- a/litellm/proxy/agent_endpoints/auth/managed_authorization.py +++ b/litellm/proxy/agent_endpoints/auth/managed_authorization.py @@ -12,6 +12,7 @@ from litellm.types.agents import AgentResponse from litellm.types.proxy.agent_identity import AgentIdentityFailure, ManagedAgentContext _MANAGED_REALTIME_ROUTES: Final = frozenset(("/realtime", "/v1/realtime", "/openai/v1/realtime")) +_MANAGED_LIVE_ROUTES: Final = frozenset(("/live/sessions", "/v1/live/sessions", "/openai/v1/live/sessions")) _MANAGED_MODEL_ROUTES: Final = frozenset( f"{prefix}/{operation}" for prefix, operation in product( @@ -87,7 +88,7 @@ def managed_agent_route_allowed(route: str, method: str | None) -> bool: if route in ("/agents", "/v1/agents"): return method in (None, "GET", "HEAD") - if route in _MANAGED_REALTIME_ROUTES: + if route in _MANAGED_REALTIME_ROUTES or route in _MANAGED_LIVE_ROUTES: return method in (None, "GET") if route in _MANAGED_MODEL_ROUTES or RouteChecks.check_route_access(route, _MANAGED_MODEL_PATHS): return method in (None, "POST") @@ -106,6 +107,8 @@ def managed_inference_request( ) -> dict[str, object]: from litellm.proxy.auth.route_checks import RouteChecks + if route in _MANAGED_LIVE_ROUTES: + return dict(body) if route in _MANAGED_REALTIME_ROUTES: model: Final = query_model or body.get("model") if not isinstance(model, str) or not model: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 9e7fa2c28e8..5568ae2995c 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1255,24 +1255,7 @@ async def common_checks( code=status.HTTP_400_BAD_REQUEST, ) - managed_policy: Final = managed_agent_policy(valid_token) - if _model and valid_token is not None and managed_policy is not None: - managed_models: Final = (managed_policy.object_permission or MappingProxyType({})).get("models", ()) - if not isinstance(managed_models, (list, tuple)) or not managed_models: - raise HTTPException(403, "This agent has no model grants") - can_object_call_model( - model=_resolve_team_alias( - _model, team_model_aliases_for_auth_check(valid_token), valid_token.team_id, llm_router - ), - llm_router=llm_router, - models=list(managed_models), - team_id=valid_token.team_id, - object_type="agent", - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - - await _check_agent_access_group_model_access(model=_model, valid_token=valid_token, llm_router=llm_router) - await _check_agent_caller_model_access( + await can_agent_call_model( model=_model, valid_token=valid_token, llm_router=llm_router, @@ -4777,7 +4760,7 @@ def _can_object_call_model( raise Exception(f"Unable to parse model, max fallback depth exceeded - received model: {model}") if isinstance(model, list): for m in model: - _can_object_call_model( + can_object_call_model( model=m, llm_router=llm_router, models=models, @@ -4877,6 +4860,40 @@ def _live_team_alias_target( return model if deleted_team_deployment else target +async def can_agent_call_model( + model: str | list[str] | None, # mutable-ok: the model checks it delegates to take list[str] + valid_token: UserAPIKeyAuth | None, + llm_router: Router | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> None: + managed_policy: Final = managed_agent_policy(valid_token) + if model and valid_token is not None and managed_policy is not None: + managed_models: Final = (managed_policy.object_permission or {}).get("models", ()) + if not isinstance(managed_models, (list, tuple)) or not managed_models: + raise HTTPException(403, "This agent has no model grants") + can_object_call_model( + model=_resolve_team_alias( + model, team_model_aliases_for_auth_check(valid_token), valid_token.team_id, llm_router + ), + llm_router=llm_router, + models=list(managed_models), + team_id=valid_token.team_id, + object_type="agent", + key_model_aliases=key_model_aliases_for_auth_check(valid_token), + ) + await _check_agent_access_group_model_access(model=model, valid_token=valid_token, llm_router=llm_router) + await _check_agent_caller_model_access( + model=model, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + async def _check_agent_access_group_model_access( model: str | list[str] | None, # mutable-ok: _can_object_call_model and the client message helper take list[str] valid_token: UserAPIKeyAuth | None, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index cdb3e685c20..7edc183c8e7 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1181,7 +1181,13 @@ async def user_api_key_auth_websocket(websocket: WebSocket) -> UserAPIKeyAuth: return await user_api_key_auth_websocket_for_model(websocket, model=websocket.query_params.get("model")) -async def user_api_key_auth_websocket_for_model(websocket: WebSocket, model: str | None) -> UserAPIKeyAuth: +async def user_api_key_auth_live_websocket(websocket: WebSocket) -> UserAPIKeyAuth: + return await user_api_key_auth_websocket_for_model(websocket, model=None) + + +async def user_api_key_auth_websocket_for_model( + websocket: WebSocket, model: str | None, *, close_on_auth_error: bool = True +) -> UserAPIKeyAuth: ws_scope: Final = websocket.scope or {} scope_headers: Final = list(ws_scope.get("headers") or []) # ``get_request_route`` falls back to ``request.url.path`` when @@ -1236,6 +1242,8 @@ async def user_api_key_auth_websocket_for_model(websocket: WebSocket, model: str try: return await user_api_key_auth(request=request, api_key=f"Bearer {api_key}") except Exception as e: + if not close_on_auth_error: + raise if is_invalid_virtual_key_error(e): raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION) log_model_access_denial(e) @@ -1243,6 +1251,16 @@ async def user_api_key_auth_websocket_for_model(websocket: WebSocket, model: str raise HTTPException(status_code=403, detail=str(e)) +async def user_api_key_auth_accepted_websocket_for_model(websocket: WebSocket, model: str) -> UserAPIKeyAuth: + """Re-authorize an already accepted WebSocket for a model it named after the upgrade. + + Runs the same key, team, member, user, project and agent checks, plus budget reservation, as + the upgrade does for ``?model=``, but raises auth errors without closing the socket so the + caller can reply with an error frame. + """ + return await user_api_key_auth_websocket_for_model(websocket, model, close_on_auth_error=False) + + def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_user_params: dict) -> UserAPIKeyAuth: valid_token.end_user_id = end_user_params.get("end_user_id") # Only overwrite token fields when the DB-derived value is not None. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 60b9ade82c0..3a831062b0a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -32,6 +32,7 @@ from itertools import chain from types import MappingProxyType, UnionType from typing import ( TYPE_CHECKING, + Annotated, Any, Final, Literal, @@ -73,6 +74,7 @@ from litellm.constants import ( LITELLM_SETTINGS_SAFE_DB_OVERRIDES, LITELLM_UI_ALLOW_HEADERS, LITELLM_UI_SESSION_DURATION, + OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS, RUNTIME_UPDATABLE_ROUTER_SETTINGS, ) from litellm.litellm_core_utils.asyncify import asyncify @@ -334,6 +336,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( mask_sensitive_keys, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.openai.live import parse_live_session_start from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._experimental.mcp_server.byok_credential_cache import byok_credential_cache @@ -347,6 +350,7 @@ from litellm.proxy.analytics_endpoints.analytics_endpoints import ( from litellm.proxy.auth.auth_checks import ( ROLE_BASED_PERMISSIONS_ADAPTER, ExperimentalUIJWTToken, + can_agent_call_model, can_key_call_resolved_model, get_team_object, log_db_metrics, @@ -391,6 +395,8 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401, RUF100 # lega _fetch_global_spend_with_event_coordination, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export fetch_global_spend_with_event_coordination, user_api_key_auth, + user_api_key_auth_accepted_websocket_for_model, + user_api_key_auth_live_websocket, user_api_key_auth_websocket, ) from litellm.proxy.batches_endpoints.endpoints import router as batches_router @@ -3917,14 +3923,8 @@ async def invalidate_spend_counter(counter_key: str) -> None: forget_spend_counter(counter_key) spend_counter_cache.in_memory_cache.delete_cache(key=counter_key) if spend_counter_cache.redis_cache is not None: - try: + with contextlib.suppress(Exception): await spend_counter_cache.redis_cache.async_delete_cache(key=counter_key) - except Exception: - verbose_proxy_logger.debug( - "Unable to delete stale spend counter %s after increment failure", - counter_key, - exc_info=True, - ) _invalidate_spend_counter: Final = invalidate_spend_counter @@ -13202,6 +13202,17 @@ async def _release_realtime_max_parallel_slot(user_api_key_dict: UserAPIKeyAuth) await release_like_http_disconnect(user_api_key_dict) +class _RealtimeErrorBody(TypedDict): + type: ReadOnly[str] + message: ReadOnly[str] + code: NotRequired[ReadOnly[str]] + + +class _RealtimeErrorEvent(TypedDict): + type: ReadOnly[str] + error: ReadOnly[_RealtimeErrorBody] + + async def _reject_realtime_session( websocket: WebSocket, user_api_key_dict: UserAPIKeyAuth, @@ -13209,13 +13220,19 @@ async def _reject_realtime_session( code: int, reason: str, error_message: str | None = None, + error_type: str = "guardrail_error", + error_code: str | None = None, ) -> None: try: if error_message is not None: try: - await websocket.send_text( - json.dumps({"type": "error", "error": {"type": "guardrail_error", "message": error_message}}) - ) + error_event: Final[_RealtimeErrorEvent] = { + "type": "error", + "error": {"type": error_type, "message": error_message} + if error_code is None + else {"type": error_type, "message": error_message, "code": error_code}, + } + await websocket.send_text(json.dumps(error_event)) except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below verbose_proxy_logger.debug("Could not send realtime pre-call error event to client; closing anyway") await websocket.close(code=code, reason=reason) @@ -13237,12 +13254,11 @@ async def realtime_websocket_endpoint( ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket), ): - requested_protocols: Final = [ - p.strip() for p in (websocket.headers.get("sec-websocket-protocol") or "").split(",") if p.strip() - ] - accept_kwargs: Final[dict] = {} - if requested_protocols: - accept_kwargs["subprotocol"] = requested_protocols[0] + requested_protocols: Final = tuple( + protocol.strip() + for protocol in (websocket.headers.get("sec-websocket-protocol") or "").split(",") + if protocol.strip() + ) route_model = model if route_model is None: @@ -13265,26 +13281,47 @@ async def realtime_websocket_endpoint( log_model_access_denial(e) await _reject_realtime_session(websocket, user_api_key_dict, code=1008, reason=e.message[:120]) return - await websocket.accept(**accept_kwargs) + await websocket.accept(subprotocol=requested_protocols[0] if requested_protocols else None) + await _route_realtime_websocket_session( + websocket, + user_api_key_dict, + route_model=route_model, + model=model, + intent=intent, + guardrails=guardrails, + ) + +async def _route_realtime_websocket_session( + websocket: WebSocket, + user_api_key_dict: UserAPIKeyAuth, + *, + route_model: str, + model: str | None, + intent: str | None, + guardrails: str | None, + request_path: str | None = None, + live_session_start: Mapping[str, object] | None = None, +) -> None: # Only use explicit parameters, not all query params query_params: Final = cast(RealtimeQueryParams, dict(_realtime_query_params_template(model, intent))) - data: dict[str, object] = { "model": route_model, "websocket": websocket, "query_params": query_params, # Only explicit params } - # Pass guardrails into data so pre-call guardrail processing picks them up if guardrails: data["guardrails"] = [g.strip() for g in guardrails.split(",") if g.strip()] + if live_session_start is not None: + data["live_session_start"] = live_session_start # Use raw ASGI headers (already lowercase bytes) to avoid extra work headers_list: Final = list(websocket.scope.get("headers") or []) - scope: Final = REALTIME_REQUEST_SCOPE_TEMPLATE.copy() scope["headers"] = headers_list + if request_path is not None: + scope["path"] = request_path request: Final = Request(scope=scope) @@ -13327,6 +13364,7 @@ async def realtime_websocket_endpoint( ) return except BaseException: + await _release_realtime_budget_reservation(user_api_key_dict) await _release_realtime_max_parallel_slot(user_api_key_dict) raise @@ -13364,6 +13402,141 @@ async def realtime_websocket_endpoint( await _release_realtime_max_parallel_slot(user_api_key_dict) +async def _reject_invalid_live_session_start( + websocket: WebSocket, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + await _reject_realtime_session( + websocket, + user_api_key_dict, + code=1008, + reason="Invalid session.start frame", + error_message="First message must be a session.start JSON object with a non-empty session.model.", + error_type="invalid_request_error", + error_code="invalid_session_start", + ) + + +LiveSessionModelAuthorizer: TypeAlias = Callable[[WebSocket, str], Awaitable[UserAPIKeyAuth]] + + +def live_session_model_authorizer() -> LiveSessionModelAuthorizer: + return user_api_key_auth_accepted_websocket_for_model + + +async def _authorize_live_session_model( + authorize_session_model: LiveSessionModelAuthorizer, + websocket: WebSocket, + model: str, +) -> UserAPIKeyAuth: + session_auth: Final = await authorize_session_model(websocket, model) + try: + await can_key_call_resolved_model( + model=model, + llm_model_list=llm_model_list, + valid_token=session_auth, + llm_router=llm_router, + ) + await can_agent_call_model( + model=model, + valid_token=session_auth, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except BaseException: + await _release_realtime_budget_reservation(session_auth) + raise + return session_auth + + +def _live_session_denial_code(error: ProxyException | HTTPException) -> str: + model_denied: Final = isinstance(error, ProxyException) and str(error.type).endswith("_model_access_denied") + return "model_not_allowed" if model_denied else "session_not_allowed" + + +@app.websocket("/openai/v1/live/sessions") +@app.websocket("/v1/live/sessions") +@app.websocket("/live/sessions") +async def live_sessions_websocket_endpoint( + websocket: WebSocket, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_live_websocket)], + authorize_session_model: Annotated[LiveSessionModelAuthorizer, Depends(live_session_model_authorizer)], +) -> None: + requested_protocols: Final = tuple( + protocol.strip() + for protocol in (websocket.headers.get("sec-websocket-protocol") or "").split(",") + if protocol.strip() + ) + try: + await websocket.accept(subprotocol=requested_protocols[0] if requested_protocols else None) + except BaseException: + await _release_realtime_budget_reservation(user_api_key_dict) + await _release_realtime_max_parallel_slot(user_api_key_dict) + raise + + try: + first_message: Final = await asyncio.wait_for( + websocket.receive_text(), + timeout=OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS, + ) + except asyncio.TimeoutError: + await _reject_invalid_live_session_start(websocket, user_api_key_dict) + return + except WebSocketDisconnect: + await _release_realtime_budget_reservation(user_api_key_dict) + await _release_realtime_max_parallel_slot(user_api_key_dict) + return + except RuntimeError: + await _reject_invalid_live_session_start(websocket, user_api_key_dict) + return + except BaseException: + await _release_realtime_budget_reservation(user_api_key_dict) + await _release_realtime_max_parallel_slot(user_api_key_dict) + raise + + session_start: Final = parse_live_session_start(first_message) + if session_start is None: + await _reject_invalid_live_session_start(websocket, user_api_key_dict) + return + + upgrade_auth: Final = user_api_key_dict.model_copy() + await _release_realtime_budget_reservation(upgrade_auth) + try: + session_auth: Final = await _authorize_live_session_model( + authorize_session_model, websocket, session_start.model + ) + except (ProxyException, HTTPException) as error: + denial_message: Final = error.message if isinstance(error, ProxyException) else str(error.detail) + if isinstance(error, ProxyException): + log_model_access_denial(error) + await _reject_realtime_session( + websocket, + upgrade_auth, + code=1008, + reason=websocket_close_reason(denial_message, fallback="Session not allowed"), + error_message=denial_message, + error_type="invalid_request_error", + error_code=_live_session_denial_code(error), + ) + return + except BaseException: + await _release_realtime_max_parallel_slot(upgrade_auth) + raise + + await _route_realtime_websocket_session( + websocket, + session_auth, + route_model=session_start.model, + model=None, + intent=None, + guardrails=None, + request_path=websocket.url.path, + live_session_start=session_start.frame, + ) + + ###################################################################### # /v1/assistant Endpoints diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 0e79cbfda13..d6bb44b1d6b 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -3,8 +3,11 @@ import asyncio import os from collections.abc import Mapping +from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, cast + +from typing_extensions import ReadOnly, TypedDict, assert_never import litellm from litellm.constants import ( @@ -37,6 +40,7 @@ from ..llms.azure.common_utils import get_azure_ad_token from ..llms.azure.realtime.handler import AzureOpenAIRealtime, azure_realtime_protocol_for_client from ..llms.bedrock.realtime.handler import BedrockRealtime from ..llms.custom_httpx.http_handler import realtime_ssl_for_url +from ..llms.openai.live.handler import OpenAILiveSessions from ..llms.openai.realtime.handler import OpenAIRealtime from ..llms.vertex_ai.audio_transcription.realtime_transformation import is_vertex_speech_to_text_model from ..llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig, vertex_realtime_config @@ -51,6 +55,7 @@ if TYPE_CHECKING: azure_realtime: Final = AzureOpenAIRealtime() openai_realtime: Final = OpenAIRealtime() +openai_live_sessions: Final = OpenAILiveSessions() bedrock_realtime: Final = BedrockRealtime() xai_realtime: Final = XAIRealtime() vertex_llm_base: Final = VertexBase() @@ -75,9 +80,9 @@ def _with_resolved_session_model(session: dict[str, object], model_name: str) -> return {**session, "model": model_name} -def _build_litellm_metadata(kwargs: dict) -> dict: +def _build_litellm_metadata(kwargs: dict) -> dict[str, object]: """Build the litellm_metadata dict for guardrail checking (internal only, not forwarded to provider).""" - metadata: Final[dict] = {**(kwargs.get("litellm_metadata") or {})} + metadata: Final[dict[str, object]] = {**(kwargs.get("litellm_metadata") or {})} guardrails: Final = (kwargs.get("metadata") or {}).get("guardrails") or kwargs.get("guardrails") or [] if guardrails: metadata["guardrails"] = guardrails @@ -345,6 +350,84 @@ async def _resolve_vertex_access_token_bounded( ) from e +class _RealtimeGuardrailRequestData(TypedDict): + litellm_metadata: ReadOnly[Mapping[str, object]] + stream: ReadOnly[Literal[True]] + + +def _realtime_guardrail_configured(litellm_metadata: Mapping[str, object]) -> bool: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import GuardrailEventHooks + + event_hooks: Final = ( + GuardrailEventHooks.realtime_input_transcription, + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ) + request_data: Final[_RealtimeGuardrailRequestData] = {"litellm_metadata": litellm_metadata, "stream": True} + return any( + isinstance(callback, CustomGuardrail) + and any(callback.should_run_guardrail(data=request_data, event_type=hook) for hook in event_hooks) + for callback in litellm.callbacks + ) + + +@dataclass(frozen=True, slots=True) +class _LiveProviderUnsupported: + provider: str + + +@dataclass(frozen=True, slots=True) +class _LiveGuardrailsUnsupported: + pass + + +_LiveSessionRejection = _LiveProviderUnsupported | _LiveGuardrailsUnsupported + + +def _live_session_rejection( + custom_llm_provider: str, litellm_metadata: Mapping[str, object] +) -> _LiveSessionRejection | None: + if custom_llm_provider != "openai": + return _LiveProviderUnsupported(provider=custom_llm_provider) + if _realtime_guardrail_configured(litellm_metadata): + return _LiveGuardrailsUnsupported() + return None + + +def _raise_live_session_rejection(rejection: _LiveSessionRejection) -> NoReturn: + match rejection: + case _LiveProviderUnsupported(provider=provider): + raise ValueError(f"OpenAI Live sessions require the openai provider, got {provider}") + case _LiveGuardrailsUnsupported(): + raise ValueError("Guardrails are not supported on OpenAI Live sessions") + case _: + assert_never(rejection) + + +async def _alive_session( + model: str, + websocket: "WebSocket", + custom_llm_provider: str, + litellm_logging_obj: LiteLLMLogging, + session_start: Mapping[str, object], + api_base: str | None, + api_key: str | None, + litellm_metadata: Mapping[str, object], +) -> None: + rejection: Final = _live_session_rejection(custom_llm_provider, litellm_metadata) + if rejection is not None: + _raise_live_session_rejection(rejection) + await openai_live_sessions.async_live_session( + model=model, + websocket=websocket, + logging_obj=litellm_logging_obj, + session_start=session_start, + api_base=api_base or litellm.api_base or "https://api.openai.com/", + api_key=api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY"), + ) + + @wrapper_client async def _arealtime( model: str, @@ -356,6 +439,7 @@ async def _arealtime( client: object | None = None, timeout: float | None = None, query_params: RealtimeQueryParams | None = None, + live_session_start: Mapping[str, object] | None = None, **kwargs, ) -> None: """ @@ -398,6 +482,19 @@ async def _arealtime( custom_llm_provider=_custom_llm_provider, ) + if live_session_start is not None: + await _alive_session( + model=model, + websocket=websocket, + custom_llm_provider=_custom_llm_provider, + litellm_logging_obj=litellm_logging_obj, + session_start=live_session_start, + api_base=dynamic_api_base or litellm_params.api_base, + api_key=dynamic_api_key, + litellm_metadata=_build_litellm_metadata(kwargs), + ) + return + provider_config: BaseRealtimeConfig | None = None if _custom_llm_provider in LlmProviders._member_map_.values(): provider_config = ProviderConfigManager.get_provider_realtime_config( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2d98bf0486f..93037d17062 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1058,6 +1058,9 @@ API_ROUTE_TO_CALL_TYPES: Final[Mapping[str, Sequence[CallTypes]]] = { "/realtime": [CallTypes.arealtime], "/v1/realtime": [CallTypes.arealtime], "/openai/v1/realtime": [CallTypes.arealtime], + "/live/sessions": (CallTypes.arealtime,), + "/v1/live/sessions": (CallTypes.arealtime,), + "/openai/v1/live/sessions": (CallTypes.arealtime,), # Provider-specific routes "/anthropic/v1/messages": [CallTypes.anthropic_messages], # Google GenAI routes diff --git a/tests/unit/llms/openai/live/__init__.py b/tests/unit/llms/openai/live/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/live/test_handler.py b/tests/unit/llms/openai/live/test_handler.py new file mode 100644 index 00000000000..ec1130376d0 --- /dev/null +++ b/tests/unit/llms/openai/live/test_handler.py @@ -0,0 +1,1208 @@ +import asyncio +import json +from collections.abc import Coroutine +from datetime import datetime +from typing import Final, NoReturn, cast + +import pytest +from pydantic import TypeAdapter +from starlette.websockets import WebSocketDisconnect +from websockets.datastructures import Headers +from websockets.exceptions import ConnectionClosed, ConnectionClosedError, ConnectionClosedOK, InvalidStatus +from websockets.frames import Close +from websockets.http11 import Response + +import litellm +from litellm.constants import REALTIME_SESSION_SUCCESS_LOGGED_KEY, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES +from litellm.cost_calculator import handle_realtime_stream_cost_calculation +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.openai.live import handler as live_handler +from litellm.llms.openai.live.handler import LiveSessionStart, OpenAILiveSessions, parse_live_session_start +from litellm.types.utils import LiteLLMRealtimeStreamLoggingObject, Usage + +_LIVE_EVENT_FRAME_ADAPTER: Final = TypeAdapter(dict[str, object]) + + +def _event_type(message: str) -> str | None: + event_type: Final = _LIVE_EVENT_FRAME_ADAPTER.validate_json(message).get("type") + return event_type if isinstance(event_type, str) else None + + +class _FakeClientWebSocket: + def __init__( + self, + messages: tuple[str, ...], + *, + disconnect_after_send_type: str | None = None, + disconnect_after_receive_type: str | None = None, + fail_on_send_type: str | None = None, + block_on_send_type: str | None = None, + receive_error_after_close: Exception | None = None, + ) -> None: + self._messages = iter(messages) + self._block_on_send_type = block_on_send_type + self._receive_error_after_close = receive_error_after_close + self._disconnect_after_send_type = disconnect_after_send_type + self._disconnect_after_receive_type = disconnect_after_receive_type + self._fail_on_send_type = fail_on_send_type + self.sent_text: asyncio.Queue[str] = asyncio.Queue() + self.closed = asyncio.Event() + self.close_info: tuple[int, str | None] | None = None + + async def receive_text(self) -> str: + try: + message: Final = next(self._messages) + if _event_type(message) == self._disconnect_after_receive_type: + self.closed.set() + return message + except StopIteration: + await self.closed.wait() + raise self._receive_error_after_close or WebSocketDisconnect(code=1000) + + async def send_text(self, data: str) -> None: + event_type: Final = _event_type(data) + if event_type == self._fail_on_send_type: + raise RuntimeError("client is disconnected") + if event_type == self._block_on_send_type: + await asyncio.Event().wait() + await self.sent_text.put(data) + if event_type == self._disconnect_after_send_type: + self.closed.set() + + async def close(self, code: int = 1000, reason: str | None = None) -> None: + self.close_info = (code, reason) + self.closed.set() + + +class _FakeUpstream: + def __init__( + self, + client_frames: tuple[str, ...], + upstream_frames: tuple[str, ...], + *, + close_response: str | None = None, + never_respond_after_close: bool = False, + block_session_close_send: bool = False, + block_close: bool = False, + fail_on_send_type: str | None = None, + send_error: Exception | None = None, + repeat_last_frame_forever: bool = False, + close_error: ConnectionClosed | None = None, + ) -> None: + self.client_frames = client_frames + self._close_error = close_error + self._repeat_last_frame_forever = repeat_last_frame_forever + self._last_frame: str | None = None + self._fail_on_send_type = fail_on_send_type + self._send_error = send_error + self._upstream_frames = iter(upstream_frames) + self._close_response = close_response + self._close_response_sent = False + self._never_respond_after_close = never_respond_after_close + self._block_session_close_send = block_session_close_send + self._block_close = block_close + self.sent_text: asyncio.Queue[str] = asyncio.Queue() + self.client_frames_sent = asyncio.Event() + self.session_close_sent = asyncio.Event() + self.closed_event = asyncio.Event() + self.closed = False + + async def send(self, data: str) -> None: + event_type: Final = _event_type(data) + if event_type == self._fail_on_send_type: + raise self._send_error or ConnectionClosedError(None, Close(code=1011, reason="upstream send failed"), None) + if event_type == "session.close" and self._block_session_close_send: + await asyncio.Event().wait() + await self.sent_text.put(data) + if event_type == "session.close": + self.session_close_sent.set() + if self.sent_text.qsize() == len(self.client_frames) + 1: + self.client_frames_sent.set() + + async def recv(self) -> str: + await self.client_frames_sent.wait() + try: + self._last_frame = next(self._upstream_frames) + return self._last_frame + except StopIteration as error: + if self._repeat_last_frame_forever and self._last_frame is not None: + await asyncio.sleep(0) + return self._last_frame + if self._close_response is not None and not self._close_response_sent: + await self.session_close_sent.wait() + self._close_response_sent = True + return self._close_response + if self._never_respond_after_close: + await self.session_close_sent.wait() + await self.closed_event.wait() + raise self._close_error or ConnectionClosedOK( + Close(code=1000, reason="upstream closed"), + Close(code=1000, reason="upstream closed"), + True, + ) from error + + async def close(self) -> None: + if self._block_close: + await asyncio.Event().wait() + self.closed = True + self.closed_event.set() + + async def __aenter__(self) -> "_FakeUpstream": + return self + + async def __aexit__(self, exc_type: object, exc: object, traceback: object) -> None: + self.closed = True + self.closed_event.set() + + +class _FailingConnection: + def __init__(self, error: Exception) -> None: + self.error = error + + async def __aenter__(self) -> NoReturn: + raise self.error + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: object | None, + ) -> None: + return None + + +class _FakeLogging: + def __init__(self, pre_call_error: Exception | None = None) -> None: + self.model_call_details: dict[str, object] = {} + self.pre_call_kwargs: dict[str, object] = {} + self.dispatch_count = 0 + self.dispatched_result: object | None = None + self.dispatched_end_time: datetime | None = None + self.prefer_async_handlers = False + self._pre_call_error = pre_call_error + + def pre_call(self, **kwargs: object) -> None: + if self._pre_call_error is not None: + raise self._pre_call_error + self.pre_call_kwargs = kwargs + + async def dispatch_success_handlers( + self, + result: object, + *, + end_time: datetime | None = None, + prefer_async_handlers: bool = False, + ) -> None: + self.dispatch_count += 1 + self.dispatched_result = result + self.dispatched_end_time = end_time + self.prefer_async_handlers = prefer_async_handlers + + +class _FakeLoggingWorker: + def __init__(self) -> None: + self.coroutine: Coroutine[object, object, None] | None = None + + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: + self.coroutine = async_coroutine + + +class _FakeConnector: + def __init__(self, upstream: _FakeUpstream) -> None: + self.upstream = upstream + self.url: str | None = None + self.kwargs: dict[str, object] = {} + + def __call__(self, url: str, **kwargs: object) -> _FakeUpstream: + self.url = url + self.kwargs = kwargs + return self.upstream + + +@pytest.mark.parametrize( + ("api_base", "expected_url"), + [ + ("https://api.openai.com/", "wss://api.openai.com/v1/live/sessions"), + ("https://api.openai.com/v1", "wss://api.openai.com/v1/live/sessions"), + ("http://localhost:8080", "ws://localhost:8080/v1/live/sessions"), + ("https://api.openai.com/v1?api-version=1", "wss://api.openai.com/v1/live/sessions"), + ], +) +def test_construct_live_url_uses_exact_path_and_drops_query(api_base: str, expected_url: str) -> None: + assert OpenAILiveSessions._construct_live_url(api_base) == expected_url + + +@pytest.mark.asyncio +async def test_async_live_session_rewrites_model_and_relays_frames_without_mutating_input() -> None: + client_frames: Final = ('{"type":"input_audio.append","audio":"client-a"}', '{"type":"session.close"}') + upstream_frames: Final = ( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + '{"type":"response.audio.delta","delta":"server-a"}', + '{"type":"session.closed","usage":{"seconds":16}}', + ) + upstream: Final = _FakeUpstream(client_frames=client_frames, upstream_frames=upstream_frames) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=client_frames) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + session_start: Final[dict[str, object]] = { + "type": "session.start", + "session": { + "model": "proxy-live-alias", + "instructions": "Be concise", + "audio": {"input": {"format": "pcm16"}}, + }, + } + original_session_start: Final = json.loads(json.dumps(session_start)) + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start=session_start, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + sent_frames: Final = tuple( + await asyncio.gather(*(upstream.sent_text.get() for _ in range(upstream.sent_text.qsize()))) + ) + forwarded_frames: Final = tuple( + await asyncio.gather(*(client.sent_text.get() for _ in range(client.sent_text.qsize()))) + ) + rewritten_start: Final = json.loads(sent_frames[0]) + assert connector.url == "wss://api.openai.com/v1/live/sessions" + assert connector.kwargs["additional_headers"] == {"Authorization": "Bearer live-test-key"} + assert connector.kwargs["max_size"] == REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES + assert "ssl" in connector.kwargs + assert rewritten_start == { + "type": "session.start", + "session": { + "model": "gpt-live-1", + "instructions": "Be concise", + "audio": {"input": {"format": "pcm16"}}, + }, + } + assert sent_frames[1:] == client_frames + assert forwarded_frames == upstream_frames + assert session_start == original_session_start + assert client.close_info == (1000, "upstream closed") + assert logger.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] is True + assert worker.coroutine is not None + await worker.coroutine + assert logger.dispatch_count == 1 + assert logger.prefer_async_handlers is True + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert tuple(event["type"] for event in logged_result.results) == ( + "session.started", + "session.closed", + ) + + +@pytest.mark.asyncio +async def test_async_live_session_coalesces_cumulative_usage_updates() -> None: + client_frames: Final = ('{"type":"session.close"}',) + upstream_frames: Final = ( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + '{"type":"response.event","event":{"type":"response.completed","response":{"id":"resp-1"}}}', + '{"type":"response.event","event":{"type":"response.failed","response":{"id":"resp-2"}}}', + '{"type":"session.usage.updated","usage":{"seconds":1}}', + '{"type":"session.usage.updated","usage":{"seconds":2}}', + '{"type":"session.usage.updated","usage":{"seconds":3}}', + '{"type":"session.closed","usage":{"seconds":3}}', + ) + upstream: Final = _FakeUpstream(client_frames=client_frames, upstream_frames=upstream_frames) + connector: Final = _FakeConnector(upstream) + worker: Final = _FakeLoggingWorker() + logger: Final = _FakeLogging() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=_FakeClientWebSocket(messages=client_frames), + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert tuple((event["type"], event.get("usage")) for event in logged_result.results) == ( + ("session.started", None), + ("response.event", None), + ("response.event", None), + ("session.usage.updated", {"seconds": 3}), + ("session.closed", {"seconds": 3}), + ) + assert tuple(logged_result.results[1:3]) == ( + { + "type": "response.event", + "event": { + "type": "response.completed", + "response": {"id": "resp-1", "model": None, "usage": None}, + }, + }, + { + "type": "response.event", + "event": { + "type": "response.failed", + "response": {"id": "resp-2", "model": None, "usage": None}, + }, + }, + ) + + +@pytest.mark.asyncio +async def test_async_live_session_requires_api_key() -> None: + with pytest.raises(ValueError, match="api_key"): + await OpenAILiveSessions().async_live_session( + model="gpt-live-1", + websocket=_FakeClientWebSocket(messages=()), + logging_obj=cast(Logging, _FakeLogging()), + session_start={"type": "session.start", "session": {"model": "proxy-live-alias"}}, + api_base="https://api.openai.com/", + api_key=None, + ) + + +@pytest.mark.asyncio +async def test_async_live_session_does_not_log_success_on_handshake_refusal() -> None: + response: Final = Response(401, "Unauthorized", Headers()) + invalid_status: Final = InvalidStatus(response) + + def refuse_connection(url: str, **kwargs: object) -> _FailingConnection: + return _FailingConnection(invalid_status) + + client: Final = _FakeClientWebSocket(messages=()) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=refuse_connection).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + try: + assert worker.coroutine is None + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in logger.model_call_details + assert logger.dispatch_count == 0 + finally: + if worker.coroutine is not None: + await worker.coroutine + + +@pytest.mark.asyncio +async def test_async_live_session_does_not_log_success_when_connection_fails_before_connect() -> None: + def fail_connection(url: str, **kwargs: object) -> _FailingConnection: + return _FailingConnection(RuntimeError("connection failed")) + + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=fail_connection).async_live_session( + model="gpt-live-1", + websocket=_FakeClientWebSocket(messages=()), + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is None + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in logger.model_call_details + assert logger.dispatch_count == 0 + + +class _RefusingLoggingWorker: + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: + async_coroutine.close() + raise RuntimeError("logging queue unavailable") + + +@pytest.mark.asyncio +async def test_async_live_session_leaves_success_unmarked_when_enqueue_fails() -> None: + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + ) + client: Final = _FakeClientWebSocket(messages=()) + logger: Final = _FakeLogging() + + await OpenAILiveSessions( + logging_worker=_RefusingLoggingWorker(), websocket_connector=_FakeConnector(upstream) + ).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in logger.model_call_details + assert client.close_info is not None + assert client.close_info[0] == 1011 + + +@pytest.mark.parametrize( + "receive_error", + [None, RuntimeError('WebSocket is not connected. Need to call "accept" first.')], + ids=["clean-disconnect", "receive-runtime-error"], +) +@pytest.mark.asyncio +async def test_async_live_session_retains_final_usage_after_client_disconnect_and_bills_voice( + receive_error: Exception | None, +) -> None: + session_closed: Final = json.dumps({"type": "session.closed", "usage": {"seconds": 7}}) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + close_response=session_closed, + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket( + messages=(), + disconnect_after_send_type="session.started", + receive_error_after_close=receive_error, + ) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert logged_result.results == [ + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.closed", "usage": {"seconds": 7}}, + ], client.sent_text.qsize() + sent_frames: Final = tuple( + await asyncio.gather(*(upstream.sent_text.get() for _ in range(upstream.sent_text.qsize()))) + ) + session_close: Final = json.dumps({"type": "session.close"}) + assert tuple(frame for frame in sent_frames if _event_type(frame) == "session.close") == (session_close,) + + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + cost: Final = handle_realtime_stream_cost_calculation( + results=logged_result.results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + ) + assert cost == pytest.approx(7 * rate) + + +@pytest.mark.asyncio +async def test_async_live_session_drains_session_closed_after_upstream_send_failure() -> None: + session_closed: Final = json.dumps({"type": "session.closed", "usage": {"seconds": 4}}) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + close_response=session_closed, + fail_on_send_type="conversation.item.create", + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=('{"type":"conversation.item.create"}',)) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert logged_result.results == [ + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.closed", "usage": {"seconds": 4}}, + ] + assert client.close_info is not None + assert client.close_info[0] == 1000 + assert upstream.session_close_sent.is_set() + + +@pytest.mark.asyncio +async def test_async_live_session_raises_after_draining_usage_when_upstream_send_fails_without_close() -> None: + session_closed: Final = json.dumps({"type": "session.closed", "usage": {"seconds": 4}}) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + close_response=session_closed, + fail_on_send_type="conversation.item.create", + send_error=OSError("upstream socket write failed"), + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=('{"type":"conversation.item.create"}',)) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert logged_result.results[-1] == {"type": "session.closed", "usage": {"seconds": 4}} + client_frames: Final = tuple( + await asyncio.gather(*(client.sent_text.get() for _ in range(client.sent_text.qsize()))) + ) + assert tuple(_event_type(frame) for frame in client_frames) == ("session.started", "session.closed", "error") + assert client.close_info is not None + assert client.close_info[0] == 1011 + + +@pytest.mark.asyncio +async def test_async_live_session_does_not_log_success_when_session_start_send_fails() -> None: + upstream: Final = _FakeUpstream(client_frames=(), upstream_frames=(), fail_on_send_type="session.start") + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=()) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is None + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in logger.model_call_details + assert client.close_info is not None + assert client.close_info[0] == 1011 + + +@pytest.mark.asyncio +async def test_async_live_session_ends_when_upstream_keeps_socket_open_after_session_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(live_handler, "OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS", 0.01) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + json.dumps({"type": "session.closed", "usage": {"seconds": 3}}), + ), + never_respond_after_close=True, + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=()) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await asyncio.wait_for( + OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ), + timeout=2, + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert logged_result.results[-1] == {"type": "session.closed", "usage": {"seconds": 3}} + assert client.close_info is not None + assert client.close_info[0] == 1000 + assert upstream.closed + + +@pytest.mark.asyncio +async def test_async_live_session_closes_client_normally_when_upstream_drops_after_session_closed() -> None: + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + json.dumps({"type": "session.closed", "usage": {"seconds": 3}}), + ), + close_error=ConnectionClosedError(None, None), + ) + client: Final = _FakeClientWebSocket(messages=()) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await asyncio.wait_for( + OpenAILiveSessions(logging_worker=worker, websocket_connector=_FakeConnector(upstream)).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ), + timeout=2, + ) + + assert client.close_info == (1000, None) + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert logged_result.results[-1] == {"type": "session.closed", "usage": {"seconds": 3}} + + +@pytest.mark.asyncio +async def test_async_live_session_forwards_upstream_close_code_before_session_closed() -> None: + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + close_error=ConnectionClosedError(Close(code=4000, reason="invalid_request_error"), None), + ) + client: Final = _FakeClientWebSocket(messages=()) + + await asyncio.wait_for( + OpenAILiveSessions( + logging_worker=_FakeLoggingWorker(), websocket_connector=_FakeConnector(upstream) + ).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, _FakeLogging()), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ), + timeout=2, + ) + + assert client.close_info == (4000, "invalid_request_error") + + +@pytest.mark.parametrize( + ("first_message", "expected"), + [ + ( + '{"type":"session.start","session":{"model":"gpt-live-1","voice":"alloy"}}', + LiveSessionStart( + frame={"type": "session.start", "session": {"model": "gpt-live-1", "voice": "alloy"}}, + model="gpt-live-1", + ), + ), + ('{"type":"session.update","session":{"model":"gpt-live-1"}}', None), + ('{"type":"session.start","session":{"model":" "}}', None), + ('{"type":"session.start","session":"gpt-live-1"}', None), + ('{"type":"session.start"}', None), + ("not json", None), + ], +) +def test_parse_live_session_start(first_message: str, expected: LiveSessionStart | None) -> None: + parsed: Final = parse_live_session_start(first_message) + if expected is None: + assert parsed is None + return + assert parsed is not None + assert parsed.model == expected.model + assert dict(parsed.frame) == dict(expected.frame) + + +@pytest.mark.asyncio +async def test_async_live_session_bounds_total_time_when_upstream_keeps_streaming_after_session_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(live_handler, "OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS", 0.05) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + json.dumps({"type": "session.closed", "usage": {"seconds": 3}}), + json.dumps({"type": "session.usage.updated", "usage": {"seconds": 3}}), + ), + repeat_last_frame_forever=True, + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=()) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await asyncio.wait_for( + OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ), + timeout=1, + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert {"type": "session.closed", "usage": {"seconds": 3}} in logged_result.results + assert client.close_info is not None + assert client.close_info[0] == 1000 + assert upstream.closed + + +@pytest.mark.asyncio +async def test_async_live_session_ends_when_client_stops_reading_after_session_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(live_handler, "OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS", 0.05) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + json.dumps({"type": "session.closed", "usage": {"seconds": 3}}), + ), + repeat_last_frame_forever=True, + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=(), block_on_send_type="session.closed") + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await asyncio.wait_for( + OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ), + timeout=1, + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert {"type": "session.closed", "usage": {"seconds": 3}} in logged_result.results + assert client.close_info is not None + assert client.close_info[0] == 1000 + assert upstream.closed + + +@pytest.mark.asyncio +async def test_async_live_session_keeps_finite_usage_after_invalid_snapshot() -> None: + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + '{"type":"session.usage.updated","usage":{"seconds":7}}', + '{"type":"session.usage.updated","usage":{"seconds":"nan"}}', + ), + never_respond_after_close=True, + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=(), disconnect_after_send_type="session.started") + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert logged_result.results == [ + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 7}}, + ] + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + cost: Final = handle_realtime_stream_cost_calculation( + results=logged_result.results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + ) + assert cost == pytest.approx(7 * rate) + + +@pytest.mark.asyncio +async def test_async_live_session_does_not_send_a_second_close_after_client_close( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(live_handler, "OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS", 0.01) + client_frames: Final = (json.dumps({"type": "session.close"}),) + upstream: Final = _FakeUpstream( + client_frames=client_frames, + upstream_frames=(), + never_respond_after_close=True, + ) + connector: Final = _FakeConnector(upstream) + worker: Final = _FakeLoggingWorker() + logger: Final = _FakeLogging() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=_FakeClientWebSocket( + messages=client_frames, + disconnect_after_receive_type="session.close", + ), + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + sent_frames: Final = tuple( + await asyncio.gather(*(upstream.sent_text.get() for _ in range(upstream.sent_text.qsize()))) + ) + assert tuple(frame for frame in sent_frames if _event_type(frame) == "session.close") == client_frames + assert worker.coroutine is not None + await worker.coroutine + + +@pytest.mark.asyncio +async def test_async_live_session_bounds_a_blocking_upstream_close_send( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(live_handler, "OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS", 0.05) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + never_respond_after_close=True, + block_session_close_send=True, + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=(), disconnect_after_send_type="session.started") + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await asyncio.wait_for( + OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ), + timeout=1, + ) + + assert worker.coroutine is not None + await worker.coroutine + assert logger.dispatch_count == 1 + + +@pytest.mark.asyncio +async def test_async_live_session_bounds_a_blocking_upstream_close( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(live_handler, "OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS", 0.05) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + close_response='{"type":"session.closed","usage":{"seconds":7}}', + block_close=True, + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=(), disconnect_after_send_type="session.started") + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await asyncio.wait_for( + OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ), + timeout=1, + ) + + assert worker.coroutine is not None + await worker.coroutine + assert logger.dispatch_count == 1 + + +@pytest.mark.asyncio +async def test_async_live_session_retains_closed_event_when_client_send_fails() -> None: + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + '{"type":"session.closed","usage":{"seconds":7}}', + ), + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=(), fail_on_send_type="session.closed") + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert {"type": "session.closed", "usage": {"seconds": 7}} in logged_result.results + + +@pytest.mark.asyncio +async def test_async_live_session_bounds_wait_for_usage_after_client_disconnect( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(live_handler, "OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS", 0.01) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + never_respond_after_close=True, + ) + connector: Final = _FakeConnector(upstream) + client: Final = _FakeClientWebSocket(messages=(), disconnect_after_send_type="session.started") + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await asyncio.wait_for( + OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ), + timeout=1, + ) + + assert upstream.closed + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + assert all(event["type"] != "session.closed" for event in logged_result.results) + + +@pytest.mark.asyncio +async def test_async_live_session_compacts_and_deduplicates_terminal_response_events() -> None: + large_output: Final = "x" * 4096 + large_instructions: Final = "y" * 4096 + response_usage: Final[dict[str, object]] = {"input_tokens": 10, "output_tokens": 5} + duplicate_response_usage: Final[dict[str, object]] = {"input_tokens": 100, "output_tokens": 50} + first_response: Final = json.dumps( + { + "type": "response.event", + "event": { + "type": "response.completed", + "response": { + "id": "resp-1", + "model": "gpt-live-1", + "usage": response_usage, + "output": large_output, + "instructions": large_instructions, + }, + }, + } + ) + duplicate_response: Final = json.dumps( + { + "type": "response.event", + "event": { + "type": "response.completed", + "response": { + "id": "resp-1", + "model": "duplicate-model", + "usage": duplicate_response_usage, + "output": large_output, + "instructions": large_instructions, + }, + }, + } + ) + response_without_id: Final = json.dumps( + { + "type": "response.event", + "event": { + "type": "response.completed", + "response": {"model": "gpt-live-1", "usage": response_usage}, + }, + } + ) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + first_response, + duplicate_response, + response_without_id, + ), + ) + connector: Final = _FakeConnector(upstream) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=connector).async_live_session( + model="gpt-live-1", + websocket=_FakeClientWebSocket(messages=()), + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + response_events: Final = tuple(event for event in logged_result.results if event["type"] == "response.event") + assert response_events == ( + { + "type": "response.event", + "event": { + "type": "response.completed", + "response": { + "id": "resp-1", + "model": "gpt-live-1", + "usage": response_usage, + }, + }, + }, + ) + + +@pytest.mark.asyncio +async def test_async_live_session_bills_a_later_response_event_with_usage_over_an_earlier_one_without() -> None: + response_usage: Final[dict[str, object]] = {"input_tokens": 10, "output_tokens": 5} + incomplete_without_usage: Final = json.dumps( + { + "type": "response.event", + "event": {"type": "response.incomplete", "response": {"id": "resp-1", "model": "gpt-live-1"}}, + } + ) + completed_with_usage: Final = json.dumps( + { + "type": "response.event", + "event": { + "type": "response.completed", + "response": {"id": "resp-1", "model": "gpt-live-1", "usage": response_usage}, + }, + } + ) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=( + '{"type":"session.started","session":{"model":"gpt-live-1"}}', + incomplete_without_usage, + completed_with_usage, + ), + ) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker, websocket_connector=_FakeConnector(upstream)).async_live_session( + model="gpt-live-1", + websocket=_FakeClientWebSocket(messages=()), + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + logged_result: Final = cast(LiteLLMRealtimeStreamLoggingObject, logger.dispatched_result) + response_events: Final = tuple(event for event in logged_result.results if event["type"] == "response.event") + assert response_events == ( + { + "type": "response.event", + "event": { + "type": "response.completed", + "response": {"id": "resp-1", "model": "gpt-live-1", "usage": response_usage}, + }, + }, + ), logged_result.results + + +@pytest.mark.asyncio +async def test_async_live_session_logs_the_relay_end_time_instead_of_the_logging_worker_time() -> None: + relay_end_time: Final = datetime(2026, 1, 1, 12, 0, 5) + upstream: Final = _FakeUpstream( + client_frames=(), + upstream_frames=('{"type":"session.started","session":{"model":"gpt-live-1"}}',), + ) + logger: Final = _FakeLogging() + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions( + logging_worker=worker, + websocket_connector=_FakeConnector(upstream), + clock=lambda: relay_end_time, + ).async_live_session( + model="gpt-live-1", + websocket=_FakeClientWebSocket(messages=()), + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "gpt-live-1"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + assert worker.coroutine is not None + await worker.coroutine + assert logger.dispatched_end_time == relay_end_time + + +@pytest.mark.asyncio +async def test_async_live_session_reports_redacted_server_error_and_bounds_close_reason() -> None: + error_message: Final = "x" * 300 + client: Final = _FakeClientWebSocket(messages=()) + logger: Final = _FakeLogging(pre_call_error=RuntimeError(error_message)) + worker: Final = _FakeLoggingWorker() + + await OpenAILiveSessions(logging_worker=worker).async_live_session( + model="gpt-live-1", + websocket=client, + logging_obj=cast(Logging, logger), + session_start={"type": "session.start", "session": {"model": "proxy-live-alias"}}, + api_base="https://api.openai.com/", + api_key="live-test-key", + ) + + error_event: Final = json.loads(client.sent_text.get_nowait()) + assert error_event == {"type": "error", "error": {"type": "server_error", "message": error_message}} + assert client.close_info is not None + assert client.close_info[0] == 1011 + assert client.close_info[1] is not None + assert len(client.close_info[1].encode("utf-8")) <= 123 + assert client.close_info[1] == "x" * 123 + try: + assert worker.coroutine is None + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in logger.model_call_details + assert logger.dispatch_count == 0 + finally: + if worker.coroutine is not None: + await worker.coroutine diff --git a/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py b/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py index eee985f0aca..fb676ae9a61 100644 --- a/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py +++ b/tests/unit/proxy/agent_endpoints/auth/test_managed_authorization.py @@ -451,6 +451,20 @@ def test_managed_realtime_requires_a_model_and_ignores_completion_defaults(route ) +@pytest.mark.parametrize("route", ["/live/sessions", "/v1/live/sessions", "/openai/v1/live/sessions"]) +def test_managed_live_upgrade_defers_the_model_to_the_session_start_frame(route: str) -> None: + from litellm.proxy.agent_endpoints.auth.managed_authorization import ( + managed_agent_route_allowed, + managed_inference_request, + ) + + assert managed_agent_route_allowed(route, "GET") + assert managed_inference_request(route, {}, {"completion_model": "allowed-default"}, "cli") == {} + assert managed_inference_request( + route, {"model": "frame-model"}, {"completion_model": "allowed-default"}, "cli" + ) == {"model": "frame-model"} + + @pytest.mark.parametrize("mode,user", [("autonomous", None), ("delegated", "verified-human")]) def test_matching_identity_revision_and_execution_mode_pass_admission(mode: str, user: str | None) -> None: context: Final = ManagedAgentContext.model_validate( diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index da8b25c2cee..3660019cc2c 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -12,6 +12,7 @@ import pytest from starlette.datastructures import URL from starlette.types import Message from litellm._logging import verbose_proxy_logger +import json import logging import litellm from litellm.proxy.auth.user_api_key_auth import ( @@ -912,6 +913,54 @@ async def test_user_api_key_auth_websocket_carries_asgi_path(): assert request_arg.scope.get("root_path") == "" +@pytest.mark.asyncio +async def test_user_api_key_auth_live_websocket_ignores_query_model(): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_live_websocket + + mock_websocket = MagicMock(spec=WebSocket) + mock_websocket.query_params = {"model": "gpt-realtime"} + mock_websocket.headers = {"authorization": "Bearer some_api_key"} + mock_websocket.scope = {"headers": [(b"authorization", b"Bearer some_api_key")]} + mock_websocket.url = URL(url="/v1/live/sessions?model=gpt-realtime") + + with patch("litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True) as mock_user_api_key_auth: + await user_api_key_auth_live_websocket(mock_websocket) + + request_arg = mock_user_api_key_auth.call_args.kwargs["request"] + assert json.loads(await request_arg.body()) == {"model": ""} + mock_user_api_key_auth.assert_called_once() + + +@pytest.mark.asyncio +async def test_accepted_websocket_reauthorizes_the_frame_model_without_closing(): + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_accepted_websocket_for_model + + mock_websocket = MagicMock(spec=WebSocket) + mock_websocket.query_params = {} + mock_websocket.headers = {"authorization": "Bearer some_api_key"} + mock_websocket.scope = { + "type": "websocket", + "path": "/v1/live/sessions", + "headers": [(b"authorization", b"Bearer some_api_key")], + } + mock_websocket.url = URL(url="/v1/live/sessions") + denial = ProxyException("denied", "user_model_access_denied", "model", 403) + + with patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True, side_effect=denial + ) as mock_user_api_key_auth: + with pytest.raises(ProxyException) as raised: + await user_api_key_auth_accepted_websocket_for_model(mock_websocket, "gpt-live-1") + + assert raised.value is denial + request_arg = mock_user_api_key_auth.call_args.kwargs["request"] + assert json.loads(await request_arg.body()) == {"model": "gpt-live-1"} + assert request_arg.scope.get("path") == "/v1/live/sessions" + assert mock_user_api_key_auth.call_args.kwargs["api_key"] == "Bearer some_api_key" + mock_websocket.close.assert_not_called() + + @pytest.mark.parametrize("enforce_rbac", [True, False]) @pytest.mark.asyncio async def test_jwt_user_api_key_auth_builder_enforce_rbac(enforce_rbac, monkeypatch): diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 79f4edd4d62..80bc3ad1a34 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -1,9 +1,9 @@ import os import sys import traceback -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping from contextlib import asynccontextmanager -from typing import Final, Literal +from typing import Final, Literal, NoReturn, cast from unittest import mock from dotenv import load_dotenv @@ -35,7 +35,7 @@ logging.basicConfig( from unittest.mock import AsyncMock, MagicMock, patch -from fastapi import FastAPI, HTTPException +from fastapi import FastAPI, HTTPException, WebSocket # test /chat/completion request to the proxy from fastapi.testclient import TestClient @@ -3296,3 +3296,496 @@ def test_get_litellm_model_info(data): ): get_litellm_model_info(model=model) get_info_mock.assert_called_once_with(data["expected"]) + + +def test_live_session_accept_failure_releases_reservations( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + + released: Final[list[str]] = [] + + async def release_budget(user_api_key_dict: object) -> None: + released.append("budget") + + async def release_slot(user_api_key_dict: object) -> None: + released.append("slot") + + async def receive() -> dict[str, object]: + return {"type": "websocket.connect"} + + async def send_failing_accept(message: Mapping[str, object]) -> None: + if message["type"] == "websocket.accept": + raise RuntimeError("accept failed") + + scope: Final = { + "type": "websocket", + "path": "/v1/live/sessions", + "raw_path": b"/v1/live/sessions", + "root_path": "", + "scheme": "ws", + "query_string": b"", + "headers": [(b"host", b"testserver"), (b"authorization", b"Bearer test-api-key")], + "client": ("testclient", 50000), + "server": ("testserver", 80), + "subprotocols": [], + } + route_call: Final = AsyncMock() + monkeypatch.setattr(proxy_server, "route_request", route_call) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", release_slot) + + with pytest.raises(RuntimeError, match="accept failed"): + asyncio.run(client_no_auth.app(scope, receive, send_failing_accept)) + + assert released == ["budget", "slot"] + route_call.assert_not_awaited() + + +def _override_live_session_authorizer( + monkeypatch: pytest.MonkeyPatch, + authorizer: Callable[[WebSocket, str], Awaitable[UserAPIKeyAuth]], +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setitem( + proxy_server.app.dependency_overrides, + proxy_server.live_session_model_authorizer, + lambda: authorizer, + ) + + +def _start_live_session_and_wait_for_reply(client: TestClient) -> None: + with client.websocket_connect("/v1/live/sessions", headers={"Authorization": "Bearer test-api-key"}) as websocket: + websocket.send_json({"type": "session.start", "session": {"model": "gpt-3.5-turbo"}}) + websocket.receive_text() + + +def test_live_session_model_check_error_releases_reservations( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + + released: Final[list[str]] = [] + + async def release_budget(user_api_key_dict: object) -> None: + released.append("budget") + + async def release_slot(user_api_key_dict: object) -> None: + released.append("slot") + + async def failing_authorizer(websocket: object, model: str) -> UserAPIKeyAuth: + raise RuntimeError("agent lookup failed") + + route_call: Final = AsyncMock() + monkeypatch.setattr(proxy_server, "route_request", route_call) + _override_live_session_authorizer(monkeypatch, failing_authorizer) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", release_slot) + + with pytest.raises(RuntimeError, match="agent lookup failed"): + _start_live_session_and_wait_for_reply(client_no_auth) + + assert released == ["budget", "slot"] + route_call.assert_not_awaited() + + +class _PreCallAborted(BaseException): + pass + + +def test_live_session_pre_call_abort_releases_reservations( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + frame_auth: Final = UserAPIKeyAuth(token="frame-authorized") + released_for_frame: Final[list[str]] = [] + + async def release_budget(user_api_key_dict: UserAPIKeyAuth) -> None: + if user_api_key_dict is frame_auth: + released_for_frame.append("budget") + + async def release_slot(user_api_key_dict: UserAPIKeyAuth) -> None: + if user_api_key_dict is frame_auth: + released_for_frame.append("slot") + + async def authorize(websocket: object, model: str) -> UserAPIKeyAuth: + return frame_auth + + class _AbortingPreCallProcessing(ProxyBaseLLMRequestProcessing): + async def common_processing_pre_call_logic(self, *args: object, **kwargs: object) -> NoReturn: + raise _PreCallAborted + + route_call: Final = AsyncMock() + monkeypatch.setattr(proxy_server, "route_request", route_call) + _override_live_session_authorizer(monkeypatch, authorize) + monkeypatch.setattr(proxy_server, "ProxyBaseLLMRequestProcessing", _AbortingPreCallProcessing) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", release_slot) + + with pytest.raises(_PreCallAborted): + _start_live_session_and_wait_for_reply(client_no_auth) + + assert released_for_frame == ["budget", "slot"] + route_call.assert_not_awaited() + + +@pytest.mark.parametrize("path", ["/openai/v1/live/sessions", "/v1/live/sessions", "/live/sessions"]) +def test_live_session_invalid_start_is_rejected_without_routing( + path: str, + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + from starlette.websockets import WebSocketDisconnect + + async def no_op(user_api_key_dict: object) -> None: + return None + + route_call: Final = AsyncMock() + monkeypatch.setattr(proxy_server, "route_request", route_call) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op) + + with client_no_auth.websocket_connect(path, headers={"Authorization": "Bearer test-api-key"}) as websocket: + websocket.send_text('{"type":"session.start","session":{"model":""}}') + assert websocket.receive_json() == { + "type": "error", + "error": { + "type": "invalid_request_error", + "code": "invalid_session_start", + "message": "First message must be a session.start JSON object with a non-empty session.model.", + }, + } + with pytest.raises(WebSocketDisconnect) as error: + websocket.receive_json() + assert error.value.code == 1008 + + route_call.assert_not_awaited() + + +@pytest.mark.parametrize( + ("denial_type", "expected_code"), + (("user_model_access_denied", "model_not_allowed"), ("budget_exceeded", "session_not_allowed")), +) +def test_live_session_frame_authorization_denial_closes_with_policy_code( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, + denial_type: str, + expected_code: str, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._types import ProxyException + from starlette.websockets import WebSocketDisconnect + + async def no_op(user_api_key_dict: object) -> None: + return None + + authorized_models: Final[list[str]] = [] + + async def deny(websocket: object, model: str) -> UserAPIKeyAuth: + authorized_models.append(model) + raise ProxyException("model access denied for this user", denial_type, "model", 403) + + route_call: Final = AsyncMock() + monkeypatch.setattr(proxy_server, "route_request", route_call) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op) + _override_live_session_authorizer(monkeypatch, deny) + + with client_no_auth.websocket_connect( + "/v1/live/sessions", + headers={"Authorization": "Bearer test-api-key"}, + ) as websocket: + websocket.send_json({"type": "session.start", "session": {"model": "gpt-4o"}}) + assert websocket.receive_json() == { + "type": "error", + "error": { + "type": "invalid_request_error", + "code": expected_code, + "message": "model access denied for this user", + }, + } + with pytest.raises(WebSocketDisconnect) as error: + websocket.receive_json() + assert error.value.code == 1008 + + assert authorized_models == ["gpt-4o"] + route_call.assert_not_awaited() + + +def test_live_session_valid_start_routes_with_model_and_frame( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from types import SimpleNamespace + + from litellm.proxy import proxy_server + from starlette.websockets import WebSocket + + async def no_op(user_api_key_dict: object) -> None: + return None + + async def complete_call() -> None: + return None + + captured: Final = SimpleNamespace(data=None, route_type=None) + + async def route( + *, + data: dict[str, object], + route_type: str, + llm_router: object, + user_model: str | None, + ) -> Coroutine[object, object, None]: + captured.data = data + captured.route_type = route_type + await cast(WebSocket, data["websocket"]).send_text("routed") + return complete_call() + + frame_auth: Final = UserAPIKeyAuth(token="frame-authorized") + + async def authorize(websocket: object, model: str) -> UserAPIKeyAuth: + return frame_auth + + monkeypatch.setattr(proxy_server, "route_request", route) + _override_live_session_authorizer(monkeypatch, authorize) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op) + start_frame: Final[dict[str, object]] = { + "type": "session.start", + "session": {"model": "gpt-3.5-turbo", "instructions": "Stay brief"}, + } + + with client_no_auth.websocket_connect( + "/v1/live/sessions", + headers={"Authorization": "Bearer test-api-key"}, + ) as websocket: + websocket.send_json(start_frame) + assert websocket.receive_text() == "routed" + + routed_data: Final = cast(dict[str, object], captured.data) + assert captured.route_type == "_arealtime" + assert routed_data["model"] == "gpt-3.5-turbo" + assert routed_data["live_session_start"] == start_frame + assert routed_data["user_api_key_dict"] is frame_auth + + +def test_live_session_ignores_query_model_and_authorizes_frame_model( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastapi import Request + + from litellm.proxy import proxy_server + from starlette.websockets import WebSocket + + async def no_op(user_api_key_dict: object) -> None: + return None + + async def complete_call() -> None: + return None + + async def route( + *, + data: dict[str, object], + route_type: str, + llm_router: object, + user_model: str | None, + ) -> Coroutine[object, object, None]: + await cast(WebSocket, data["websocket"]).send_text("routed") + return complete_call() + + upgrade_auth: Final = UserAPIKeyAuth(token="upgrade", models=["gpt-live-1"]) + + async def authenticate(request: Request, api_key: str) -> UserAPIKeyAuth: + if await request.body() != b'{"model": ""}': + raise HTTPException(status_code=403, detail="Model is not allowed") + return upgrade_auth + + auth_call: Final = AsyncMock(side_effect=authenticate) + monkeypatch.setattr("litellm.proxy.auth.user_api_key_auth.user_api_key_auth", auth_call) + frame_auth: Final = UserAPIKeyAuth(token="frame-authorized", models=["gpt-live-1"]) + authorize: Final = AsyncMock(return_value=frame_auth) + monkeypatch.setattr(proxy_server, "route_request", route) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op) + _override_live_session_authorizer(monkeypatch, authorize) + + with client_no_auth.websocket_connect( + "/v1/live/sessions?model=gpt-realtime", + headers={"Authorization": "Bearer test-api-key"}, + ) as websocket: + websocket.send_json({"type": "session.start", "session": {"model": "gpt-live-1"}}) + assert websocket.receive_text() == "routed" + + request_arg: Final = auth_call.await_args.kwargs["request"] + assert json.loads(asyncio.run(request_arg.body())) == {"model": ""} + assert authorize.await_args.args[1] == "gpt-live-1" + + +def test_live_session_reauthorized_token_still_needs_the_model_on_the_key( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from starlette.websockets import WebSocketDisconnect + + from litellm.proxy import proxy_server + + released_tokens: Final[list[str | None]] = [] + + async def release_budget(user_api_key_dict: UserAPIKeyAuth) -> None: + released_tokens.append(user_api_key_dict.token) + + async def no_op(user_api_key_dict: object) -> None: + return None + + async def custom_auth_without_common_checks(websocket: object, model: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth(token="custom-auth", models=["gpt-4o"]) + + route_call: Final = AsyncMock() + monkeypatch.setattr(proxy_server, "route_request", route_call) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op) + _override_live_session_authorizer(monkeypatch, custom_auth_without_common_checks) + + with client_no_auth.websocket_connect( + "/v1/live/sessions", + headers={"Authorization": "Bearer test-api-key"}, + ) as websocket: + websocket.send_json({"type": "session.start", "session": {"model": "gpt-3.5-turbo"}}) + error_event: Final = websocket.receive_json() + with pytest.raises(WebSocketDisconnect) as error: + websocket.receive_json() + + assert error_event["error"]["code"] == "model_not_allowed", error_event + assert error.value.code == 1008 + assert "custom-auth" in released_tokens + route_call.assert_not_awaited() + + +def test_live_session_releases_the_upgrade_reservation_before_the_frame_reserves( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from starlette.websockets import WebSocket + + from litellm.proxy import proxy_server + + released_reservations: Final[list[object]] = [] + + async def release_budget(user_api_key_dict: UserAPIKeyAuth) -> None: + released_reservations.append(user_api_key_dict.budget_reservation) + + async def no_op(user_api_key_dict: object) -> None: + return None + + async def complete_call() -> None: + return None + + async def route( + *, + data: dict[str, object], + route_type: str, + llm_router: object, + user_model: str | None, + ) -> Coroutine[object, object, None]: + await cast(WebSocket, data["websocket"]).send_text("routed") + return complete_call() + + frame_auth: Final = UserAPIKeyAuth(token="frame-authorized", budget_reservation={"reserved_for": "frame model"}) + + async def upgrade_auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth(token="upgrade", budget_reservation={"reserved_for": "query model"}) + + released_before_frame_auth: Final[list[tuple[object, ...]]] = [] + + async def authorize(websocket: object, model: str) -> UserAPIKeyAuth: + released_before_frame_auth.append(tuple(released_reservations)) + return frame_auth + + monkeypatch.setattr(proxy_server, "route_request", route) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op) + monkeypatch.setitem(proxy_server.app.dependency_overrides, proxy_server.user_api_key_auth_live_websocket, upgrade_auth) + _override_live_session_authorizer(monkeypatch, authorize) + + with client_no_auth.websocket_connect( + "/v1/live/sessions?model=gpt-4o", + headers={"Authorization": "Bearer test-api-key"}, + ) as websocket: + websocket.send_json({"type": "session.start", "session": {"model": "gpt-3.5-turbo"}}) + assert websocket.receive_text() == "routed" + + assert released_before_frame_auth == [({"reserved_for": "query model"},)], released_reservations + + +def test_live_session_preserves_pretty_printed_start_frame( + client_no_auth: TestClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from types import SimpleNamespace + + from starlette.websockets import WebSocket + + from litellm.proxy import proxy_server + + async def no_op(user_api_key_dict: object) -> None: + return None + + async def complete_call() -> None: + return None + + captured: Final = SimpleNamespace(data=None, route_type=None) + + async def route( + *, + data: dict[str, object], + route_type: str, + llm_router: object, + user_model: str | None, + ) -> Coroutine[object, object, None]: + captured.data = data + captured.route_type = route_type + await cast(WebSocket, data["websocket"]).send_text("routed") + return complete_call() + + frame_auth: Final = UserAPIKeyAuth(token="frame-authorized") + + async def authorize(websocket: object, model: str) -> UserAPIKeyAuth: + return frame_auth + + monkeypatch.setattr(proxy_server, "route_request", route) + _override_live_session_authorizer(monkeypatch, authorize) + monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op) + monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op) + + start_frame: Final = "\r\n".join( + [ + "{", + ' "type": "session.start",', + ' "session": {"model": "gpt-live-1", "instructions": "a\\nb"}', + "}", + ] + ) + expected_start: Final[dict[str, object]] = { + "type": "session.start", + "session": {"model": "gpt-live-1", "instructions": "a\nb"}, + } + + with client_no_auth.websocket_connect( + "/v1/live/sessions", + headers={"Authorization": "Bearer test-api-key"}, + ) as websocket: + websocket.send_text(start_frame) + assert websocket.receive_text() == "routed" + + routed_data: Final = cast(dict[str, object], captured.data) + assert captured.route_type == "_arealtime" + assert routed_data["model"] == "gpt-live-1" + assert routed_data["live_session_start"] == expected_start diff --git a/tests/unit/realtime_api/test_main.py b/tests/unit/realtime_api/test_main.py index 9d11d2191fa..52433665841 100644 --- a/tests/unit/realtime_api/test_main.py +++ b/tests/unit/realtime_api/test_main.py @@ -1,6 +1,7 @@ import asyncio import json import time +from collections.abc import Mapping from types import TracebackType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -33,6 +34,11 @@ class FakeLogging: pass +class _CallCapture: + def __init__(self) -> None: + self.values: Mapping[str, object] | None = None + + def test_resolves_top_level_session_model(): resolved = _with_resolved_session_model({"model": "alias/gpt-realtime"}, "gpt-realtime") assert resolved == {"model": "gpt-realtime"} @@ -238,6 +244,191 @@ def test_client_secret_forwards_nested_transcription_model_untouched(monkeypatch assert session["input_audio_transcription"]["model"] == "whisper-1" +@pytest.mark.asyncio +async def test_arealtime_live_session_dispatches_to_openai_live_handler(monkeypatch: pytest.MonkeyPatch) -> None: + captured: Final = _CallCapture() + + def mock_get_llm_provider( + model: str, + api_base: str | None, + api_key: str | None, + ) -> tuple[str, str, str | None, str | None]: + return "gpt-live-1", "openai", "resolved-key", "https://live-gateway.example/v1" + + class FakeLiveHandler: + async def async_live_session(self, **kwargs: object) -> None: + captured.values = kwargs + + monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) + monkeypatch.setattr(realtime_main, "openai_live_sessions", FakeLiveHandler()) + start_frame: Final[dict[str, object]] = { + "type": "session.start", + "session": {"model": "live-alias", "instructions": "Stay brief"}, + } + websocket: Final = MagicMock() + logging_obj: Final = FakeLogging() + + await realtime_main._arealtime.__wrapped__( + model="live-alias", + websocket=websocket, + litellm_logging_obj=logging_obj, + live_session_start=start_frame, + ) + + assert captured.values is not None + assert captured.values["model"] == "gpt-live-1" + assert captured.values["websocket"] is websocket + assert captured.values["logging_obj"] is logging_obj + assert captured.values["session_start"] == start_frame + assert captured.values["api_base"] == "https://live-gateway.example/v1" + assert captured.values["api_key"] == "resolved-key" + + +@pytest.mark.asyncio +async def test_arealtime_live_session_rejects_matching_guardrail_before_dispatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class RecordingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__( + guardrail_name="live-blocker", + event_hook=GuardrailEventHooks.realtime_input_transcription, + default_on=True, + ) + self.request_data: dict[str, object] | None = None + + def should_run_guardrail(self, data: dict[str, object], event_type: GuardrailEventHooks) -> bool: + self.request_data = data + return super().should_run_guardrail(data, event_type) + + class FakeLiveHandler: + def __init__(self) -> None: + self.called = False + + async def async_live_session(self, **kwargs: object) -> None: + self.called = True + + guardrail: Final = RecordingGuardrail() + live_handler: Final = FakeLiveHandler() + + def mock_get_llm_provider( + model: str, + api_base: str | None, + api_key: str | None, + ) -> tuple[str, str, str | None, str | None]: + return "gpt-live-1", "openai", "resolved-key", "https://live-gateway.example/v1" + + monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) + monkeypatch.setattr(realtime_main, "openai_live_sessions", live_handler) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + with pytest.raises(ValueError, match="Guardrails are not supported on OpenAI Live sessions"): + await realtime_main._arealtime.__wrapped__( + model="live-alias", + websocket=MagicMock(), + litellm_logging_obj=FakeLogging(), + litellm_metadata={"session_id": "live-session"}, + live_session_start={"type": "session.start", "session": {"model": "live-alias"}}, + ) + + assert guardrail.request_data == {"litellm_metadata": {"session_id": "live-session"}, "stream": True} + assert live_handler.called is False + + +@pytest.mark.asyncio +async def test_arealtime_live_session_rejects_streaming_only_pre_call_guardrail( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class FakeLiveHandler: + def __init__(self) -> None: + self.called = False + + async def async_live_session(self, **kwargs: object) -> None: + self.called = True + + guardrail: Final = CustomGuardrail( + guardrail_name="streaming-only-blocker", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + stream_scope="streaming", + ) + live_handler: Final = FakeLiveHandler() + + def mock_get_llm_provider( + model: str, + api_base: str | None, + api_key: str | None, + ) -> tuple[str, str, str | None, str | None]: + return "gpt-live-1", "openai", "resolved-key", "https://live-gateway.example/v1" + + monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) + monkeypatch.setattr(realtime_main, "openai_live_sessions", live_handler) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + with pytest.raises(ValueError, match="Guardrails are not supported on OpenAI Live sessions"): + await realtime_main._arealtime.__wrapped__( + model="live-alias", + websocket=MagicMock(), + litellm_logging_obj=FakeLogging(), + litellm_metadata={"session_id": "live-session"}, + live_session_start={"type": "session.start", "session": {"model": "live-alias"}}, + ) + + assert live_handler.called is False + + +@pytest.mark.asyncio +async def test_arealtime_live_session_rejects_non_openai_provider(monkeypatch: pytest.MonkeyPatch) -> None: + def mock_get_llm_provider( + model: str, + api_base: str | None, + api_key: str | None, + ) -> tuple[str, str, str | None, str | None]: + return "claude-live", "anthropic", None, None + + monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) + + with pytest.raises(ValueError, match="anthropic"): + await realtime_main._arealtime.__wrapped__( + model="anthropic/claude-live", + websocket=MagicMock(), + litellm_logging_obj=FakeLogging(), + live_session_start={"type": "session.start", "session": {"model": "claude-live"}}, + ) + + +@pytest.mark.asyncio +async def test_arealtime_without_live_start_preserves_openai_realtime_dispatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: Final = _CallCapture() + + def mock_get_llm_provider( + model: str, + api_base: str | None, + api_key: str | None, + ) -> tuple[str, str, str | None, str | None]: + return "gpt-realtime", "openai", "resolved-key", "https://realtime-gateway.example/v1" + + class FakeRealtimeHandler: + async def async_realtime(self, **kwargs: object) -> None: + captured.values = kwargs + + monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) + monkeypatch.setattr(realtime_main, "openai_realtime", FakeRealtimeHandler()) + + await realtime_main._arealtime.__wrapped__( + model="gpt-realtime", + websocket=MagicMock(), + litellm_logging_obj=FakeLogging(), + ) + + assert captured.values is not None + assert captured.values["model"] == "gpt-realtime" + assert captured.values["api_key"] == "resolved-key" + assert captured.values["api_base"] == "https://realtime-gateway.example/v1" + + class _CapturingConnect: def __init__(self) -> None: self.url: str | None = None diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 5d5d53d793e..93f2352db60 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -8690,3 +8690,264 @@ def test_batch_cost_calculator(): cost = completion_cost(**args) assert cost > 0 + + +def _live_session_results(*events: dict[str, object]) -> OpenAIRealtimeStreamList: + return cast(OpenAIRealtimeStreamList, list(events)) + + +def _live_logging_object( + start_time: datetime.datetime, + end_time: datetime.datetime, +) -> Logging: + return cast( + Logging, + SimpleNamespace(model_call_details={"start_time": start_time, "end_time": end_time}), + ) + + +def test_realtime_keeps_first_reported_session_model_for_pricing(_local_model_cost_map: None) -> None: + first_model: Final = "gpt-realtime-2" + later_model: Final = "gpt-realtime-mini" + combined_usage: Final = Usage(prompt_tokens=120, completion_tokens=60, total_tokens=180) + first_model_cost: Final = handle_realtime_stream_cost_calculation( + results=_live_session_results({"type": "session.created", "session": {"model": first_model}}), + combined_usage_object=combined_usage, + custom_llm_provider="openai", + litellm_model_name="unmapped-realtime-alias", + ) + first_model_with_later_update_cost: Final = handle_realtime_stream_cost_calculation( + results=_live_session_results( + {"type": "session.created", "session": {"model": first_model}}, + {"type": "session.started", "session": {"model": later_model}}, + ), + combined_usage_object=combined_usage, + custom_llm_provider="openai", + litellm_model_name="unmapped-realtime-alias", + ) + later_model_cost: Final = handle_realtime_stream_cost_calculation( + results=_live_session_results({"type": "session.started", "session": {"model": later_model}}), + combined_usage_object=combined_usage, + custom_llm_provider="openai", + litellm_model_name="unmapped-realtime-alias", + ) + + assert first_model_with_later_update_cost == pytest.approx(first_model_cost) + assert first_model_cost != pytest.approx(later_model_cost) + + +def test_live_session_charges_latest_closed_voice_usage(_local_model_cost_map: None) -> None: + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 12}}, + {"type": "session.usage.updated", "usage": {"seconds": 15}}, + {"type": "session.closed", "usage": {"seconds": 16}}, + ) + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + ) + + assert rate > 0 + assert cost == pytest.approx(16 * rate) + + +def test_live_session_uses_last_reported_duration_when_wall_clock_is_short(_local_model_cost_map: None) -> None: + start_time: Final = datetime.datetime(2025, 1, 1) + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 15}}, + ) + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + litellm_logging_obj=_live_logging_object(start_time, start_time + datetime.timedelta(seconds=5)), + ) + + assert cost == pytest.approx(15 * rate) + + +def test_live_session_uses_latest_cumulative_usage_snapshot_not_sum(_local_model_cost_map: None) -> None: + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 12}}, + {"type": "session.usage.updated", "usage": {"seconds": 15}}, + ) + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + ) + + assert cost == pytest.approx(15 * rate) + + +def test_live_session_ignores_non_finite_usage_snapshot_after_finite_one(_local_model_cost_map: None) -> None: + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 15}}, + {"type": "session.usage.updated", "usage": {"seconds": float("nan")}}, + ) + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + ) + + assert cost == pytest.approx(15 * rate) + + +def test_live_session_ignores_non_finite_closed_usage_snapshot(_local_model_cost_map: None) -> None: + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 15}}, + {"type": "session.closed", "usage": {"seconds": float("inf")}}, + ) + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + ) + + assert cost == pytest.approx(15 * rate) + + +def test_live_session_skips_non_finite_rate_for_next_candidate( + _local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "live-non-finite-rate", + {"input_cost_per_second": float("nan"), "litellm_provider": "openai", "mode": "realtime"}, + ) + monkeypatch.setitem( + litellm.model_cost, + "live-fallback-rate", + {"input_cost_per_second": 0.25, "litellm_provider": "openai", "mode": "realtime"}, + ) + litellm.get_model_info.cache_clear() + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "live-non-finite-rate"}}, + {"type": "session.usage.updated", "usage": {"seconds": 4}}, + ) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="live-fallback-rate", + ) + + assert cost == pytest.approx(4 * 0.25) + + +def test_live_session_uses_wall_clock_duration_when_greater_than_reported(_local_model_cost_map: None) -> None: + start_time: Final = datetime.datetime(2025, 1, 1) + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 15}}, + ) + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + litellm_logging_obj=_live_logging_object(start_time, start_time + datetime.timedelta(seconds=20)), + ) + + assert cost == pytest.approx(20 * rate) + + +def test_live_session_deduplicates_delegated_responses_and_charges_voice(_local_model_cost_map: None) -> None: + delegated_response: Final[dict[str, object]] = { + "type": "response.event", + "delegation_id": "delegation-1", + "event": { + "type": "response.completed", + "response": { + "id": "resp-1", + "model": "gpt-4o-mini", + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + }, + } + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.closed", "usage": {"seconds": 12}}, + delegated_response, + delegated_response, + ) + live_rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + backend_model: Final = litellm.model_cost["gpt-4o-mini"] + expected_backend_cost: Final = ( + 10 * cast(float, backend_model["input_cost_per_token"]) + + 5 * cast(float, backend_model["output_cost_per_token"]) + ) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + ) + + assert cost == pytest.approx(12 * live_rate + expected_backend_cost) + + +def test_live_session_zero_closed_seconds_do_not_use_usage_fallback(_local_model_cost_map: None) -> None: + start_time: Final = datetime.datetime(2025, 1, 1) + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 15}}, + {"type": "session.closed", "usage": {"seconds": 0}}, + ) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + litellm_logging_obj=_live_logging_object(start_time, start_time + datetime.timedelta(seconds=20)), + ) + + assert cost == 0.0 + + +def test_live_session_uses_latest_snapshot_when_closed_usage_is_missing(_local_model_cost_map: None) -> None: + start_time: Final = datetime.datetime(2025, 1, 1) + results: Final = _live_session_results( + {"type": "session.started", "session": {"model": "gpt-live-1"}}, + {"type": "session.usage.updated", "usage": {"seconds": 15}}, + {"type": "session.closed", "reason": "server_end"}, + ) + rate: Final = cast(float, litellm.model_cost["gpt-live-1"]["input_cost_per_second"]) + + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=Usage(), + custom_llm_provider="openai", + litellm_model_name="gpt-live-1", + litellm_logging_obj=_live_logging_object(start_time, start_time + datetime.timedelta(seconds=20)), + ) + + assert cost == pytest.approx(15 * rate) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f44c425c2c0..06170f09504 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -9142,6 +9142,26 @@ export interface paths { patch?: never; trace?: never; }; + "/live/sessions": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * WebSocket: live_sessions_websocket_endpoint + * @description WebSocket connection endpoint + */ + get: operations["websocket_live_sessions_websocket_endpoint_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/login": { parameters: { query?: never; @@ -10822,6 +10842,26 @@ export interface paths { patch?: never; trace?: never; }; + "/openai/v1/live/sessions": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * WebSocket: live_sessions_websocket_endpoint + * @description WebSocket connection endpoint + */ + get: operations["websocket_live_sessions_websocket_endpoint_get_3"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/openai/v1/realtime": { parameters: { query?: never; @@ -20268,6 +20308,26 @@ export interface paths { patch?: never; trace?: never; }; + "/v1/live/sessions": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * WebSocket: live_sessions_websocket_endpoint + * @description WebSocket connection endpoint + */ + get: operations["websocket_live_sessions_websocket_endpoint_get_2"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/v1/logs": { parameters: { query?: never; @@ -62788,6 +62848,24 @@ export interface operations { }; }; }; + websocket_live_sessions_websocket_endpoint_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description WebSocket Protocol Switched */ + 101: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; login_login_post: { parameters: { query?: never; @@ -65399,6 +65477,24 @@ export interface operations { }; }; }; + websocket_live_sessions_websocket_endpoint_get_3: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description WebSocket Protocol Switched */ + 101: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; websocket_realtime_websocket_endpoint_get_3: { parameters: { query?: never; @@ -77635,6 +77731,24 @@ export interface operations { }; }; }; + websocket_live_sessions_websocket_endpoint_get_2: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description WebSocket Protocol Switched */ + 101: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; ingest_otlp_traces_v1_logs_post: { parameters: { query?: never; From ec28acf322cf48f32a5c0c26524c961be351a117 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 13 Oct 2026 14:39:52 -0700 Subject: [PATCH 015/325] fix(mcp): resolve a bare tool name on a single-server scope against that server (#46225) A JSON-RPC tools/call on /{server_name}/mcp (or /mcp with x-mcp-servers naming one server) that uses the upstream's bare tool name answered "Tool not found" on every proxy process that had not yet listed that server, because the per-process tool-name map stays empty for a server whose upstream only lists tools with the caller's own credentials, and mcp SDK 2 no longer lists before an uncached call. call_mcp_tool now hands the single scoped server to execute_mcp_tool as requested_server_id, so the existing REST contract lists the server with the caller's headers and routes the bare name to it, and a bare name shared by two servers goes to the one the path names Co-authored-by: aayushbaluni <73417844+aayushbaluni@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_experimental/mcp_server/operations.py | 16 ++++ .../mcp_server/test_mcp_server.py | 41 ++++------ .../mcp_server/test_operations.py | 82 +++++++++++++++++++ 3 files changed, 112 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 5087bc6df09..889ca709eb9 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -2451,6 +2451,20 @@ async def fire_mcp_tool_call_failure_logging( ) +def _single_scoped_server( + mcp_servers: Sequence[str] | None, allowed_mcp_servers: Sequence[MCPServer] +) -> MCPServer | None: + """The one server a ``/{server_name}/mcp`` path or ``x-mcp-servers`` header scoped this call to. + + A bare tool name on such a request resolves against that server, the way a REST ``server_id`` + does, instead of against the per-process tool-name map, which stays empty for a server whose + upstream only lists tools with the caller's own credentials until this process has listed it. + """ + if mcp_servers is None or len(mcp_servers) != 1 or len(allowed_mcp_servers) != 1: + return None + return allowed_mcp_servers[0] + + @client @catalog_operation(lambda: global_mcp_server_manager) async def call_mcp_tool( @@ -2512,6 +2526,7 @@ async def call_mcp_tool( detail="User not allowed to call this tool.", ) + scoped_server: Final = _single_scoped_server(mcp_servers, allowed_mcp_servers) # Delegate to execute_mcp_tool for execution response = await execute_mcp_tool( name=name, @@ -2528,6 +2543,7 @@ async def call_mcp_tool( input_responses=input_responses, request_state=request_state, incoming_bearer_token=incoming_bearer_token, + requested_server_id=scoped_server.server_id if scoped_server is not None else None, **kwargs, ) except Exception as e: diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py index 3345683517b..de711d737d0 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py @@ -3000,11 +3000,7 @@ async def test_call_mcp_tool_uses_manager_permission_lookup(): "get_mcp_server_by_id", return_value=mock_server, ), - patch.object( - global_mcp_server_manager, - "get_mcp_server_from_tool_name", - return_value=mock_server, - ) as mock_get_server, + patch.object(global_mcp_server_manager, "server_exposes_tool", return_value=True), patch( "litellm.proxy._experimental.mcp_server.operations.global_mcp_tool_registry" ) as mock_tool_registry, @@ -3029,24 +3025,18 @@ async def test_call_mcp_tool_uses_manager_permission_lookup(): assert result == expected_response mock_get_allowed.assert_awaited_once() - # We call `_get_mcp_server_from_tool_name` multiple times: - # - for logging/metadata - # - for resolving the server during dispatch - # - and inside `call_tool` for guardrails/hooks - # The exact count isn't important, only that it is used. - assert mock_get_server.call_count >= 2 - # First call should use the prefixed tool name - assert ( - mock_get_server.call_args_list[0][0][0] - == f"{mock_server.name}/gmail_send_email" - ) + # The request names exactly one server, so that server is the dispatch target + # and the tool name reaches it untouched (the "/" is not a LiteLLM prefix). + mock_handle_managed.assert_awaited_once() + assert mock_handle_managed.await_args.kwargs["server_name"] == mock_server.name + assert mock_handle_managed.await_args.kwargs["name"] == f"{mock_server.name}/gmail_send_email" @pytest.mark.asyncio async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permissions(): """ - Ensure `call_mcp_tool` correctly resolves the MCP server for an unprefixed tool - name and enforces server-level permissions using that resolved server. + Ensure `call_mcp_tool` resolves an unprefixed tool name on a request scoped to + one server against that server and enforces server-level permissions with it. """ from litellm.proxy._experimental.mcp_server.server import ( call_mcp_tool, @@ -3076,11 +3066,7 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission "get_mcp_server_by_id", return_value=mock_server, ), - patch.object( - global_mcp_server_manager, - "get_mcp_server_from_tool_name", - return_value=mock_server, - ) as mock_get_server, + patch.object(global_mcp_server_manager, "server_exposes_tool", return_value=True), patch( "litellm.proxy._experimental.mcp_server.operations.global_mcp_tool_registry" ) as mock_tool_registry, @@ -3097,7 +3083,7 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission mock_tool_registry.get_tool.return_value = None mock_handle_managed.return_value = expected_response - # Call with UNPREFIXED tool name; server should be resolved via mapping + # Call with UNPREFIXED tool name; the single server the request names is the target result = await call_mcp_tool( name="gmail_send_email", arguments={"body": "hello"}, @@ -3106,11 +3092,12 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission assert result == expected_response mock_get_allowed.assert_awaited_once() - # We should resolve the server at least once using the unprefixed name - assert mock_get_server.call_count >= 1 - assert mock_get_server.call_args_list[0][0][0] == "gmail_send_email" + mock_handle_managed.assert_awaited_once() + assert mock_handle_managed.await_args.kwargs["server_name"] == mock_server.name + assert mock_handle_managed.await_args.kwargs["name"] == "gmail_send_email" # Permissions check should be invoked with the resolved server name mock_is_allowed.assert_called_once() + assert mock_is_allowed.call_args.kwargs["server_name"] == mock_server.name @pytest.mark.asyncio diff --git a/tests/unit/proxy/_experimental/mcp_server/test_operations.py b/tests/unit/proxy/_experimental/mcp_server/test_operations.py index dcc4eb1a570..af547ebf6c7 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_operations.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_operations.py @@ -565,6 +565,88 @@ async def test_tools_call_warmup_does_not_consume_mcp_server_rpm() -> None: assert [tool.name for tool in listing.tools] == ["echo"] + +def _scoped_rig( + monkeypatch: pytest.MonkeyPatch, +) -> tuple[MCPServer, MCPServer, AsyncMock, AsyncMock]: + """Two config servers that both expose a bare ``list_rules`` on a manager that has listed neither.""" + from mcp.types import CallToolResult, TextContent + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + + manager: Final = MCPServerManager() + memory: Final = MCPServer( + server_id="id-memory", name="memory_server", server_name="memory_server", transport=MCPTransport.http + ) + other: Final = MCPServer( + server_id="id-other", name="other_server", server_name="other_server", transport=MCPTransport.http + ) + manager.config_mcp_servers = {server.server_id: server for server in (memory, other)} + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(operations, "global_mcp_server_manager", manager) + monkeypatch.setattr( + manager, "get_allowed_mcp_servers", AsyncMock(return_value=[memory.server_id, other.server_id]) + ) + monkeypatch.setattr(global_mcp_tool_registry, "published_tools", {}) + listing: Final = AsyncMock() + monkeypatch.setattr(operations, "_get_tools_from_mcp_servers", listing) + handled: Final = AsyncMock( + return_value=CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) + ) + monkeypatch.setattr(operations, "_handle_managed_mcp_tool", handled) + return memory, other, listing, handled + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope", ["memory_server", "other_server"]) +async def test_bare_tool_name_on_a_single_server_scope_routes_to_that_server_cold( + monkeypatch: pytest.MonkeyPatch, scope: str +) -> None: + """``/{server_name}/mcp`` (or ``x-mcp-servers`` naming one server) with the upstream's bare tool name + must list and call the scoped server on a process whose tool-name map has never seen it, and when two + servers share the bare name the scope decides (BerriAI/litellm#44831).""" + memory, other, listing, handled = _scoped_rig(monkeypatch) + expected: Final = memory if scope == "memory_server" else other + + async with operations.global_mcp_server_manager.catalog.operation(): + result: Final = await operations.call_mcp_tool( + name="list_rules", + arguments={}, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-scoped"), + mcp_servers=[scope], + ) + + assert result.is_error is False + assert listing.await_args.kwargs["mcp_servers"] == [expected.server_id] + assert handled.await_args.kwargs["server_name"] == expected.name + assert handled.await_args.kwargs["name"] == "list_rules" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mcp_servers", [None, ["memory_server", "other_server"]]) +async def test_bare_tool_name_without_a_single_server_scope_is_still_unresolved_cold( + monkeypatch: pytest.MonkeyPatch, mcp_servers: list[str] | None +) -> None: + """The scope shortcut only applies when the request names exactly one server; an ambiguous bare name + keeps today's answer rather than being guessed onto the first allowed server.""" + from fastapi import HTTPException + + _memory, _other, listing, handled = _scoped_rig(monkeypatch) + + async with operations.global_mcp_server_manager.catalog.operation(): + with pytest.raises(HTTPException) as unresolved: + await operations.call_mcp_tool( + name="list_rules", + arguments={}, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-scoped"), + mcp_servers=mcp_servers, + ) + + assert unresolved.value.status_code == 404 + listing.assert_not_awaited() + handled.assert_not_awaited() + @pytest.mark.asyncio async def test_tools_call_pre_call_check_enforces_mcp_server_rpm() -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager From 92faffa78d0e92caa8b0dbe45ff4b1ea54d56be4 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 13 Oct 2026 14:47:35 -0700 Subject: [PATCH 016/325] fix(proxy): bound SIGTERM connection wait by GRACEFUL_SHUTDOWN_TIMEOUT (#46186) Resolves LIT-9448 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- helm/litellm-helm/values.yaml | 8 +-- litellm/constants.py | 1 + litellm/proxy/proxy_cli.py | 12 +++- .../shutdown/graceful_shutdown_manager.py | 7 ++- .../proxy/shutdown/graceful_uvicorn_worker.py | 10 +++ tests/unit/proxy/test_proxy_cli.py | 62 +++++++++++++++++++ 6 files changed, 92 insertions(+), 8 deletions(-) create mode 100644 litellm/proxy/shutdown/graceful_uvicorn_worker.py diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index 48358010485..11954d121a5 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -440,10 +440,10 @@ db: # Prefer the native /health/drain preStop hook over a fixed `sleep`: it marks # the pod NotReady and blocks only until in-flight requests actually finish # (bounded by GRACEFUL_SHUTDOWN_TIMEOUT, default 30s), instead of always -# waiting the worst-case duration. The drain runs once (the preStop hook and -# the SIGTERM handler share it), so set terminationGracePeriodSeconds a few -# seconds above GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before -# SIGKILL. +# waiting the worst-case duration. After SIGTERM the server waits up to +# GRACEFUL_SHUTDOWN_TIMEOUT again for connections still open and then cancels +# them, so set terminationGracePeriodSeconds a few seconds above twice +# GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before SIGKILL. # # /health/drain is off by default; enable it with # general_settings.enable_drain_endpoint: true. The kubelet calls preStop diff --git a/litellm/constants.py b/litellm/constants.py index 6a6816eb4fa..a7ea4a2975b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -104,6 +104,7 @@ S3_PARTITION_GRANULARITY_ENV_VAR: Final = "S3_PARTITION_GRANULARITY" MAX_FILE_LIST_LIMIT: Final = 10000 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)) DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)) +GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS: Final = 10 budget_reservation_disabled_info_emitted = False DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1)) DEFAULT_SQS_BATCH_SIZE: Final = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512)) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 7764d877b98..3c5f55d87a8 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -308,12 +308,14 @@ class ProxyInitializationHelpers: import litellm from litellm._logging import _get_uvicorn_json_log_config, resolve_log_level + from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager uvicorn_args: Final = { "app": "litellm.proxy.proxy_server:app", "host": host, "port": port, "server_header": False, + "timeout_graceful_shutdown": GracefulShutdownManager.get_timeout(), } if log_config is not None: print(f"Using log_config: {log_config}") @@ -571,6 +573,12 @@ class ProxyInitializationHelpers: """ Run litellm with `gunicorn` """ + import math + + from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS + from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager + from litellm.proxy.shutdown.graceful_uvicorn_worker import build_graceful_shutdown_uvicorn_worker + if os.name == "nt": pass else: @@ -627,10 +635,12 @@ class ProxyInitializationHelpers: return self.application print(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} with {num_workers} workers\033[0m\n") + graceful_shutdown_timeout: Final = GracefulShutdownManager.get_timeout() gunicorn_options: Final = { "bind": f"{host}:{port}", "workers": num_workers, # default is 1 - "worker_class": "uvicorn.workers.UvicornWorker", + "worker_class": build_graceful_shutdown_uvicorn_worker(graceful_shutdown_timeout), + "graceful_timeout": math.ceil(graceful_shutdown_timeout) + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS, "preload": True, # Add the preload flag, "accesslog": "-", # Log to stdout "timeout": 600, # default to very high number, bedrock/anthropic.claude-v2:1 can take 30+ seconds for the 1st chunk to come in diff --git a/litellm/proxy/shutdown/graceful_shutdown_manager.py b/litellm/proxy/shutdown/graceful_shutdown_manager.py index d07ef05abf7..d7a07ec34e5 100644 --- a/litellm/proxy/shutdown/graceful_shutdown_manager.py +++ b/litellm/proxy/shutdown/graceful_shutdown_manager.py @@ -2,9 +2,10 @@ Application-level graceful shutdown coordination for the LiteLLM proxy. Kubernetes terminates a pod by sending ``SIGTERM`` and, after -``terminationGracePeriodSeconds``, ``SIGKILL``. By default LiteLLM delegates -the signal to uvicorn and tears down immediately, dropping any in-flight -requests (streaming, batch inference, long-lived calls). +``terminationGracePeriodSeconds``, ``SIGKILL``. On ``SIGTERM`` uvicorn stops +accepting connections and waits up to ``GRACEFUL_SHUTDOWN_TIMEOUT`` for open +ones, cancelling whatever is still running, before the lifespan teardown that +runs this drain. A fixed ``preStop`` sleep can not solve this: it has to be sized for the *worst-case* request, so it either wastes time on every routine shutdown or is diff --git a/litellm/proxy/shutdown/graceful_uvicorn_worker.py b/litellm/proxy/shutdown/graceful_uvicorn_worker.py new file mode 100644 index 00000000000..6ae5ae60c3f --- /dev/null +++ b/litellm/proxy/shutdown/graceful_uvicorn_worker.py @@ -0,0 +1,10 @@ +from uvicorn.workers import UvicornWorker + + +def build_graceful_shutdown_uvicorn_worker(timeout_graceful_shutdown: float) -> type[UvicornWorker]: + class GracefulShutdownUvicornWorker(UvicornWorker): + def __init__(self, *args: object) -> None: + self.CONFIG_KWARGS = {**UvicornWorker.CONFIG_KWARGS, "timeout_graceful_shutdown": timeout_graceful_shutdown} + super().__init__(*args) + + return GracefulShutdownUvicornWorker diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 33f0216ff85..9566b6e30a2 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -1,4 +1,5 @@ import inspect +import logging import os from contextlib import nullcontext from pathlib import Path @@ -22,8 +23,10 @@ import yaml from uvicorn.config import LOOP_FACTORIES from uvicorn.importer import import_from_string +from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS from litellm.proxy import proxy_cli from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server +from litellm.proxy.shutdown.graceful_shutdown_manager import DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT @pytest.fixture(autouse=True) @@ -153,6 +156,22 @@ class TestProxyInitializationHelpers: ) assert args["timeout_worker_healthcheck"] == 15 + @pytest.mark.parametrize( + ("raw_timeout", "expected"), + [("7.5", 7.5), ("0", 0.0), (None, DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT)], + ) + def test_uvicorn_graceful_shutdown_is_bounded_by_graceful_shutdown_timeout( + self, monkeypatch, raw_timeout, expected + ): + if raw_timeout is None: + monkeypatch.delenv("GRACEFUL_SHUTDOWN_TIMEOUT", raising=False) + else: + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", raw_timeout) + + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000) + + assert uvicorn.Config(**args).timeout_graceful_shutdown == expected + @staticmethod def _uvicorn_access_info_enabled(args: dict) -> bool: import logging @@ -2050,6 +2069,49 @@ class TestProxyInitializationHelpers: assert captured["options"]["max_requests"] == 1000 assert captured["options"]["max_requests_jitter"] == 50 + @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") + def test_gunicorn_workers_bound_graceful_shutdown_by_graceful_shutdown_timeout(self, monkeypatch): + pytest.importorskip("gunicorn") + import gunicorn.config + import gunicorn.glogging + from uvicorn.workers import UvicornWorker + + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "7.5") + captured: dict = {} + + def capture_run(self): + captured["options"] = dict(self.options) + + with ( + patch("gunicorn.app.base.BaseApplication.run", capture_run), + patch( # test-quality-ok: option tests must not start a thread or change the pytest worker's child ownership + "litellm.proxy.proxy_cli.start_query_engine_reaper" + ), + ): + ProxyInitializationHelpers._run_gunicorn_server( + host="127.0.0.1", + port=4011, + app=MagicMock(), + num_workers=2, + ssl_certfile_path=None, + ssl_keyfile_path=None, + ) + + for name in ("uvicorn.error", "uvicorn.access"): + uvicorn_logger = logging.getLogger(name) + for attr in ("handlers", "level", "propagate"): + monkeypatch.setattr(uvicorn_logger, attr, getattr(uvicorn_logger, attr)) + gunicorn_cfg = gunicorn.config.Config() + worker = captured["options"]["worker_class"]( + 0, os.getpid(), [], MagicMock(), 30, gunicorn_cfg, gunicorn.glogging.Logger(gunicorn_cfg) + ) + + assert isinstance(worker, UvicornWorker) + assert worker.config.timeout_graceful_shutdown == 7.5 + assert worker.config.loop == UvicornWorker.CONFIG_KWARGS["loop"] + assert worker.config.http == UvicornWorker.CONFIG_KWARGS["http"] + assert captured["options"]["graceful_timeout"] == 8 + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS + @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") def test_gunicorn_master_is_reserved_for_forking_before_it_runs(self, fork_reservation): """preload forks workers from the master, so native routes are forbidden there first""" From 7bd108139f0260226215558d1db5f9528ceffb6b Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 13 Oct 2026 15:23:07 -0700 Subject: [PATCH 017/325] fix(openai): carry cache_control to prompt_cache_breakpoint on /v1/responses to gpt-5.6 (#46105) * fix(openai): carry cache_control to prompt_cache_breakpoint on /v1/responses to gpt-5.6 Native /v1/responses used to strip every Anthropic cache_control mark, so a gpt-5.6 deployment in explicit mode never wrote a cache. Marks on system, developer and user parts and on function_call_output now become prompt_cache_breakpoint on the same block when the deployment reads breakpoints, and are stripped exactly as before everywhere else * fix(openai): read str providers and base_url in the prompt cache breakpoint gate Responses configs generated from a provider JSON file, and openai_like itself, name their provider as a plain str, so .value crashed every native Responses request through them. The gate also only read api_base, while responses() keeps a caller's base_url as an extra field on the params, so an injected mark on a custom host was rewritten to a breakpoint * test(openai): cover the pass-through branches of the prompt cache breakpoint swap A part type OpenAI does not cache keeps no breakpoint and an item mark skips it, and an item or content of a shape OpenAI will reject reaches the strip untouched instead of failing inside the proxy * fix(prompt_caching): keep a client-written prompt_cache_breakpoint when a cache_control mark lands on the same block --- .../prompt_templates/common_utils.py | 3 +- .../llms/openai/responses/transformation.py | 119 ++++++- ..._responses_prompt_cache_breakpoint_wire.py | 229 +++++++++++++ ...ore_utils_prompt_templates_common_utils.py | 26 ++ .../test_openai_responses_transformation.py | 310 +++++++++++++++++- 5 files changed, 678 insertions(+), 9 deletions(-) create mode 100644 tests/integration/providers/test_openai_responses_prompt_cache_breakpoint_wire.py diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 591cae6dfa4..3c1b3cf964b 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1788,7 +1788,8 @@ _MarkedT: Final = TypeVar("_MarkedT", bound=Mapping[str, object]) def with_prompt_cache_breakpoint(target: _MarkedT, marker: object) -> _MarkedT: - if marker is None: + """A breakpoint the client already wrote on the block wins, TTL included.""" + if marker is None or target.get("prompt_cache_breakpoint") is not None: return target marked: Final = {**target, "prompt_cache_breakpoint": marker} return cast(_MarkedT, marked) # cast-ok: same block shape as the input plus the marker key diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index abd3d03b139..7c159c6a6a2 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -10,14 +10,17 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger +from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( safe_convert_created_field, ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( + cache_control_as_prompt_cache_breakpoint, drop_non_python_regex_patterns, flatten_combinators_and_drop_non_python_regex_patterns, + with_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -51,6 +54,75 @@ _MODEL_FAMILIES_REJECTING_TOP_LEVEL_SCHEMA_COMBINATORS: Final = ("gpt-4", "gpt-3 _PROVIDERS_WITH_OPENAI_SCHEMA_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) _PROVIDERS_REPLAYING_ONLY_THEIR_OWN_REASONING: Final = _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS +_ROLES_READING_PROMPT_CACHE_BREAKPOINTS: Final = frozenset({"system", "developer", "user"}) +_PART_TYPES_READING_PROMPT_CACHE_BREAKPOINTS: Final = frozenset({"input_text", "input_image", "input_file"}) +_INPUT_ITEM: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_CONTENT_PARTS: Final[TypeAdapter[Sequence[Mapping[str, object]]]] = TypeAdapter(Sequence[Mapping[str, object]]) + + +class _DeploymentHost(BaseModel): + """`responses()` builds the params from raw kwargs, where callers may spell the host `base_url`.""" + + model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True) + + api_base: object = None + base_url: object = None + + +def _provider_name(provider: LlmProviders | str) -> str: + """Configs generated from a provider JSON file return the slug as a plain str.""" + return provider.value if isinstance(provider, LlmProviders) else provider + + +def _item_with_prompt_cache_breakpoints(item: object) -> object: + """Marks on assistant output and on tools are left for the strip that follows: OpenAI writes no cache there.""" + try: + fields: Final = _INPUT_ITEM.validate_python(item) + except ValidationError: + return item + if fields.get("type") == "function_call_output": + return _with_prompt_cache_breakpoints_under(fields, "output") + if fields.get("type", "message") == "message" and fields.get("role") in _ROLES_READING_PROMPT_CACHE_BREAKPOINTS: + return _with_prompt_cache_breakpoints_under(fields, "content") + return item + + +def _with_prompt_cache_breakpoints_under(item: Mapping[str, object], key: str) -> Mapping[str, object]: + value: Final = item.get(key) + item_mark: Final = item.get("cache_control") + unmarked_item: Final = {name: field for name, field in item.items() if name != "cache_control"} + if isinstance(value, str): + if item_mark is None: + return item + text_part: Final[Mapping[str, object]] = {"type": "input_text", "text": value} + return {**unmarked_item, key: [with_prompt_cache_breakpoint(text_part, PromptCacheBreakpoint(mode="explicit"))]} + try: + parts: Final = _CONTENT_PARTS.validate_python(value) + except ValidationError: + return item + swapped_parts: Final = tuple(_part_with_prompt_cache_breakpoint(part) for part in parts) + return {**unmarked_item, key: list(_with_breakpoint_on_last_reading_part(swapped_parts, item_mark))} + + +def _part_with_prompt_cache_breakpoint(part: Mapping[str, object]) -> Mapping[str, object]: + if part.get("type") not in _PART_TYPES_READING_PROMPT_CACHE_BREAKPOINTS: + return part + return cache_control_as_prompt_cache_breakpoint(part) + + +def _with_breakpoint_on_last_reading_part( + parts: Sequence[Mapping[str, object]], item_mark: object +) -> Sequence[Mapping[str, object]]: + if item_mark is None: + return parts + reading_indexes: Final = [ + index for index, part in enumerate(parts) if part.get("type") in _PART_TYPES_READING_PROMPT_CACHE_BREAKPOINTS + ] + if not reading_indexes: + return parts + last: Final = reading_indexes[-1] + marked_last: Final = with_prompt_cache_breakpoint(parts[last], PromptCacheBreakpoint(mode="explicit")) + return (*parts[:last], marked_last, *parts[last + 1 :]) class _ReasoningSupportEntry(LiteLLMBaseModel): @@ -293,19 +365,22 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> dict: - """Strip Anthropic-only `cache_control` markers before sending to OpenAI. - - OpenAI's Responses API rejects unknown fields on input content blocks - with HTTP 400 ("Unknown parameter: 'input[0].content[0].cache_control'"). - Chat Completions strips these in - `remove_cache_control_flag_from_messages_and_tools`; mirror that here. + """Anthropic-only `cache_control` marks become `prompt_cache_breakpoint` on the same block + where the deployment reads them (gpt-5.6 and newer on OpenAI's API, or a host opted in + through `prompt_cache_options`) and are stripped everywhere else, because OpenAI's + Responses API rejects unknown fields on input content blocks with HTTP 400 + ("Unknown parameter: 'input[0].content[0].cache_control'"). """ + request_params: Final = cast( # cast-ok: the base signature leaves the params bag untyped + "Mapping[str, object]", response_api_optional_request_params + ) replay_safe_input, sanitized_tools = self._prepared_input_and_tools( model=model, input=input, tools=response_api_optional_request_params.get("tools"), litellm_params=litellm_params, + prompt_cache_options=request_params.get("prompt_cache_options"), ) if sanitized_tools is not None: response_api_optional_request_params["tools"] = sanitized_tools @@ -321,10 +396,18 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): input: str | ResponseInputParam, tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None, litellm_params: GenericLiteLLMParams, + prompt_cache_options: object = None, ) -> tuple[str | ResponseInputParam, Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None]: validated_input: Final = self._validate_input_param(self._drop_bridge_minted_reasoning_items(input)) + marked_input: Final = ( + self._input_with_prompt_cache_breakpoints(validated_input) + if self._reads_prompt_cache_breakpoints( + model=model, litellm_params=litellm_params, prompt_cache_options=prompt_cache_options + ) + else validated_input + ) stripped_input, stripped_tools = self.remove_cache_control_flag_from_input_and_tools( - model=model, input=validated_input, tools=tools + model=model, input=marked_input, tools=tools ) object_schema_tools: Final = self._tools_with_object_parameters(model=model, tools=stripped_tools) sanitized_tools: Final = self._sanitized_tool_schemas_for_openai( @@ -332,6 +415,24 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): ) return self._drop_foreign_tool_call_item_ids(stripped_input), sanitized_tools + def _reads_prompt_cache_breakpoints( + self, model: str, litellm_params: GenericLiteLLMParams, prompt_cache_options: object + ) -> bool: + host: Final = _DeploymentHost.model_validate(litellm_params) + return AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint( + model=model, + custom_llm_provider=_provider_name(self.custom_llm_provider), + api_base=host.api_base or host.base_url, + prompt_cache_options=prompt_cache_options, + ) + + @staticmethod + def _input_with_prompt_cache_breakpoints(input: str | ResponseInputParam) -> str | ResponseInputParam: + if not isinstance(input, list): + return input + marked_items: Final = [_item_with_prompt_cache_breakpoints(item) for item in input] + return cast("ResponseInputParam", marked_items) # cast-ok: items keep their shape, marks renamed in place + def _tools_with_object_parameters( self, model: str, tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None ) -> Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None: @@ -951,11 +1052,15 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): compact_path: Final = parsed_url.path.rstrip("/") + "/compact" url: Final = str(parsed_url.copy_with(path=compact_path)) + request_params: Final = cast( # cast-ok: the base signature leaves the params bag untyped + "Mapping[str, object]", response_api_optional_request_params + ) replay_safe_input, sanitized_tools = self._prepared_input_and_tools( model=model, input=input, tools=response_api_optional_request_params.get("tools"), litellm_params=litellm_params, + prompt_cache_options=request_params.get("prompt_cache_options"), ) if sanitized_tools is not None: response_api_optional_request_params["tools"] = sanitized_tools diff --git a/tests/integration/providers/test_openai_responses_prompt_cache_breakpoint_wire.py b/tests/integration/providers/test_openai_responses_prompt_cache_breakpoint_wire.py new file mode 100644 index 00000000000..f64023d1af7 --- /dev/null +++ b/tests/integration/providers/test_openai_responses_prompt_cache_breakpoint_wire.py @@ -0,0 +1,229 @@ +import uuid +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from typing import Final + +import httpx +import openai +import pytest +from integration._support import prompt_cache_breakpoint as pcb +from integration._support import responses_vendor as rv +from integration._support.client import Gateway, gateway_from_environment +from integration._support.wire import Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(120) + +_EPHEMERAL: Final[Mapping[str, JsonValue]] = {"type": "ephemeral"} +_EPHEMERAL_1H: Final[Mapping[str, JsonValue]] = {"type": "ephemeral", "ttl": "1h"} +_OPTED_IN: Final[Mapping[str, JsonValue]] = {"prompt_cache_options": {"mode": "explicit"}} +_LOOKUP_TOOL: Final[Mapping[str, JsonValue]] = { + "type": "function", + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, +} + + +@dataclass(frozen=True, slots=True) +class _Native: + gateway: Gateway + wire: Wire + reading: str + stripping: str + spend: pcb.SpendLogs + + +@pytest.fixture(scope="module") +def native() -> Iterator[_Native]: + with ( + wire_server(pcb.respond) as wire, + gateway_from_environment() as gateway, + gateway.scenario() as scenario, + pcb.spend_logs() as spend, + ): + api_base: Final = f"{wire.url}/v1" + yield _Native( + gateway, + wire, + scenario.model(model=pcb.MODEL, api_base=api_base, **_OPTED_IN), + scenario.model(model=pcb.MODEL, api_base=api_base), + spend, + ) + + +def _v1(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + "/v1" + + +def _input_text(text: str) -> dict[str, JsonValue]: + return {"type": "input_text", "text": text} + + +def _marked( + block: Mapping[str, JsonValue], cache_control: Mapping[str, JsonValue] = _EPHEMERAL +) -> dict[str, JsonValue]: + return {**block, "cache_control": dict(cache_control)} + + +def _marked_history(marker: str) -> list[JsonValue]: + return [ + { + "type": "message", + "role": "developer", + "content": [_input_text("Answer briefly"), _marked(_input_text("Stable rules"), _EPHEMERAL_1H)], + }, + _marked({"role": "user", "content": pcb.prompt(marker)}), + ] + + +def _responses(native: _Native, model: str, input: Sequence[JsonValue], *, stream: bool = False) -> httpx.Response: + return native.gateway.request( + "POST", "/v1/responses", {"model": model, "input": list(input), "stream": stream, **pcb.NO_CACHE} + ) + + +def _answered(response: httpx.Response, marker: str) -> str: + assert response.status_code == 200, response.text + assert rv.answer(marker) in response.text, response.text + return response.headers["x-litellm-call-id"] + + +def _wire_body(request: Request, *, stream: bool = False) -> dict[str, JsonValue]: + body: Final = pcb.body_of(request) + assert body["model"] == "gpt-6.1-sol", body + assert (body.get("stream") is True) is stream, body + assert b"cache_control" not in request.body, request.body + return body + + +def _assert_marked_history_on_wire(body: Mapping[str, JsonValue], marker: str) -> None: + assert body["prompt_cache_options"] == {"mode": "explicit"}, body + assert body["input"] == [ + { + "type": "message", + "role": "developer", + "content": [ + _input_text("Answer briefly"), + pcb.marked(_input_text("Stable rules"), pcb.EXPLICIT), + ], + }, + {"role": "user", "content": [pcb.marked(_input_text(pcb.prompt(marker)), pcb.EXPLICIT)]}, + ], body + + +def test_openai_sdk_cache_control_reaches_the_wire_as_a_prompt_cache_breakpoint_on_the_same_block( + native: _Native, +) -> None: + marker: Final = uuid.uuid4().hex + with openai.OpenAI(api_key=native.gateway.key, base_url=_v1(native.gateway), max_retries=0) as client: + raw: Final = client.responses.with_raw_response.create( + model=native.reading, input=_marked_history(marker), extra_body=dict(pcb.NO_CACHE) + ) + response: Final = raw.parse() + assert pcb.answers(response.id, marker), response + assert response.output_text == rv.answer(marker), response + _assert_marked_history_on_wire(_wire_body(pcb.posted(native.wire, marker)), marker) + native.spend.landed(native.reading, raw.headers["x-litellm-call-id"], marker) + + +def test_stream_carries_the_marks_the_same_way(native: _Native) -> None: + marker: Final = uuid.uuid4().hex + call_id: Final = _answered(_responses(native, native.reading, _marked_history(marker), stream=True), marker) + _assert_marked_history_on_wire(_wire_body(pcb.posted(native.wire, marker), stream=True), marker) + native.spend.landed(native.reading, call_id, marker) + + +@pytest.mark.parametrize( + ("output", "expected"), + ( + ("found it", [pcb.marked(_input_text("found it"), pcb.EXPLICIT)]), + ( + [_input_text("first"), _input_text("last")], + [_input_text("first"), pcb.marked(_input_text("last"), pcb.EXPLICIT)], + ), + ), + ids=("string-output", "list-output"), +) +def test_cache_control_on_a_function_call_output_marks_its_last_part( + native: _Native, output: JsonValue, expected: JsonValue +) -> None: + marker: Final = uuid.uuid4().hex + history: Final[list[JsonValue]] = [ + {"role": "user", "content": pcb.prompt(marker)}, + {"type": "function_call", "call_id": "call_1", "name": "lookup", "arguments": "{}"}, + _marked({"type": "function_call_output", "call_id": "call_1", "output": output}), + ] + call_id: Final = _answered(_responses(native, native.reading, history), marker) + request: Final = pcb.posted(native.wire, marker) + _wire_body(request) + assert pcb.function_output(pcb.input_items(request), "call_1") == expected, ( + "OpenAI caches up to a breakpoint on a function_call_output part (prompt caching guide, read 2026-10-12)" + ) + native.spend.landed(native.reading, call_id, marker) + + +def test_cache_control_on_an_assistant_turn_and_on_a_tool_is_dropped(native: _Native) -> None: + marker: Final = uuid.uuid4().hex + history: Final[list[JsonValue]] = [ + {"role": "user", "content": "first question"}, + { + "type": "message", + "role": "assistant", + "content": [_marked({"type": "output_text", "text": "earlier answer"})], + }, + {"role": "user", "content": pcb.prompt(marker)}, + ] + response: Final = native.gateway.request( + "POST", + "/v1/responses", + {"model": native.reading, "input": history, "tools": [_marked(_LOOKUP_TOOL)], **pcb.NO_CACHE}, + ) + call_id: Final = _answered(response, marker) + request: Final = pcb.posted(native.wire, marker) + body: Final = _wire_body(request) + assert b"prompt_cache_breakpoint" not in request.body, request.body + assert pcb.content_of(pcb.input_items(request), "assistant") == [ + {"type": "output_text", "text": "earlier answer"} + ], "OpenAI accepts a breakpoint on output_text but writes no cache for it (live gpt-5.6, 2026-10-12)" + assert body["tools"] == [dict(_LOOKUP_TOOL)], body + native.spend.landed(native.reading, call_id, marker) + + +def test_cache_control_is_stripped_where_the_deployment_does_not_read_breakpoints(native: _Native) -> None: + marker: Final = uuid.uuid4().hex + call_id: Final = _answered(_responses(native, native.stripping, _marked_history(marker)), marker) + request: Final = pcb.posted(native.wire, marker) + body: Final = _wire_body(request) + assert "prompt_cache_options" not in body, body + assert b"prompt_cache_breakpoint" not in request.body, request.body + assert body["input"] == [ + { + "type": "message", + "role": "developer", + "content": [_input_text("Answer briefly"), _input_text("Stable rules")], + }, + {"role": "user", "content": pcb.prompt(marker)}, + ], body + native.spend.landed(native.stripping, call_id, marker) + + +def test_a_breakpoint_the_client_already_wrote_is_left_alone_beside_a_converted_mark(native: _Native) -> None: + marker: Final = uuid.uuid4().hex + history: Final[list[JsonValue]] = [ + { + "role": "user", + "content": [ + pcb.marked(_input_text(pcb.prompt(marker)), pcb.EXPLICIT_30M), + _marked(_input_text("and nothing else")), + ], + } + ] + call_id: Final = _answered(_responses(native, native.reading, history), marker) + request: Final = pcb.posted(native.wire, marker) + _wire_body(request) + (item,) = pcb.input_items(request) + assert item["content"] == [ + pcb.marked(_input_text(pcb.prompt(marker)), pcb.EXPLICIT_30M), + pcb.marked(_input_text("and nothing else"), pcb.EXPLICIT), + ], item + native.spend.landed(native.reading, call_id, marker) diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index c040b161ba4..25d929ae025 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -30,6 +30,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_encrypted_reasoning_from_messages, system_messages_first, update_messages_with_model_file_ids, + with_prompt_cache_breakpoint, ) _ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' @@ -2280,3 +2281,28 @@ def test_a_block_without_cache_control_is_returned_untouched() -> None: unmarked: Final = {"type": "text", "text": "rules"} assert cache_control_as_prompt_cache_breakpoint(unmarked) is unmarked + + +def test_a_breakpoint_the_client_wrote_survives_a_message_level_mark_ttl_included() -> None: + block: Final = { + "type": "input_text", + "text": "rules", + "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}, + } + + assert with_prompt_cache_breakpoint(block, {"mode": "explicit"}) is block + + +def test_a_breakpoint_the_client_wrote_survives_a_per_part_cache_control() -> None: + both: Final = { + "type": "input_text", + "text": "rules", + "cache_control": {"type": "ephemeral"}, + "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}, + } + + assert cache_control_as_prompt_cache_breakpoint(both) == { + "type": "input_text", + "text": "rules", + "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}, + } diff --git a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py index 2121e04dcac..2a35031202a 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py @@ -1,5 +1,5 @@ import json -from collections.abc import Mapping +from collections.abc import Callable, Mapping from types import SimpleNamespace from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -2303,6 +2303,314 @@ class TestToolSchemaRegexPatternWiring: assert result["tools"][0] is tool +_OPENAI_API: Final = GenericLiteLLMParams(api_base="https://api.openai.com/v1") +_OTHER_HOST: Final = GenericLiteLLMParams(api_base="https://gateway.example/v1") +_EXPLICIT_OPT_IN: Final = {"prompt_cache_options": {"mode": "explicit"}} +_EPHEMERAL: Final = {"type": "ephemeral"} +_EXPLICIT_BREAKPOINT: Final = {"mode": "explicit"} + + +def _marked_input(cache_control: dict[str, str] | None) -> list[dict[str, object]]: + mark: Final = {} if cache_control is None else {"cache_control": cache_control} + return [ + { + "role": "developer", + "content": [ + {"type": "input_text", "text": "Be concise."}, + {"type": "input_text", "text": "Be helpful.", **mark}, + ], + }, + {"type": "message", "role": "user", "content": "look", **mark}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "looking", **mark}]}, + {"type": "function_call", "call_id": "call_1", "name": "lookup", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_1", "output": "found", **mark}, + {"type": "function_call", "call_id": "call_2", "name": "lookup", "arguments": "{}"}, + { + "type": "function_call_output", + "call_id": "call_2", + "output": [{"type": "input_text", "text": "first"}, {"type": "input_text", "text": "last"}], + **mark, + }, + ] + + +def _marked_tool(cache_control: dict[str, str] | None) -> dict[str, object]: + mark: Final = {} if cache_control is None else {"cache_control": cache_control} + return {"type": "function", "name": "lookup", "parameters": {"type": "object", "properties": {}}, **mark} + + +def _openai_like_responses_config() -> OpenAIResponsesAPIConfig: + from litellm.llms.openai_like.responses.transformation import OpenAILikeResponsesConfig + + return OpenAILikeResponsesConfig() + + +def _generated_sail_responses_config() -> OpenAIResponsesAPIConfig: + from litellm.llms.openai_like.dynamic_config import create_responses_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + sail = JSONProviderRegistry.get("sail") + assert sail is not None + return create_responses_config_class(sail)() + + +class TestCacheControlBecomesPromptCacheBreakpoint: + """Clients that speak Anthropic's dialect mark blocks with `cache_control`. On a deployment that + reads OpenAI's `prompt_cache_breakpoint` the mark moves onto the same block instead of being + stripped, so explicit-mode prompt caching keeps working through the proxy.""" + + @pytest.fixture(autouse=True) + def _no_openai_api_base_override(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", None) + + def _transform( + self, + model: str, + litellm_params: GenericLiteLLMParams, + cache_control: dict[str, str] | None, + extra_params: dict[str, object] | None = None, + config: OpenAIResponsesAPIConfig | None = None, + ) -> dict[str, object]: + return (config or OpenAIResponsesAPIConfig()).transform_responses_api_request( + model=model, + input=_marked_input(cache_control), + response_api_optional_request_params={"tools": [_marked_tool(cache_control)], **(extra_params or {})}, + litellm_params=litellm_params, + headers={}, + ) + + @pytest.mark.parametrize( + ("litellm_params", "extra_params"), + [ + pytest.param(_OPENAI_API, {}, id="api.openai.com"), + pytest.param( + GenericLiteLLMParams(base_url="https://api.openai.com/v1"), {}, id="api.openai.com-as-base_url" + ), + pytest.param(GenericLiteLLMParams(), {}, id="default-host"), + pytest.param(_OTHER_HOST, _EXPLICIT_OPT_IN, id="other-host-opted-in"), + ], + ) + def test_marks_reach_gpt_5_6_as_explicit_breakpoints_on_the_same_blocks( + self, litellm_params: GenericLiteLLMParams, extra_params: dict[str, object] + ) -> None: + result = self._transform("gpt-5.6", litellm_params, {**_EPHEMERAL, "ttl": "1h"}, extra_params) + + assert result["input"] == [ + { + "role": "developer", + "content": [ + {"type": "input_text", "text": "Be concise."}, + {"type": "input_text", "text": "Be helpful.", "prompt_cache_breakpoint": _EXPLICIT_BREAKPOINT}, + ], + }, + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "look", "prompt_cache_breakpoint": _EXPLICIT_BREAKPOINT}], + }, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "looking"}]}, + {"type": "function_call", "call_id": "call_1", "name": "lookup", "arguments": "{}"}, + { + "type": "function_call_output", + "call_id": "call_1", + "output": [{"type": "input_text", "text": "found", "prompt_cache_breakpoint": _EXPLICIT_BREAKPOINT}], + }, + {"type": "function_call", "call_id": "call_2", "name": "lookup", "arguments": "{}"}, + { + "type": "function_call_output", + "call_id": "call_2", + "output": [ + {"type": "input_text", "text": "first"}, + {"type": "input_text", "text": "last", "prompt_cache_breakpoint": _EXPLICIT_BREAKPOINT}, + ], + }, + ] + assert result["tools"] == [_marked_tool(None)] + assert result.get("prompt_cache_options") == extra_params.get("prompt_cache_options") + + def test_per_part_marks_win_over_the_item_mark_and_keep_their_own_block(self) -> None: + result = OpenAIResponsesAPIConfig().transform_responses_api_request( + model="gpt-5.6", + input=[ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "a", "cache_control": _EPHEMERAL}, + {"type": "input_image", "image_url": "https://img.example/x.png", "detail": "auto"}, + {"type": "input_text", "text": "b"}, + ], + "cache_control": _EPHEMERAL, + } + ], + response_api_optional_request_params={}, + litellm_params=_OPENAI_API, + headers={}, + ) + + assert result["input"] == [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "a", "prompt_cache_breakpoint": _EXPLICIT_BREAKPOINT}, + {"type": "input_image", "image_url": "https://img.example/x.png", "detail": "auto"}, + {"type": "input_text", "text": "b", "prompt_cache_breakpoint": _EXPLICIT_BREAKPOINT}, + ], + } + ] + + @pytest.mark.parametrize( + ("model", "litellm_params"), + [ + pytest.param("gpt-4o", _OPENAI_API, id="model-older-than-gpt-5.6"), + pytest.param("gpt-5.6", _OTHER_HOST, id="other-host-not-opted-in"), + pytest.param( + "gpt-5.6", GenericLiteLLMParams(base_url="https://gateway.example/v1"), id="other-host-as-base_url" + ), + ], + ) + def test_marks_are_still_stripped_where_the_deployment_does_not_read_breakpoints( + self, model: str, litellm_params: GenericLiteLLMParams + ) -> None: + marked = self._transform(model, litellm_params, _EPHEMERAL) + unmarked = self._transform(model, litellm_params, None) + + assert marked == unmarked + assert "prompt_cache_breakpoint" not in json.dumps(marked) + + def test_non_openai_subclass_still_strips_on_a_gpt_5_6_deployment(self) -> None: + from litellm.llms.hosted_vllm.responses.transformation import HostedVLLMResponsesAPIConfig + + config = HostedVLLMResponsesAPIConfig() + marked = self._transform("gpt-5.6", _OPENAI_API, _EPHEMERAL, _EXPLICIT_OPT_IN, config=config) + unmarked = self._transform("gpt-5.6", _OPENAI_API, None, _EXPLICIT_OPT_IN, config=config) + + assert marked == unmarked + assert "prompt_cache_breakpoint" not in json.dumps(marked["input"]) + + @pytest.mark.parametrize( + "make_config", [_openai_like_responses_config, _generated_sail_responses_config], ids=["openai_like", "sail"] + ) + def test_a_config_naming_its_provider_as_a_plain_str_still_strips_on_a_gpt_5_6_deployment( + self, make_config: Callable[[], OpenAIResponsesAPIConfig] + ) -> None: + config = make_config() + assert type(config.custom_llm_provider) is str + + marked = self._transform("gpt-5.6", _OPENAI_API, _EPHEMERAL, _EXPLICIT_OPT_IN, config=config) + unmarked = self._transform("gpt-5.6", _OPENAI_API, None, _EXPLICIT_OPT_IN, config=config) + + assert marked == unmarked + assert "prompt_cache_breakpoint" not in json.dumps(marked["input"]) + + def test_a_breakpoint_the_client_wrote_on_the_same_block_wins_over_the_mark_ttl_included(self) -> None: + kept = {"mode": "explicit", "ttl": "30m"} + result = OpenAIResponsesAPIConfig().transform_responses_api_request( + model="gpt-5.6", + input=[ + { + "role": "user", + "content": [ + { + "type": "input_text", + "text": "a", + "prompt_cache_breakpoint": kept, + "cache_control": _EPHEMERAL, + } + ], + }, + { + "role": "user", + "content": [{"type": "input_text", "text": "b", "prompt_cache_breakpoint": kept}], + "cache_control": _EPHEMERAL, + }, + ], + response_api_optional_request_params={}, + litellm_params=_OPENAI_API, + headers={}, + ) + + assert result["input"] == [ + {"role": "user", "content": [{"type": "input_text", "text": "a", "prompt_cache_breakpoint": kept}]}, + {"role": "user", "content": [{"type": "input_text", "text": "b", "prompt_cache_breakpoint": kept}]}, + ] + + def test_marks_on_parts_openai_does_not_cache_are_dropped_and_the_item_mark_skips_them(self) -> None: + audio = {"type": "input_audio", "input_audio": {"data": "UklGRg==", "format": "wav"}} + result = OpenAIResponsesAPIConfig().transform_responses_api_request( + model="gpt-5.6", + input=[ + {"role": "user", "content": [{**audio, "cache_control": _EPHEMERAL}], "cache_control": _EPHEMERAL}, + { + "role": "user", + "content": [{"type": "input_text", "text": "a"}, {**audio, "cache_control": _EPHEMERAL}], + "cache_control": _EPHEMERAL, + }, + ], + response_api_optional_request_params={}, + litellm_params=_OPENAI_API, + headers={}, + ) + + assert result["input"] == [ + {"role": "user", "content": [audio]}, + { + "role": "user", + "content": [ + {"type": "input_text", "text": "a", "prompt_cache_breakpoint": _EXPLICIT_BREAKPOINT}, + audio, + ], + }, + ] + + def test_items_and_content_of_a_shape_openai_will_reject_pass_through_to_the_strip(self) -> None: + """Validating the body is upstream's job. The swap leaves what it cannot read alone, so the + caller gets OpenAI's 400 for it instead of a failure inside the proxy.""" + result = OpenAIResponsesAPIConfig().transform_responses_api_request( + model="gpt-5.6", + input=[ + "a bare string where an item belongs", + { + "type": "function_call_output", + "call_id": "call_1", + "output": {"result": 1}, + "cache_control": _EPHEMERAL, + }, + {"role": "user", "content": "look", "cache_control": _EPHEMERAL}, + ], + response_api_optional_request_params={}, + litellm_params=_OPENAI_API, + headers={}, + ) + + assert result["input"] == [ + "a bare string where an item belongs", + {"type": "function_call_output", "call_id": "call_1", "output": {"result": 1}}, + { + "role": "user", + "content": [{"type": "input_text", "text": "look", "prompt_cache_breakpoint": _EXPLICIT_BREAKPOINT}], + }, + ] + + def test_compact_request_swaps_marks_the_same_way(self) -> None: + _, data = OpenAIResponsesAPIConfig().transform_compact_response_api_request( + model="gpt-5.6", + input=[{"role": "user", "content": "look", "cache_control": _EPHEMERAL}], + response_api_optional_request_params=dict(_EXPLICIT_OPT_IN), + litellm_params=_OTHER_HOST, + headers={}, + api_base="https://gateway.example/v1/responses", + ) + + assert data["input"] == [ + { + "role": "user", + "content": [{"type": "input_text", "text": "look", "prompt_cache_breakpoint": _EXPLICIT_BREAKPOINT}], + } + ] + + class TestReasoningFollowsModelSupport: """Responses API clients like Codex send `reasoning` on every request, and OpenAI 400s it on non-reasoning models like gpt-4o. drop_params must strip it there, the same way the From 271b4f03066b8ff4ecec8bb80b89deb464b35a63 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Tue, 13 Oct 2026 15:30:14 -0700 Subject: [PATCH 018/325] refactor(rust): one stream codec shape per API format (#45996) Chat Completions, Messages and Responses each get a streaming module with the same three names, event_stream, decode and encode, built on a shared SSE helper in base_llm/sse.rs that owns the [DONE] marker. Callers use the module path instead of vendor-prefixed names. Encoding matches what the Python proxy writes to clients: Messages frames carry an event line, the OpenAI formats are bare data frames. Responses gets one event type for SSE and WebSocket. ResponsesEventType names every event on OpenAI's streaming-events reference, with unknown names riding in Other, and is_terminal and is_failure follow Python's iterator. The framer adds a bulk decode_all and an SseBlockCodec that keeps each event's wire bytes. An event the stream ends inside is dropped, as the SSE spec and the openai SDK's decoder do. ChatCompletionChunk no longer requires id, created or object, since compatible hosts omit them. --- litellm-rust/crates/framer/src/sse.rs | 283 ++++++++++++++-- litellm-rust/crates/framer/tests/sse.rs | 20 +- .../crates/inference-messages/src/handler.rs | 7 +- .../inference-responses/src/websocket.rs | 12 - .../src/formats/chat_completions/streaming.rs | 72 ++++- .../src/formats/messages/streaming.rs | 19 ++ .../src/formats/responses/AGENTS.md | 2 +- .../llms-types/src/formats/responses/mod.rs | 1 + .../src/formats/responses/streaming.rs | 305 ++++++++++++++++++ .../formats/responses/streaming_websocket.rs | 125 +------ litellm-rust/crates/llms/AGENTS.md | 7 + .../llms/src/anthropic/chat/transformation.rs | 4 +- .../llms/src/base_llm/chat/streaming.rs | 125 ++++++- .../llms/src/base_llm/messages/streaming.rs | 232 ++++++++----- litellm-rust/crates/llms/src/base_llm/mod.rs | 1 + .../crates/llms/src/base_llm/responses/mod.rs | 1 + .../llms/src/base_llm/responses/streaming.rs | 135 ++++++++ .../src/base_llm/responses/transformation.rs | 14 +- litellm-rust/crates/llms/src/base_llm/sse.rs | 102 ++++++ .../llms/src/bedrock/chat/invoke_handler.rs | 4 +- .../anthropic_claude3_transformation.rs | 6 +- .../src/openai/responses/transformation.rs | 8 +- 22 files changed, 1219 insertions(+), 266 deletions(-) create mode 100644 litellm-rust/crates/llms-types/src/formats/responses/streaming.rs create mode 100644 litellm-rust/crates/llms/src/base_llm/responses/streaming.rs create mode 100644 litellm-rust/crates/llms/src/base_llm/sse.rs diff --git a/litellm-rust/crates/framer/src/sse.rs b/litellm-rust/crates/framer/src/sse.rs index 6fee1cfab7f..2e7686e4343 100644 --- a/litellm-rust/crates/framer/src/sse.rs +++ b/litellm-rust/crates/framer/src/sse.rs @@ -1,6 +1,6 @@ use std::str; -use bytes::{Buf, BufMut, BytesMut}; +use bytes::{Buf, BufMut, Bytes, BytesMut}; use tokio_util::codec::{Decoder, Encoder}; use crate::SseError; @@ -29,38 +29,107 @@ impl Decoder for SseCodec { return Ok(None); } while let Some(end) = block_end(src) { - let block = src.split_to(end); - let pending = lines(&block) - .map(|(line, _)| line) - .take_while(|line| !line.is_empty()) - .try_fold(Pending::default(), Pending::apply)?; - if let Some(event) = pending.dispatch() { + if let Some(event) = event_in(&src.split_to(end))? { return Ok(Some(event)); } } Ok(None) } - fn decode_eof(&mut self, _pending: &mut BytesMut) -> Result, SseError> { - Ok(None) + /// An event the stream ended inside of is dropped, as the SSE spec requires. + fn decode_eof(&mut self, pending: &mut BytesMut) -> Result, SseError> { + self.decode(pending) } } +pub fn decode_all(wire: &[u8]) -> Result, SseError> { + let mut codec = SseCodec::default(); + let mut buffer = BytesMut::from(wire); + let mut events = Vec::new(); + while let Some(event) = codec.decode_eof(&mut buffer)? { + events.push(event); + } + Ok(events) +} + impl SseCodec { fn skip_bom(&mut self, src: &mut BytesMut) -> bool { - if self.past_bom { - return true; - } - if src.starts_with(BOM) { - src.advance(BOM.len()); - } else if BOM.starts_with(src) { - return false; - } - self.past_bom = true; - true + skip_bom(&mut self.past_bom, src) } } +fn event_in(block: &[u8]) -> Result, SseError> { + Ok(lines(block) + .map(|(line, _)| line) + .take_while(|line| !line.is_empty()) + .try_fold(Pending::default(), Pending::apply)? + .dispatch()) +} + +/// One event together with the exact wire bytes that carried it, for relays that forward a +/// provider's bytes untouched but still need to see each event go by. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SseBlock { + pub bytes: Bytes, + pub event: SseEvent, +} + +#[derive(Clone, Copy, Debug, Default)] +pub struct SseBlockCodec { + past_bom: bool, +} + +impl Decoder for SseBlockCodec { + type Item = SseBlock; + type Error = SseError; + + fn decode(&mut self, src: &mut BytesMut) -> Result, SseError> { + self.decode_block(src, false) + } + + fn decode_eof(&mut self, pending: &mut BytesMut) -> Result, SseError> { + self.decode_block(pending, true) + } +} + +impl SseBlockCodec { + fn decode_block( + &mut self, + src: &mut BytesMut, + eof: bool, + ) -> Result, SseError> { + if !skip_bom(&mut self.past_bom, src) { + return Ok(None); + } + while let Some(end) = block_end(src) { + if !eof && end == src.len() && src.ends_with(b"\r") { + return Ok(None); + } + let block = src.split_to(end); + if let Some(event) = event_in(&block)? { + return Ok(Some(SseBlock { + bytes: block.freeze(), + event, + })); + } + } + Ok(None) + } +} + +fn skip_bom(past_bom: &mut bool, src: &mut BytesMut) -> bool { + if *past_bom { + return true; + } + if src.starts_with(BOM) { + src.advance(BOM.len()); + } else if BOM.starts_with(src) { + return false; + } + *past_bom = true; + true +} + fn block_end(bytes: &[u8]) -> Option { lines(bytes) .find(|(line, _)| line.is_empty()) @@ -168,3 +237,179 @@ impl Encoder for SseCodec { Ok(()) } } + +#[cfg(test)] +mod tests { + use futures_util::{StreamExt, TryStreamExt, stream}; + use rstest::rstest; + + use super::*; + use crate::frames; + + fn pieces(wire: &[u8], size: usize) -> Vec { + wire.chunks(size).map(Bytes::copy_from_slice).collect() + } + + async fn blocks(wire: &[u8], size: usize) -> Result, SseError> { + frames( + stream::iter(pieces(wire, size).into_iter().map(Ok::<_, std::io::Error>)), + SseBlockCodec::default(), + ) + .try_collect() + .await + } + + #[rstest] + #[case::lf( + b": keepalive\n\nevent: delta\ndata: a\ndata: b\n\n", + b"event: delta\ndata: a\ndata: b\n\n" + )] + #[case::crlf( + b": keepalive\r\n\r\nevent: delta\r\ndata: a\r\ndata: b\r\n\r\n", + b"event: delta\r\ndata: a\r\ndata: b\r\n\r\n" + )] + #[case::cr( + b": keepalive\r\revent: delta\rdata: a\rdata: b\r\r", + b"event: delta\rdata: a\rdata: b\r\r" + )] + #[case::bom(b"\xEF\xBB\xBFdata: \xe2\x82\xac\n\n", b"data: \xe2\x82\xac\n\n")] + #[case::empty_data(b"event: ignored\n\ndata:\n\n", b"data:\n\n")] + #[case::extensions( + b"id: 1\nretry: 10\nvendor: x\ndata: a\n\n", + b"id: 1\nretry: 10\nvendor: x\ndata: a\n\n" + )] + #[tokio::test] + async fn blocks_carry_the_wire_bytes_and_the_event_they_dispatch( + #[case] wire: &[u8], + #[case] expected: &[u8], + #[values(1, 3, 4096)] size: usize, + ) { + let blocks = blocks(wire, size).await.unwrap(); + let expected_events = decode_all(expected).unwrap(); + assert_eq!(expected_events.len(), 1); + assert_eq!( + blocks, + vec![SseBlock { + bytes: Bytes::copy_from_slice(expected), + event: expected_events[0].clone(), + }] + ); + assert_eq!(expected_events, decode_all(wire).unwrap()); + } + + #[rstest] + #[case::empty(b"", vec![])] + #[case::unterminated(b"data: partial\n", vec![])] + #[case::partial_tail(b"data: complete\n\ndata: partial\n", vec![&b"data: complete\n\n"[..]])] + #[case::multiple(b"data: first\n\ndata: second\n\n", vec![&b"data: first\n\n"[..], &b"data: second\n\n"[..]])] + #[case::trailing_comment(b"data: first\n\n: bye\n\n", vec![&b"data: first\n\n"[..]])] + #[tokio::test] + async fn blocks_dispatch_complete_events_and_drop_a_cut_tail( + #[case] wire: &[u8], + #[case] expected: Vec<&[u8]>, + ) { + let mut blocks = Box::pin(frames( + stream::iter(pieces(wire, 1).into_iter().map(Ok::<_, std::io::Error>)), + SseBlockCodec::default(), + )); + for expected in expected { + assert_eq!(blocks.next().await.unwrap().unwrap().bytes, expected); + } + assert!(blocks.next().await.is_none()); + } + + #[rstest] + fn bulk_decode_preserves_every_field() { + let events = decode_all(b"id: 1\nretry: 20\nevent: delta\ndata: a\ndata: b\n\n").unwrap(); + assert_eq!( + events, + vec![SseEvent { + event: Some("delta".into()), + data: "a\nb".into(), + id: Some("1".into()), + retry: Some(20), + }] + ); + } + + #[rstest] + #[case::only_a_tail(b"data: partial\n", vec![])] + #[case::tail_after_an_event(b"data: complete\n\ndata: partial\n", vec!["complete"])] + #[case::no_terminator_at_all(b"data: partial", vec![])] + fn bulk_decode_drops_the_event_a_stream_ends_inside( + #[case] wire: &[u8], + #[case] expected: Vec<&str>, + ) { + let data: Vec = decode_all(wire) + .unwrap() + .into_iter() + .map(|e| e.data) + .collect(); + assert_eq!(data, expected); + } + + #[rstest] + #[case::trailing_blank_lines(b"data: x\n\n\n\n")] + #[case::trailing_comment(b"data: x\n\n: bye\n\n")] + #[case::lone_cr_at_eof(b"data: x\r\r")] + fn bulk_decode_accepts_a_stream_that_ends_between_events(#[case] wire: &[u8]) { + assert_eq!(decode_all(wire).unwrap().len(), 1); + } + + #[rstest] + fn bulk_decode_rejects_invalid_utf8() { + assert!(matches!( + decode_all(b"data: \xff\n\n"), + Err(SseError::InvalidUtf8(_)) + )); + } + + #[rstest] + #[tokio::test] + async fn block_framing_preserves_transport_errors() { + let input = stream::iter([ + Ok(Bytes::from_static(b"data: first\n\n")), + Err(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + "reset", + )), + ]); + let mut blocks = Box::pin(frames(input, SseBlockCodec::default())); + assert_eq!( + blocks.next().await.unwrap().unwrap().bytes, + Bytes::from_static(b"data: first\n\n") + ); + let Some(Err(SseError::Body(error))) = blocks.next().await else { + panic!("expected body error"); + }; + let cause = error + .get_ref() + .unwrap() + .downcast_ref::() + .unwrap(); + assert_eq!(cause.kind(), std::io::ErrorKind::ConnectionReset); + assert!(blocks.next().await.is_none()); + } + + #[rstest] + #[tokio::test] + async fn invalid_utf8_stops_block_framing_after_prior_events() { + let mut blocks = Box::pin(frames( + stream::iter( + pieces(b"data: first\n\ndata: \xff\n\n", 1) + .into_iter() + .map(Ok::<_, std::io::Error>), + ), + SseBlockCodec::default(), + )); + assert_eq!( + blocks.next().await.unwrap().unwrap().bytes, + Bytes::from_static(b"data: first\n\n") + ); + assert!(matches!( + blocks.next().await, + Some(Err(SseError::InvalidUtf8(_))) + )); + assert!(blocks.next().await.is_none()); + } +} diff --git a/litellm-rust/crates/framer/tests/sse.rs b/litellm-rust/crates/framer/tests/sse.rs index bf7a56525cf..7bf2bfac29f 100644 --- a/litellm-rust/crates/framer/tests/sse.rs +++ b/litellm-rust/crates/framer/tests/sse.rs @@ -97,12 +97,11 @@ async fn dispatch_follows_the_data_buffer(#[case] wire: &[u8], #[case] expected: } #[rstest] -#[case::unterminated_single(b"data: partial\n", vec![])] -#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])] #[case::lone_cr_terminates_at_eof(b"data: x\r\r", vec![event(None, "x")])] -#[case::lone_cr_line_then_eof(b"data: x\r", vec![])] +#[case::trailing_blank_lines(b"data: x\n\n\n", vec![event(None, "x")])] +#[case::trailing_comment(b"data: x\n\n: bye\n\n", vec![event(None, "x")])] #[tokio::test] -async fn eof_dispatches_only_terminated_events( +async fn eof_between_events_ends_the_stream_cleanly( #[case] wire: &[u8], #[case] expected: Vec, ) { @@ -112,6 +111,19 @@ async fn eof_dispatches_only_terminated_events( ); } +#[rstest] +#[case::unterminated_single(b"data: partial\n", vec![])] +#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])] +#[case::lone_cr_line_then_eof(b"data: x\r", vec![])] +#[case::partial_bom(b"\xEF\xBB", vec![])] +#[tokio::test] +async fn eof_inside_an_event_drops_it_after_the_terminated_events( + #[case] wire: &[u8], + #[case] expected: Vec, +) { + assert_eq!(collect(every(wire, 1)).await.unwrap(), expected); +} + #[rstest] #[case::inside_the_first_line(vec![&b"data: a\r"[..], &b"\ndata: b\r\n\r\n"[..]])] #[case::inside_the_blank_line(vec![&b"data: a\r\ndata: b\r\n\r"[..], &b"\n"[..]])] diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 178f96c4c33..c4032baa98f 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -7,7 +7,7 @@ use litellm_http::transport::Error as TransportError; use litellm_llms::base_llm::{ auth::{Authenticated, resolve_auth}, messages::{ - streaming::{ByteStream, StreamDecoder, encode_anthropic_sse}, + streaming::{self, ByteStream, StreamDecoder}, transformation::BaseMessagesConfig, }, }; @@ -245,7 +245,7 @@ fn decoded_chunks( let Some(event) = events.try_next().await? else { return Ok(None); }; - let chunk = encode_anthropic_sse(&event)?; + let chunk = streaming::encode(&event)?; log_chunk(provider, "client_response", &chunk); Ok(Some((chunk, events))) }) @@ -271,7 +271,6 @@ fn log_chunk(provider: &str, stage: &str, data: &bytes::Bytes) { #[cfg(test)] mod tests { - use litellm_llms::base_llm::messages::streaming::anthropic_sse_event_stream; use rstest::rstest; use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any}; @@ -300,7 +299,7 @@ mod tests { .await .unwrap(); let MessagesCallResponse::Stream { mut chunks, .. } = - streaming_response(response, Some(anthropic_sse_event_stream), "test") + streaming_response(response, Some(streaming::event_stream), "test") else { panic!("a streaming response returns chunks"); }; diff --git a/litellm-rust/crates/inference-responses/src/websocket.rs b/litellm-rust/crates/inference-responses/src/websocket.rs index 5d088bd2d63..213a9a95831 100644 --- a/litellm-rust/crates/inference-responses/src/websocket.rs +++ b/litellm-rust/crates/inference-responses/src/websocket.rs @@ -2,7 +2,6 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; use futures_util::{SinkExt, StreamExt}; use litellm_http::websocket::{UpstreamWebSocket, connect_upstream}; -use litellm_llms_types::formats::responses::streaming_websocket::ResponsesWsEventType; use tokio::sync::Mutex; use tokio_tungstenite::tungstenite::{ Message, @@ -12,17 +11,6 @@ use tokio_tungstenite::tungstenite::{ use super::Error; -pub fn is_terminal_event(event_type: &ResponsesWsEventType) -> bool { - matches!( - event_type, - ResponsesWsEventType::ResponseCreated - | ResponsesWsEventType::ResponseCompleted - | ResponsesWsEventType::ResponseFailed - | ResponsesWsEventType::ResponseIncomplete - | ResponsesWsEventType::Error - ) -} - #[derive(Clone)] pub struct ResponsesWebSocketConnection { socket: Arc>>, diff --git a/litellm-rust/crates/llms-types/src/formats/chat_completions/streaming.rs b/litellm-rust/crates/llms-types/src/formats/chat_completions/streaming.rs index 520b064d10d..f76efb2d132 100644 --- a/litellm-rust/crates/llms-types/src/formats/chat_completions/streaming.rs +++ b/litellm-rust/crates/llms-types/src/formats/chat_completions/streaming.rs @@ -33,16 +33,82 @@ pub struct ChatCompletionStreamingChoice { pub logprobs: Option, } +/// Compatible hosts omit `id`, `created` and `object` on chunks, and Python's +/// `ModelResponseStream` fills them in, so none of the envelope fields is required here. #[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Default)] pub struct ChatCompletionChunk { - pub id: String, - pub created: u64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub created: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub model: Option, - pub object: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub object: Option, + #[serde(default)] pub choices: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] pub usage: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub provider_specific_fields: Option>, + #[serde(flatten)] + pub extra: Map, +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::{Value, json}; + + use super::*; + + #[rstest] + #[case::openai(json!({ + "id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "m", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi", "future": 1}, "finish_reason": null, "logprobs": null}], + "usage": null + }))] + #[case::usage_only(json!({ + "id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "m", "choices": [], + "usage": { + "prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3, + "prompt_tokens_details": {"cached_tokens": 0, "cache_creation_tokens": 0, "text_tokens": 1} + } + }))] + #[case::bare_host(json!({"choices": [{"index": 0, "delta": {"content": "hi"}}]}))] + #[case::host_extensions(json!({"id": "x", "choices": [], "citations": ["a"], "future": {"k": null}}))] + fn chunks_round_trip_without_losing_fields(#[case] wire: Value) { + let chunk: ChatCompletionChunk = serde_json::from_value(wire.clone()).unwrap(); + assert_eq!(serde_json::to_value(chunk).unwrap(), without_nulls(wire)); + } + + /// Typed optional fields serialize by omission; extension fields keep their nulls. + fn without_nulls(value: Value) -> Value { + let Value::Object(fields) = value else { + return value; + }; + Value::Object( + fields + .into_iter() + .filter(|(_, value)| !value.is_null()) + .map(|(key, value)| match (key.as_str(), value) { + ("choices", Value::Array(choices)) => ( + key, + Value::Array(choices.into_iter().map(without_nulls).collect()), + ), + (_, value) => (key, value), + }) + .collect(), + ) + } + + #[rstest] + #[case::choices_not_a_list(json!({"choices": {}}))] + #[case::choice_without_delta(json!({"choices": [{"index": 0}]}))] + #[case::created_not_a_number(json!({"created": "now", "choices": []}))] + #[case::non_object(json!([]))] + fn malformed_chunks_are_rejected(#[case] wire: Value) { + assert!(serde_json::from_value::(wire).is_err()); + } } diff --git a/litellm-rust/crates/llms-types/src/formats/messages/streaming.rs b/litellm-rust/crates/llms-types/src/formats/messages/streaming.rs index 77c9d8cb742..b2f70ee3753 100644 --- a/litellm-rust/crates/llms-types/src/formats/messages/streaming.rs +++ b/litellm-rust/crates/llms-types/src/formats/messages/streaming.rs @@ -109,7 +109,9 @@ pub struct MessagesStreamError { } #[macro_rules_attribute::apply(crate::wire_type)] +#[derive(strum::AsRefStr)] #[serde(tag = "type", rename_all = "snake_case")] +#[strum(serialize_all = "snake_case")] pub enum MessagesStreamEvent { MessageStart { message: MessagesStreamMessage, @@ -175,6 +177,23 @@ mod tests { assert_eq!(serde_json::to_value(event).unwrap(), wire); } + #[rstest] + #[case::start(json!({"type": "message_start", "message": {"id": "m", "type": "message", "role": "assistant", "model": "x", "content": [], "stop_reason": null, "stop_sequence": null, "usage": {}}}))] + #[case::block_start(json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}))] + #[case::block_delta(json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "a"}}))] + #[case::block_stop(json!({"type": "content_block_stop", "index": 0}))] + #[case::delta(json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}}))] + #[case::stop(json!({"type": "message_stop"}))] + #[case::ping(json!({"type": "ping"}))] + #[case::error(json!({"type": "error", "error": {"type": "overloaded_error", "message": "busy"}}))] + fn the_event_name_is_the_serialized_type_tag(#[case] wire: Value) { + let event: MessagesStreamEvent = serde_json::from_value(wire.clone()).unwrap(); + assert_eq!( + Some(event.as_ref()), + wire.get("type").and_then(Value::as_str) + ); + } + #[rstest] #[case::invalid_index(json!({"type": "content_block_stop", "index": "zero"}))] #[case::missing_delta(json!({"type": "content_block_delta", "index": 0}))] diff --git a/litellm-rust/crates/llms-types/src/formats/responses/AGENTS.md b/litellm-rust/crates/llms-types/src/formats/responses/AGENTS.md index 606c1862ad7..0584b50c000 100644 --- a/litellm-rust/crates/llms-types/src/formats/responses/AGENTS.md +++ b/litellm-rust/crates/llms-types/src/formats/responses/AGENTS.md @@ -1,7 +1,7 @@ # rules - Responses originated at OpenAI and is now offered by other hosts and gateways. Types here describe the format, not OpenAI's deployment of it -- Served over HTTP and over the WebSocket mode. `ResponsesWsEvent` is a wire event and lives here +- Served over HTTP and over the WebSocket mode. `ResponsesStreamEvent` and `ResponsesEventType` in `streaming.rs` are the wire event for both - `ResponsesWsTransformResult` wraps a provider transformation's output, so it lives in `llms/src/base_llm/responses/` - Hosts that accept only part of the format, or add output item types, are handled in `llms/src//responses/`. Keep unknown output items passing through instead of rejecting them - Responses-to-Chat emulation and model-specific parameter rewriting are transformations, not format rules diff --git a/litellm-rust/crates/llms-types/src/formats/responses/mod.rs b/litellm-rust/crates/llms-types/src/formats/responses/mod.rs index 7fb0324dd43..c72cb755fea 100644 --- a/litellm-rust/crates/llms-types/src/formats/responses/mod.rs +++ b/litellm-rust/crates/llms-types/src/formats/responses/mod.rs @@ -1,5 +1,6 @@ mod output; mod response; +pub mod streaming; pub mod streaming_websocket; pub use output::*; diff --git a/litellm-rust/crates/llms-types/src/formats/responses/streaming.rs b/litellm-rust/crates/llms-types/src/formats/responses/streaming.rs new file mode 100644 index 00000000000..32631ebfeb5 --- /dev/null +++ b/litellm-rust/crates/llms-types/src/formats/responses/streaming.rs @@ -0,0 +1,305 @@ +use serde_json::{Map, Value}; + +/// Every event name on https://developers.openai.com/api/reference/resources/responses/streaming-events +/// as of 2026-10-11, plus `response.create` from the WebSocket mode. Names OpenAI adds later ride in +/// `Other` untouched. +#[derive( + Clone, + Debug, + PartialEq, + Eq, + strum::AsRefStr, + strum::EnumIter, + strum::EnumString, + strum::Display, + serde_with::DeserializeFromStr, + serde_with::SerializeDisplay, +)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))] +pub enum ResponsesEventType { + #[strum(serialize = "response.create")] + ResponseCreate, + #[strum(serialize = "response.created")] + ResponseCreated, + #[strum(serialize = "response.queued")] + ResponseQueued, + #[strum(serialize = "response.in_progress")] + ResponseInProgress, + #[strum(serialize = "response.completed")] + ResponseCompleted, + #[strum(serialize = "response.failed")] + ResponseFailed, + #[strum(serialize = "response.incomplete")] + ResponseIncomplete, + #[strum(serialize = "error")] + Error, + #[strum(serialize = "response.audio.delta")] + AudioDelta, + #[strum(serialize = "response.audio.done")] + AudioDone, + #[strum(serialize = "response.audio.transcript.delta")] + AudioTranscriptDelta, + #[strum(serialize = "response.audio.transcript.done")] + AudioTranscriptDone, + #[strum(serialize = "response.code_interpreter_call.completed")] + CodeInterpreterCallCompleted, + #[strum(serialize = "response.code_interpreter_call.in_progress")] + CodeInterpreterCallInProgress, + #[strum(serialize = "response.code_interpreter_call.interpreting")] + CodeInterpreterCallInterpreting, + #[strum(serialize = "response.code_interpreter_call_code.delta")] + CodeInterpreterCallCodeDelta, + #[strum(serialize = "response.code_interpreter_call_code.done")] + CodeInterpreterCallCodeDone, + #[strum(serialize = "response.compaction.compacting")] + CompactionCompacting, + #[strum(serialize = "response.content_part.added")] + ContentPartAdded, + #[strum(serialize = "response.content_part.done")] + ContentPartDone, + #[strum(serialize = "response.custom_tool_call_input.delta")] + CustomToolCallInputDelta, + #[strum(serialize = "response.custom_tool_call_input.done")] + CustomToolCallInputDone, + #[strum(serialize = "response.file_search_call.completed")] + FileSearchCallCompleted, + #[strum(serialize = "response.file_search_call.in_progress")] + FileSearchCallInProgress, + #[strum(serialize = "response.file_search_call.searching")] + FileSearchCallSearching, + #[strum(serialize = "response.function_call_arguments.delta")] + FunctionCallArgumentsDelta, + #[strum(serialize = "response.function_call_arguments.done")] + FunctionCallArgumentsDone, + #[strum(serialize = "response.image_generation_call.completed")] + ImageGenerationCallCompleted, + #[strum(serialize = "response.image_generation_call.generating")] + ImageGenerationCallGenerating, + #[strum(serialize = "response.image_generation_call.in_progress")] + ImageGenerationCallInProgress, + #[strum(serialize = "response.image_generation_call.partial_image")] + ImageGenerationCallPartialImage, + #[strum(serialize = "response.mcp_call.completed")] + McpCallCompleted, + #[strum(serialize = "response.mcp_call.failed")] + McpCallFailed, + #[strum(serialize = "response.mcp_call.in_progress")] + McpCallInProgress, + #[strum(serialize = "response.mcp_call_arguments.delta")] + McpCallArgumentsDelta, + #[strum(serialize = "response.mcp_call_arguments.done")] + McpCallArgumentsDone, + #[strum(serialize = "response.mcp_list_tools.completed")] + McpListToolsCompleted, + #[strum(serialize = "response.mcp_list_tools.failed")] + McpListToolsFailed, + #[strum(serialize = "response.mcp_list_tools.in_progress")] + McpListToolsInProgress, + #[strum(serialize = "response.output_item.added")] + OutputItemAdded, + #[strum(serialize = "response.output_item.done")] + OutputItemDone, + #[strum(serialize = "response.output_text.annotation.added")] + OutputTextAnnotationAdded, + #[strum(serialize = "response.output_text.delta")] + OutputTextDelta, + #[strum(serialize = "response.output_text.done")] + OutputTextDone, + #[strum(serialize = "response.reasoning_summary_part.added")] + ReasoningSummaryPartAdded, + #[strum(serialize = "response.reasoning_summary_part.done")] + ReasoningSummaryPartDone, + #[strum(serialize = "response.reasoning_summary_text.delta")] + ReasoningSummaryTextDelta, + #[strum(serialize = "response.reasoning_summary_text.done")] + ReasoningSummaryTextDone, + #[strum(serialize = "response.reasoning_text.delta")] + ReasoningTextDelta, + #[strum(serialize = "response.reasoning_text.done")] + ReasoningTextDone, + #[strum(serialize = "response.refusal.delta")] + RefusalDelta, + #[strum(serialize = "response.refusal.done")] + RefusalDone, + #[strum(serialize = "response.shell_call_command.added")] + ShellCallCommandAdded, + #[strum(serialize = "response.shell_call_command.delta")] + ShellCallCommandDelta, + #[strum(serialize = "response.shell_call_command.done")] + ShellCallCommandDone, + #[strum(serialize = "response.shell_call_output_content.delta")] + ShellCallOutputContentDelta, + #[strum(serialize = "response.shell_call_output_content.done")] + ShellCallOutputContentDone, + #[strum(serialize = "response.web_search_call.completed")] + WebSearchCallCompleted, + #[strum(serialize = "response.web_search_call.in_progress")] + WebSearchCallInProgress, + #[strum(serialize = "response.web_search_call.searching")] + WebSearchCallSearching, + #[strum(default, transparent)] + Other(String), +} + +impl ResponsesEventType { + /// The last event of a response. `response.completed`, `response.incomplete` and + /// `response.failed` each close the response and carry its final body; `error` aborts it. + pub fn is_terminal(&self) -> bool { + matches!( + self, + Self::ResponseCompleted | Self::ResponseFailed | Self::ResponseIncomplete | Self::Error + ) + } + + /// A terminal event after which the response is not a success. + pub fn is_failure(&self) -> bool { + matches!( + self, + Self::ResponseFailed | Self::ResponseIncomplete | Self::Error + ) + } +} + +#[macro_rules_attribute::apply(crate::wire_type)] +pub struct ResponsesStreamEvent { + #[serde(rename = "type")] + pub event_type: ResponsesEventType, + #[serde(flatten)] + pub data: Map, +} + +impl ResponsesStreamEvent { + pub fn model(&self) -> Option<&str> { + let model = self.data.get("model").and_then(Value::as_str); + if model.is_some() { + return model; + } + self.data + .get("response") + .and_then(Value::as_object) + .and_then(|response| response.get("model")) + .and_then(Value::as_str) + } + + pub fn is_response_create(&self) -> bool { + self.event_type == ResponsesEventType::ResponseCreate + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::{Value, json}; + use strum::IntoEnumIterator; + + use super::*; + + #[rstest] + fn every_named_event_type_round_trips() { + let named: Vec<_> = ResponsesEventType::iter() + .filter(|event_type| !matches!(event_type, ResponsesEventType::Other(_))) + .collect(); + assert!(named.len() > 50); + for event_type in named { + let wire = event_type.to_string(); + assert_eq!(wire, event_type.as_ref()); + assert_eq!(serde_json::to_value(&event_type).unwrap(), json!(wire)); + assert_eq!(wire.parse::().unwrap(), event_type); + assert_eq!( + serde_json::from_value::(json!(wire)).unwrap(), + event_type + ); + } + } + + #[rstest] + #[case::unknown("future.event", ResponsesEventType::Other("future.event".into()))] + #[case::empty("", ResponsesEventType::Other(String::new()))] + #[case::case_sensitive( + "Response.Completed", + ResponsesEventType::Other("Response.Completed".into()) + )] + #[case::escaped("future\"\\\n", ResponsesEventType::Other("future\"\\\n".into()))] + fn unknown_event_types_round_trip_through_other( + #[case] wire: &str, + #[case] expected: ResponsesEventType, + ) { + assert_eq!( + serde_json::from_value::(json!(wire)).unwrap(), + expected + ); + assert_eq!(serde_json::to_value(&expected).unwrap(), json!(wire)); + assert_eq!(expected.as_ref(), wire); + assert_eq!(expected.to_string(), wire); + } + + /// Python's `ResponsesAPIStreamingIterator` stores the response on `response.completed`, + /// `response.incomplete` and `response.failed`, and raises on `response.failed` and `error`. + #[rstest] + #[case::completed(ResponsesEventType::ResponseCompleted, true, false)] + #[case::incomplete(ResponsesEventType::ResponseIncomplete, true, true)] + #[case::failed(ResponsesEventType::ResponseFailed, true, true)] + #[case::error(ResponsesEventType::Error, true, true)] + #[case::create(ResponsesEventType::ResponseCreate, false, false)] + #[case::created(ResponsesEventType::ResponseCreated, false, false)] + #[case::in_progress(ResponsesEventType::ResponseInProgress, false, false)] + #[case::delta(ResponsesEventType::OutputTextDelta, false, false)] + #[case::item_done(ResponsesEventType::OutputItemDone, false, false)] + #[case::unknown(ResponsesEventType::Other("response.future".into()), false, false)] + #[case::tool_failed(ResponsesEventType::McpCallFailed, false, false)] + fn terminal_and_failure_events_match_the_python_iterator( + #[case] event_type: ResponsesEventType, + #[case] terminal: bool, + #[case] failure: bool, + ) { + assert_eq!(event_type.is_terminal(), terminal); + assert_eq!(event_type.is_failure(), failure); + } + + #[rstest] + #[case::number("17")] + #[case::boolean("true")] + #[case::null("null")] + #[case::array("[]")] + #[case::object("{}")] + fn event_types_reject_non_strings(#[case] wire: &str) { + assert!(serde_json::from_str::(wire).is_err()); + } + + #[rstest] + #[case::text(json!({"type":"response.output_text.delta","delta":"hello","sequence_number":3,"future":null}))] + #[case::tool(json!({"type":"response.function_call_arguments.delta","delta":"{\"x\":","item_id":"fc_1","output_index":0}))] + #[case::unknown(json!({"type":"future.event","payload":{"items":[1,null]},"future":true}))] + #[case::terminal(json!({"type":"response.completed","response":{"id":"r","output":[],"future":null}}))] + fn events_preserve_every_payload_field(#[case] wire: Value) { + let event: ResponsesStreamEvent = serde_json::from_value(wire.clone()).unwrap(); + assert_eq!(serde_json::to_value(event).unwrap(), wire); + } + + #[rstest] + #[case::missing_type(json!({"delta":"hello"}))] + #[case::numeric_type(json!({"type":7}))] + #[case::null_type(json!({"type":null}))] + #[case::non_object(json!([]))] + fn malformed_events_are_rejected(#[case] wire: Value) { + assert!(serde_json::from_value::(wire).is_err()); + } + + #[rstest] + #[case::flat(json!({"type":"response.create","model":"model-a"}), Some("model-a"))] + #[case::nested(json!({"type":"response.create","response":{"model":"model-b"}}), Some("model-b"))] + #[case::none(json!({"type":"response.created","response":{}}), None)] + fn model_reads_flat_and_nested_shapes(#[case] wire: Value, #[case] expected: Option<&str>) { + let event: ResponsesStreamEvent = serde_json::from_value(wire).unwrap(); + assert_eq!(event.model(), expected); + } + + #[cfg(feature = "schema")] + #[rstest] + fn event_type_schema_accepts_unknown_strings() { + let schema = schemars::schema_for!(ResponsesEventType).to_value(); + assert_eq!(schema.get("type"), Some(&json!("string"))); + assert!(schema.get("enum").is_none()); + } +} diff --git a/litellm-rust/crates/llms-types/src/formats/responses/streaming_websocket.rs b/litellm-rust/crates/llms-types/src/formats/responses/streaming_websocket.rs index 741540b5741..f1331de4e46 100644 --- a/litellm-rust/crates/llms-types/src/formats/responses/streaming_websocket.rs +++ b/litellm-rust/crates/llms-types/src/formats/responses/streaming_websocket.rs @@ -2,68 +2,6 @@ use super::ResponsesOutputItem; use crate::recognized::Recognized; use serde_json::{Map, Value}; -#[derive( - Clone, - Debug, - PartialEq, - Eq, - strum::AsRefStr, - strum::EnumString, - strum::Display, - serde_with::DeserializeFromStr, - serde_with::SerializeDisplay, -)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[cfg_attr(feature = "schema", schemars(with = "String"))] -pub enum ResponsesWsEventType { - #[strum(serialize = "response.create")] - ResponseCreate, - #[strum(serialize = "response.created")] - ResponseCreated, - #[strum(serialize = "response.completed")] - ResponseCompleted, - #[strum(serialize = "response.failed")] - ResponseFailed, - #[strum(serialize = "response.incomplete")] - ResponseIncomplete, - #[strum(serialize = "error")] - Error, - #[strum(default, transparent)] - Other(String), -} - -impl ResponsesWsEventType { - pub fn as_str(&self) -> &str { - self.as_ref() - } -} - -#[macro_rules_attribute::apply(crate::wire_type)] -pub struct ResponsesWsEvent { - #[serde(rename = "type")] - pub event_type: ResponsesWsEventType, - #[serde(flatten)] - pub data: Map, -} - -impl ResponsesWsEvent { - pub fn model(&self) -> Option<&str> { - let model = self.data.get("model").and_then(Value::as_str); - if model.is_some() { - return model; - } - self.data - .get("response") - .and_then(Value::as_object) - .and_then(|response| response.get("model")) - .and_then(Value::as_str) - } - - pub fn is_response_create(&self) -> bool { - self.event_type == ResponsesWsEventType::ResponseCreate - } -} - #[macro_rules_attribute::apply(crate::wire_type)] #[derive(Eq)] pub struct ResponsesErrorFrame { @@ -112,8 +50,7 @@ mod tests { use crate::formats::responses::ResponsesOutputItem; use crate::{ - formats::responses::streaming_websocket::{ResponsesEventResponse, ResponsesWsEventType}, - recognized::Recognized, + formats::responses::streaming_websocket::ResponsesEventResponse, recognized::Recognized, test_support::*, }; @@ -134,66 +71,6 @@ mod tests { ); } - #[rstest] - #[case::flat(serde_json::json!({"type":"response.create","model":"model-a"}), Some("model-a"))] - #[case::nested(serde_json::json!({"type":"response.create","response":{"model":"model-b"}}), Some("model-b"))] - fn model_reads_flat_and_nested_create_shapes( - #[case] payload: serde_json::Value, - #[case] expected: Option<&str>, - ) { - let event: ResponsesWsEvent = serde_json::from_value(payload).expect("valid event"); - assert_eq!(event.model(), expected); - } - - #[rstest] - #[case::create("response.create", ResponsesWsEventType::ResponseCreate)] - #[case::created("response.created", ResponsesWsEventType::ResponseCreated)] - #[case::completed("response.completed", ResponsesWsEventType::ResponseCompleted)] - #[case::failed("response.failed", ResponsesWsEventType::ResponseFailed)] - #[case::incomplete("response.incomplete", ResponsesWsEventType::ResponseIncomplete)] - #[case::error("error", ResponsesWsEventType::Error)] - #[case::unknown( - "response.output_text.delta", - ResponsesWsEventType::Other("response.output_text.delta".to_string()) - )] - #[case::empty("", ResponsesWsEventType::Other(String::new()))] - #[case::case_sensitive( - "Response.Completed", - ResponsesWsEventType::Other("Response.Completed".into()) - )] - #[case::escaped("future\"\\\n", ResponsesWsEventType::Other("future\"\\\n".into()))] - fn websocket_event_type_round_trips( - #[case] wire: &str, - #[case] expected: ResponsesWsEventType, - ) { - let serialized = serde_json::to_string(&expected).unwrap(); - assert_eq!(serialized, serde_json::to_string(wire).unwrap()); - assert_eq!( - serde_json::from_str::(&serialized).unwrap(), - expected - ); - } - - #[rstest] - #[case::number("17")] - #[case::boolean("true")] - #[case::null("null")] - #[case::array("[]")] - #[case::object("{}")] - fn websocket_event_type_rejects_non_strings(#[case] wire: &str) { - assert!(serde_json::from_str::(wire).is_err()); - } - - #[cfg(feature = "schema")] - #[rstest] - fn websocket_event_type_schema_is_open_string() { - let schema = schemars::schema_for!(ResponsesWsEventType); - assert_eq!( - schema.to_value().get("type"), - Some(&serde_json::json!("string")) - ); - } - #[rstest] fn nested_event_response_exposes_typed_output_and_preserves_extensions() { let wire = json!({ diff --git a/litellm-rust/crates/llms/AGENTS.md b/litellm-rust/crates/llms/AGENTS.md index ec0f6ffe95a..3bcc65cc123 100644 --- a/litellm-rust/crates/llms/AGENTS.md +++ b/litellm-rust/crates/llms/AGENTS.md @@ -32,6 +32,13 @@ - Provider-wide docs (auth, errors, regions): `src//` - One host's endpoint docs: `src///` +## Streaming + +- Each format's `src/base_llm//streaming.rs` exposes the same three names: `event_stream` (bytes to typed events), `decode` (a whole stream to typed events) and `encode` (one typed event to the SSE frame LiteLLM's proxy sends clients). Callers use the module path, never a vendor prefix +- SSE framing, `[DONE]` and the generic helpers live in `src/base_llm/sse.rs`. Messages frames are `event:` plus `data:`; Chat Completions and Responses frames are bare `data:`, as Python's proxy writes them +- Decoding is strict: a malformed event, or `[DONE]` on a Messages stream, is `InvalidResponse` and ends the stream. Python never fails a stream this way; the divergence is deliberate and `src/base_llm/sse.rs` lists what Python does per format. A stream cut inside an event drops that event, as the SSE spec requires +- What ends or fails a stream belongs on the format's event type (`is_terminal`, `is_failure`). The cache's success check in `inference/src/caching.rs` still matches type strings; move it onto the typed events rather than adding new string matches + ## Python pairs - Python paths identify counterparts but do not dictate Rust module names or class hierarchy. Preserve behavior and concepts, not structure diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index ce759f6299e..91baa247ac8 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -30,7 +30,7 @@ use crate::{ Unsupported, ValidatedEnvironment, unsupported_message, unsupported_param, }, }, - messages::streaming::anthropic_sse_event_stream, + messages::streaming::event_stream as messages_event_stream, }, }; @@ -196,7 +196,7 @@ impl BaseConfig for AnthropicConfig { fn model_response_iterator(&self, shape: StreamShape) -> Option { Some(ChatStream::new( - anthropic_sse_event_stream, + messages_event_stream, ModelResponseIterator::new(shape), )) } diff --git a/litellm-rust/crates/llms/src/base_llm/chat/streaming.rs b/litellm-rust/crates/llms/src/base_llm/chat/streaming.rs index 2b4a8d084e6..2faa5699751 100644 --- a/litellm-rust/crates/llms/src/base_llm/chat/streaming.rs +++ b/litellm-rust/crates/llms/src/base_llm/chat/streaming.rs @@ -1,15 +1,32 @@ use std::collections::HashMap; +use bytes::Bytes; use futures_util::{StreamExt, stream::BoxStream}; use litellm_llms_types::formats::chat_completions::ChatCompletionChunk; use crate::{ Error, - base_llm::base_model_iterator::{ByteStream, StreamError, StreamTransformer, transform_stream}, + base_llm::{ + base_model_iterator::{ByteStream, StreamError, StreamTransformer, transform_stream}, + sse::{self, is_done}, + }, }; pub type ChatChunkStream = BoxStream<'static, Result>; +/// OpenAI-compatible SSE: bare `data:` frames closed by `data: [DONE]`. +pub fn event_stream(bytes: ByteStream) -> ChatChunkStream { + sse::event_stream(bytes, is_done) +} + +pub fn decode(wire: &[u8]) -> Result, Error> { + sse::decode(wire, is_done) +} + +pub fn encode(chunk: &ChatCompletionChunk) -> Result { + sse::encode_data(chunk) +} + /// What Python's `map_openai_params` decides about the stream and `completion` /// hands to `ModelResponseIterator`: it is settled while the request is built, /// never re-derived from the body. @@ -52,3 +69,109 @@ impl ChatStream { (self.run)(bytes) } } + +#[cfg(test)] +mod tests { + use futures_util::{StreamExt, TryStreamExt, stream}; + use rstest::rstest; + use serde_json::{Value, json}; + + use super::*; + + fn in_pieces(wire: &[u8], size: usize) -> ByteStream { + let pieces: Vec = wire.chunks(size).map(Bytes::copy_from_slice).collect(); + stream::iter(pieces.into_iter().map(Ok)).boxed() + } + + #[rstest] + #[case::role(json!({"id":"c","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{"role":"assistant","content":""}}]}))] + #[case::text(json!({"id":"c","choices":[{"index":0,"delta":{"content":"hi\nthere"}}]}))] + #[case::tool(json!({"id":"c","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"t","type":"function","function":{"name":"f","arguments":"{\"x\":"}}]}}]}))] + #[case::finish(json!({"id":"c","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}))] + #[case::usage(json!({"id":"c","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3,"prompt_tokens_details":{"cached_tokens":0,"cache_creation_tokens":0,"text_tokens":1}}}))] + #[case::host_extension(json!({"id":"c","choices":[],"citations":[{"url":"u"}]}))] + fn encoded_chunks_are_bare_data_frames_that_decode_back(#[case] wire: Value) { + let chunk: ChatCompletionChunk = serde_json::from_value(wire.clone()).unwrap(); + let encoded = encode(&chunk).unwrap(); + assert!(encoded.starts_with(b"data: ")); + assert!(encoded.ends_with(b"\n\n")); + assert_eq!( + serde_json::from_slice::(&encoded[b"data: ".len()..encoded.len() - 2]).unwrap(), + wire + ); + assert_eq!(decode(&encoded).unwrap(), vec![chunk]); + } + + #[rstest] + #[case::lf("\n")] + #[case::crlf("\r\n")] + #[case::cr("\r")] + #[tokio::test] + async fn streams_split_anywhere_skip_comments_and_done( + #[case] newline: &str, + #[values(1, 5, 4096)] size: usize, + ) { + let wire = format!( + ": keepalive{newline}{newline}data: {{\"id\":\"c\",\"choices\":[{{\"index\":0,\"delta\":{{\"content\":\"a\"}}}}]}}{newline}{newline}data:{{\"id\":\"c\",\"choices\":[{{\"index\":0,\"delta\":{{}},\"finish_reason\":\"stop\"}}]}}{newline}{newline}data: [DONE]{newline}{newline}" + ); + let streamed = event_stream(in_pieces(wire.as_bytes(), size)) + .try_collect::>() + .await + .unwrap(); + assert_eq!(streamed, decode(wire.as_bytes()).unwrap()); + assert_eq!(streamed.len(), 2); + assert_eq!(streamed[0].choices[0].delta.content.as_deref(), Some("a")); + assert_eq!( + streamed[1].choices[0].finish_reason.as_deref(), + Some("stop") + ); + } + + #[rstest] + #[case::spaced(b"data: [DONE]\n\n")] + #[case::unspaced(b"data:[DONE]\n\n")] + #[case::trailing_space(b"data: [DONE] \n\n")] + fn every_done_spelling_closes_the_stream_without_a_chunk(#[case] wire: &[u8]) { + assert_eq!(decode(wire).unwrap(), vec![]); + } + + #[rstest] + #[case::invalid_json(b"data: {\n\n")] + #[case::choices_not_a_list(b"data: {\"choices\":{}}\n\n")] + #[case::invalid_utf8(b"data: \xff\n\n")] + #[tokio::test] + async fn malformed_streams_are_invalid_responses(#[case] wire: &[u8]) { + assert!(matches!(decode(wire), Err(Error::InvalidResponse(_)))); + let streamed = event_stream(in_pieces(wire, 1)).collect::>().await; + assert!(matches!( + streamed.last(), + Some(Err(Error::InvalidResponse(_))) | None + )); + } + + #[rstest] + #[tokio::test] + async fn a_body_error_ends_the_stream_after_delivered_chunks() { + let mut chunks = event_stream( + stream::iter([ + Ok(Bytes::from_static( + b"data: {\"id\":\"c\",\"choices\":[]}\n\n", + )), + Err(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + "reset", + )), + ]) + .boxed(), + ); + assert_eq!( + chunks.next().await.unwrap().unwrap().id.as_deref(), + Some("c") + ); + assert!(matches!( + chunks.next().await, + Some(Err(Error::InvalidResponse(_))) + )); + assert!(chunks.next().await.is_none()); + } +} diff --git a/litellm-rust/crates/llms/src/base_llm/messages/streaming.rs b/litellm-rust/crates/llms/src/base_llm/messages/streaming.rs index 07ccb2fb911..8c301e67de3 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/streaming.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/streaming.rs @@ -1,38 +1,32 @@ use bytes::Bytes; -use futures_util::{StreamExt, stream::BoxStream}; -use litellm_framer::{frames, sse::SseCodec}; use litellm_llms_types::formats::messages::streaming::MessagesStreamEvent; -use crate::Error; pub use crate::base_llm::base_model_iterator::ByteStream; +use crate::{Error, base_llm::sse}; -pub type EventStream = BoxStream<'static, Result>; +pub type EventStream = sse::EventStream; pub type StreamDecoder = fn(ByteStream) -> EventStream; -pub fn anthropic_sse_event_stream(bytes: ByteStream) -> EventStream { - Box::pin(frames(bytes, SseCodec::default()).map(|event| { - let event = event.map_err(|error| { - Error::InvalidResponse(crate::ErrorDetail::failed("stream framing", error)) - })?; - serde_json::from_str(&event.data).map_err(|error| { - Error::InvalidResponse(crate::ErrorDetail::invalid("Anthropic stream event", error)) - }) - })) +fn no_control_frames(_: &litellm_framer::sse::SseEvent) -> bool { + false } -pub fn encode_anthropic_sse(event: &MessagesStreamEvent) -> Result { - let data = serde_json::to_value(event).map_err(|error| { - Error::InvalidResponse(crate::ErrorDetail::invalid("Anthropic stream event", error)) - })?; - let name = data - .get("type") - .and_then(serde_json::Value::as_str) - .ok_or_else(|| { - Error::InvalidResponse( - "Anthropic stream event is invalid: stream event has no type".into(), - ) - })?; - Ok(Bytes::from(format!("event: {name}\ndata: {data}\n\n"))) +pub fn event_stream(bytes: ByteStream) -> EventStream { + sse::event_stream(bytes, no_control_frames) +} + +pub fn decode(wire: &[u8]) -> Result, Error> { + sse::decode(wire, no_control_frames) +} + +/// Anthropic names every frame after the event's `type`, and LiteLLM's proxy does the same +/// when it re-emits translated events. +pub fn encode(event: &MessagesStreamEvent) -> Result { + let data = serde_json::to_string(event).map_err(sse::invalid)?; + Ok(Bytes::from(format!( + "event: {}\ndata: {data}\n\n", + event.as_ref() + ))) } #[cfg(test)] @@ -41,6 +35,7 @@ mod tests { use litellm_llms_types::formats::messages::streaming::{ MessagesContentBlockDelta, MessagesStreamUsage, }; + use rstest::rstest; use serde_json::json; use super::*; @@ -48,51 +43,11 @@ mod tests { const TEXT_DELTA: &str = r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#; - fn in_pieces(wire: &[u8]) -> ByteStream { - let pieces: Vec = wire.chunks(3).map(Bytes::copy_from_slice).collect(); + fn in_pieces(wire: &[u8], size: usize) -> ByteStream { + let pieces: Vec = wire.chunks(size).map(Bytes::copy_from_slice).collect(); stream::iter(pieces.into_iter().map(Ok)).boxed() } - #[tokio::test] - async fn sse_frames_split_anywhere_decode_into_typed_events() { - let wire = format!("event: content_block_delta\ndata: {TEXT_DELTA}\n\n"); - let events = anthropic_sse_event_stream(in_pieces(wire.as_bytes())) - .try_collect::>() - .await - .unwrap(); - - assert_eq!( - events, - vec![MessagesStreamEvent::ContentBlockDelta { - index: 0, - delta: MessagesContentBlockDelta::TextDelta { - text: "hello".into(), - }, - }] - ); - } - - #[tokio::test] - async fn decodes_citations_delta_events() { - let wire = concat!( - "event: content_block_delta\n", - r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#, - "\n\n", - ); - let events = anthropic_sse_event_stream(in_pieces(wire.as_bytes())) - .try_collect::>() - .await - .unwrap(); - - assert!(matches!( - events.as_slice(), - [MessagesStreamEvent::ContentBlockDelta { - delta: MessagesContentBlockDelta::Citations { .. }, - .. - }] - )); - } - fn events() -> Vec { vec![ MessagesStreamEvent::Ping, @@ -110,27 +65,80 @@ mod tests { ] } - #[tokio::test] - async fn encoded_events_decode_back_to_themselves() { - let wire = events() - .iter() - .map(encode_anthropic_sse) - .collect::, _>>() - .unwrap(); + fn wire(newline: &str) -> String { + format!( + "event: content_block_delta{newline}data: {TEXT_DELTA}{newline}{newline}: keepalive{newline}{newline}event: message_stop{newline}data: {{\"type\":\"message_stop\"}}{newline}{newline}" + ) + } - let decoded = anthropic_sse_event_stream(stream::iter(wire.into_iter().map(Ok)).boxed()) + #[rstest] + #[case::lf("\n")] + #[case::crlf("\r\n")] + #[case::cr("\r")] + #[tokio::test] + async fn frames_split_anywhere_decode_into_typed_events( + #[case] newline: &str, + #[values(1, 3, 4096)] size: usize, + ) { + let wire = wire(newline); + let expected = vec![ + MessagesStreamEvent::ContentBlockDelta { + index: 0, + delta: MessagesContentBlockDelta::TextDelta { + text: "hello".into(), + }, + }, + MessagesStreamEvent::MessageStop { usage: None }, + ]; + let streamed = event_stream(in_pieces(wire.as_bytes(), size)) .try_collect::>() .await .unwrap(); + assert_eq!(streamed, expected); + assert_eq!(decode(wire.as_bytes()).unwrap(), expected); + } - assert_eq!(decoded, events()); + #[tokio::test] + async fn decodes_citations_delta_events() { + let wire = concat!( + "event: content_block_delta\n", + r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#, + "\n\n", + ); + let events = event_stream(in_pieces(wire.as_bytes(), 3)) + .try_collect::>() + .await + .unwrap(); + assert!(matches!( + events.as_slice(), + [MessagesStreamEvent::ContentBlockDelta { + delta: MessagesContentBlockDelta::Citations { .. }, + .. + }] + )); + } + + #[tokio::test] + async fn encoded_events_decode_back_to_themselves() { + let wire: Vec = events() + .iter() + .map(encode) + .collect::, _>>() + .unwrap() + .into_iter() + .flatten() + .collect(); + let streamed = event_stream(in_pieces(&wire, 3)) + .try_collect::>() + .await + .unwrap(); + assert_eq!(streamed, events()); + assert_eq!(decode(&wire).unwrap(), events()); } #[test] fn an_event_is_named_by_its_type() { - let encoded = - encode_anthropic_sse(&MessagesStreamEvent::MessageStop { usage: None }).unwrap(); - + let encoded = encode(&MessagesStreamEvent::MessageStop { usage: None }).unwrap(); assert_eq!( encoded, Bytes::from(format!( @@ -139,4 +147,68 @@ mod tests { )) ); } + + #[rstest] + #[case::json(b"data: {\n\n")] + #[case::unknown_event(b"data: {\"type\":\"future_event\"}\n\n")] + #[case::malformed_tool(b"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":3}}\n\n")] + #[case::utf8(b"data: \xff\n\n")] + #[case::done_is_not_an_anthropic_frame(b"data: [DONE]\n\n")] + #[tokio::test] + async fn malformed_streams_are_invalid_responses(#[case] wire: &[u8]) { + assert!(matches!(decode(wire), Err(Error::InvalidResponse(_)))); + let streamed = event_stream(in_pieces(wire, 1)).collect::>().await; + assert!(matches!( + streamed.last(), + Some(Err(Error::InvalidResponse(_))) | None + )); + } + + #[rstest] + fn decoding_keeps_tool_deltas_and_block_extensions() { + let wire = concat!( + "data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"t\",\"name\":\"lookup\",\"input\":{},\"future\":null}}\n\n", + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"x\\\":\"}}\n\n", + ); + let decoded = decode(wire.as_bytes()).unwrap(); + let [ + MessagesStreamEvent::ContentBlockStart { content_block, .. }, + MessagesStreamEvent::ContentBlockDelta { + delta: MessagesContentBlockDelta::InputJsonDelta { partial_json }, + .. + }, + ] = decoded.as_slice() + else { + panic!("expected tool events"); + }; + assert_eq!( + content_block.extra.get("future"), + Some(&serde_json::Value::Null) + ); + assert_eq!(partial_json, "{\"x\":"); + } + + #[rstest] + #[tokio::test] + async fn a_body_error_ends_the_stream_after_delivered_events() { + let mut events = event_stream( + stream::iter([ + Ok(Bytes::from_static(b"data: {\"type\":\"ping\"}\n\n")), + Err(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + "reset", + )), + ]) + .boxed(), + ); + assert_eq!( + events.next().await.unwrap().unwrap(), + MessagesStreamEvent::Ping + ); + assert!(matches!( + events.next().await, + Some(Err(Error::InvalidResponse(_))) + )); + assert!(events.next().await.is_none()); + } } diff --git a/litellm-rust/crates/llms/src/base_llm/mod.rs b/litellm-rust/crates/llms/src/base_llm/mod.rs index 9a6a0659d7d..4c04c23a0bb 100644 --- a/litellm-rust/crates/llms/src/base_llm/mod.rs +++ b/litellm-rust/crates/llms/src/base_llm/mod.rs @@ -5,3 +5,4 @@ pub mod chat; pub mod messages; pub mod ocr; pub mod responses; +pub mod sse; diff --git a/litellm-rust/crates/llms/src/base_llm/responses/mod.rs b/litellm-rust/crates/llms/src/base_llm/responses/mod.rs index f239b6921fa..fa7df180f50 100644 --- a/litellm-rust/crates/llms/src/base_llm/responses/mod.rs +++ b/litellm-rust/crates/llms/src/base_llm/responses/mod.rs @@ -1 +1,2 @@ +pub mod streaming; pub mod transformation; diff --git a/litellm-rust/crates/llms/src/base_llm/responses/streaming.rs b/litellm-rust/crates/llms/src/base_llm/responses/streaming.rs new file mode 100644 index 00000000000..8310f0b6d9d --- /dev/null +++ b/litellm-rust/crates/llms/src/base_llm/responses/streaming.rs @@ -0,0 +1,135 @@ +use bytes::Bytes; +use litellm_llms_types::formats::responses::streaming::ResponsesStreamEvent; + +pub use crate::base_llm::base_model_iterator::ByteStream; +use crate::{ + Error, + base_llm::sse::{self, is_done}, +}; + +pub type EventStream = sse::EventStream; + +pub fn event_stream(bytes: ByteStream) -> EventStream { + sse::event_stream(bytes, is_done) +} + +pub fn decode(wire: &[u8]) -> Result, Error> { + sse::decode(wire, is_done) +} + +/// A bare `data:` frame. Python's proxy serializes every Responses event this way and never +/// writes an `event:` line, so clients see one shape whichever runtime served them. +pub fn encode(event: &ResponsesStreamEvent) -> Result { + sse::encode_data(event) +} + +#[cfg(test)] +mod tests { + use futures_util::{StreamExt, TryStreamExt, stream}; + use litellm_llms_types::formats::responses::streaming::ResponsesEventType; + use rstest::rstest; + use serde_json::{Value, json}; + + use super::*; + + fn in_pieces(wire: &[u8], size: usize) -> ByteStream { + let pieces: Vec = wire.chunks(size).map(Bytes::copy_from_slice).collect(); + stream::iter(pieces.into_iter().map(Ok)).boxed() + } + + #[rstest] + #[case::created(json!({"type":"response.created","response":{"id":"r","output":[]}}))] + #[case::text(json!({"type":"response.output_text.delta","delta":"hello\nworld","sequence_number":3}))] + #[case::tool(json!({"type":"response.function_call_arguments.delta","delta":"{\"x\":","item_id":"fc_1"}))] + #[case::unknown(json!({"type":"future.event","payload":[1,null],"future":{"nested":true}}))] + #[case::completed(json!({"type":"response.completed","response":{"id":"r","usage":{"input_tokens":3}}}))] + #[case::failed(json!({"type":"response.failed","response":{"error":{"code":"failure"}}}))] + #[case::incomplete(json!({"type":"response.incomplete","response":{"incomplete_details":{"reason":"limit"}}}))] + #[case::error(json!({"type":"error","error":{"message":"boom"}}))] + fn encoded_events_are_bare_data_frames_that_decode_back(#[case] wire: Value) { + let event: ResponsesStreamEvent = serde_json::from_value(wire.clone()).unwrap(); + let encoded = encode(&event).unwrap(); + assert!(encoded.starts_with(b"data: ")); + assert!(encoded.ends_with(b"\n\n")); + assert_eq!( + serde_json::from_slice::(&encoded[b"data: ".len()..encoded.len() - 2]).unwrap(), + wire + ); + assert_eq!(decode(&encoded).unwrap(), vec![event]); + } + + #[rstest] + #[case::lf("\n")] + #[case::crlf("\r\n")] + #[case::cr("\r")] + #[tokio::test] + async fn streams_split_anywhere_skip_comments_and_done( + #[case] newline: &str, + #[values(1, 5, 4096)] size: usize, + ) { + let wire = format!( + "\u{feff}event: response.created{newline}data: {{\"type\":\"response.created\"}}{newline}{newline}: keepalive{newline}{newline}event: future.event{newline}data: {{\"type\":\"future.event\",\"payload\":null}}{newline}{newline}data: [DONE]{newline}{newline}" + ); + let streamed = event_stream(in_pieces(wire.as_bytes(), size)) + .try_collect::>() + .await + .unwrap(); + assert_eq!(streamed, decode(wire.as_bytes()).unwrap()); + assert_eq!(streamed.len(), 2); + assert_eq!(streamed[0].event_type, ResponsesEventType::ResponseCreated); + assert_eq!( + streamed[1].event_type, + ResponsesEventType::Other("future.event".into()) + ); + assert_eq!(streamed[1].data.get("payload"), Some(&Value::Null)); + } + + #[rstest] + #[case::spaced(b"data: [DONE]\n\n")] + #[case::unspaced(b"data:[DONE]\n\n")] + #[case::trailing_space(b"data: [DONE] \n\n")] + fn every_done_spelling_closes_the_stream_without_an_event(#[case] wire: &[u8]) { + assert_eq!(decode(wire).unwrap(), vec![]); + } + + #[rstest] + #[case::invalid_json(b"data: {\n\n")] + #[case::missing_type(b"data: {}\n\n")] + #[case::numeric_type(b"data: {\"type\":3}\n\n")] + #[case::invalid_utf8(b"data: \xff\n\n")] + #[tokio::test] + async fn malformed_streams_are_invalid_responses(#[case] wire: &[u8]) { + assert!(matches!(decode(wire), Err(Error::InvalidResponse(_)))); + let streamed = event_stream(in_pieces(wire, 1)).collect::>().await; + assert!(matches!( + streamed.last(), + Some(Err(Error::InvalidResponse(_))) | None + )); + } + + #[rstest] + #[tokio::test] + async fn a_body_error_ends_the_stream_after_delivered_events() { + let mut events = event_stream( + stream::iter([ + Ok(Bytes::from_static( + b"data: {\"type\":\"response.created\"}\n\n", + )), + Err(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + "reset", + )), + ]) + .boxed(), + ); + assert_eq!( + events.next().await.unwrap().unwrap().event_type, + ResponsesEventType::ResponseCreated + ); + assert!(matches!( + events.next().await, + Some(Err(Error::InvalidResponse(_))) + )); + assert!(events.next().await.is_none()); + } +} diff --git a/litellm-rust/crates/llms/src/base_llm/responses/transformation.rs b/litellm-rust/crates/llms/src/base_llm/responses/transformation.rs index 30899692fb6..93eaad0ebc0 100644 --- a/litellm-rust/crates/llms/src/base_llm/responses/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/responses/transformation.rs @@ -1,5 +1,5 @@ use litellm_llms_types::formats::responses::{ - ResponsesApiResponse, streaming_websocket::ResponsesWsEvent, + ResponsesApiResponse, streaming::ResponsesStreamEvent, }; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; @@ -38,11 +38,11 @@ pub trait BaseResponsesApiConfig: Sync { #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct ResponsesWsTransformResult { - pub events: Vec, + pub events: Vec, } impl ResponsesWsTransformResult { - pub fn passthrough(event: ResponsesWsEvent) -> Self { + pub fn passthrough(event: ResponsesStreamEvent) -> Self { Self { events: vec![event], } @@ -58,18 +58,18 @@ pub trait ResponsesWebSocketProviderConfig: Sync { fn transform_ws_request( &self, - event: &ResponsesWsEvent, + event: &ResponsesStreamEvent, model: &str, ) -> Result; fn transform_ws_response( &self, - event: &ResponsesWsEvent, + event: &ResponsesStreamEvent, model: &str, ) -> Result; } -pub fn enforce_model(event: &ResponsesWsEvent, model: &str) -> ResponsesWsEvent { +pub fn enforce_model(event: &ResponsesStreamEvent, model: &str) -> ResponsesStreamEvent { if !event.is_response_create() { return event.clone(); } @@ -103,7 +103,7 @@ pub fn enforce_model(event: &ResponsesWsEvent, model: &str) -> ResponsesWsEvent mod tests { use super::*; - fn event(value: serde_json::Value) -> ResponsesWsEvent { + fn event(value: serde_json::Value) -> ResponsesStreamEvent { serde_json::from_value(value).expect("valid event") } diff --git a/litellm-rust/crates/llms/src/base_llm/sse.rs b/litellm-rust/crates/llms/src/base_llm/sse.rs new file mode 100644 index 00000000000..5ff2e3adddd --- /dev/null +++ b/litellm-rust/crates/llms/src/base_llm/sse.rs @@ -0,0 +1,102 @@ +//! SSE decoding and encoding shared by every format's `streaming` module. +//! +//! DIVERGENCE: decoding is strict. An event whose `data` is not a valid event of the format is +//! `InvalidResponse` and ends the stream, and so is `[DONE]` on a Messages stream. Python +//! never fails a stream on a malformed event: +//! - Chat Completions (`BaseModelResponseIterator`) and Responses +//! (`BaseResponsesAPIStreamingIterator`) skip the event and keep going +//! - Messages (`BaseAnthropicMessagesStreamingIterator`) forwards the provider's bytes to the +//! client without parsing them +//! - Anthropic to Chat (`anthropic/chat/handler.py` `ModelResponseIterator`) buffers the invalid +//! lines on the guess that one payload was split across them +//! +//! Skipping loses content or a tool-call fragment with nothing to show for it, so Rust reports it. + +use bytes::Bytes; +use futures_util::{StreamExt, TryStreamExt, future::ready, stream::BoxStream}; +use litellm_framer::{ + frames, + sse::{SseCodec, SseEvent, decode_all}, +}; +use serde::{Serialize, de::DeserializeOwned}; + +use crate::{Error, ErrorDetail, base_llm::base_model_iterator::ByteStream}; + +pub type EventStream = BoxStream<'static, Result>; + +/// Which frames are control frames rather than events of the format. +pub type IsControl = fn(&SseEvent) -> bool; + +/// OpenAI's end-of-stream marker on Chat Completions and Responses streams. It closes the +/// stream and is not an event; Anthropic Messages streams never carry it. +pub const DONE_FRAME: &[u8] = b"data: [DONE]\n\n"; + +pub fn is_done(event: &SseEvent) -> bool { + event.data.trim() == "[DONE]" +} + +pub fn event_stream( + bytes: ByteStream, + is_control: IsControl, +) -> EventStream { + frames(bytes, SseCodec::default()) + .map_err(framing) + .try_filter(move |event| ready(!is_control(event))) + .and_then(|event| ready(parse(&event))) + .boxed() +} + +/// Decodes a whole stream at once. An event the stream ends inside is dropped. +pub fn decode(wire: &[u8], is_control: IsControl) -> Result, Error> { + decode_all(wire) + .map_err(framing)? + .iter() + .filter(|event| !is_control(event)) + .map(parse) + .collect() +} + +/// A bare `data:` frame, which is what LiteLLM's proxy sends to clients on OpenAI formats. +pub fn encode_data(event: &E) -> Result { + let data = serde_json::to_string(event).map_err(invalid)?; + Ok(Bytes::from(format!("data: {data}\n\n"))) +} + +fn parse(event: &SseEvent) -> Result { + serde_json::from_str(&event.data).map_err(invalid) +} + +pub fn invalid(error: serde_json::Error) -> Error { + Error::InvalidResponse(ErrorDetail::invalid("stream event", error)) +} + +fn framing(error: litellm_framer::SseError) -> Error { + Error::InvalidResponse(ErrorDetail::failed("stream framing", error)) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::exact("[DONE]", true)] + #[case::padded(" [DONE] ", true)] + #[case::json("{\"type\":\"x\"}", false)] + #[case::mentions("[DONE] was sent", false)] + #[case::empty("", false)] + fn done_is_an_exact_marker(#[case] data: &str, #[case] done: bool) { + let event = SseEvent { + data: data.into(), + ..SseEvent::default() + }; + assert_eq!(is_done(&event), done); + } + + #[rstest] + fn the_done_frame_we_send_decodes_as_a_control_frame() { + let events: Vec = decode(DONE_FRAME, is_done).unwrap(); + assert!(events.is_empty()); + } +} diff --git a/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs b/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs index 373ad4ca4f5..878c424ae2d 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs @@ -102,7 +102,7 @@ mod tests { use litellm_llms_types::formats::messages::streaming::MessagesContentBlockDelta; use super::*; - use crate::base_llm::messages::streaming::anthropic_sse_event_stream; + use crate::base_llm::messages::streaming; fn in_pieces(wire: &[u8]) -> ByteStream { let pieces: Vec = wire.chunks(3).map(Bytes::copy_from_slice).collect(); @@ -147,7 +147,7 @@ mod tests { .try_collect::>() .await .unwrap(); - let from_sse = anthropic_sse_event_stream(in_pieces(sse.as_bytes())) + let from_sse = streaming::event_stream(in_pieces(sse.as_bytes())) .try_collect::>() .await .unwrap(); diff --git a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs index 4f8d4f01697..b467ec723a5 100644 --- a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs @@ -332,7 +332,7 @@ mod tests { use litellm_auth_aws::constants::DEFAULT_BEDROCK_REGION; use super::*; - use crate::base_llm::messages::streaming::encode_anthropic_sse; + use crate::base_llm::messages::streaming; fn event(value: Value) -> MessagesStreamEvent { serde_json::from_value(value).unwrap() @@ -493,7 +493,7 @@ mod tests { .boxed(); let sse = bedrock_anthropic_messages_event_stream(bytes) - .map_ok(|event| encode_anthropic_sse(&event).unwrap()) + .map_ok(|event| streaming::encode(&event).unwrap()) .try_collect::>() .await .unwrap() @@ -508,7 +508,7 @@ mod tests { )), ] .iter() - .flat_map(|event| encode_anthropic_sse(event).unwrap()) + .flat_map(|event| streaming::encode(event).unwrap()) .collect(); assert_eq!(sse, expected); } diff --git a/litellm-rust/crates/llms/src/openai/responses/transformation.rs b/litellm-rust/crates/llms/src/openai/responses/transformation.rs index 959be2a9a4a..3f66c0224d6 100644 --- a/litellm-rust/crates/llms/src/openai/responses/transformation.rs +++ b/litellm-rust/crates/llms/src/openai/responses/transformation.rs @@ -1,5 +1,5 @@ use litellm_llms_types::formats::responses::{ - ResponsesApiResponse, streaming_websocket::ResponsesWsEvent, + ResponsesApiResponse, streaming::ResponsesStreamEvent, }; use serde_json::{Map, Value}; @@ -34,7 +34,7 @@ impl ResponsesWebSocketProviderConfig for OpenAiResponsesApiConfig { fn transform_ws_request( &self, - event: &ResponsesWsEvent, + event: &ResponsesStreamEvent, model: &str, ) -> Result { Ok(ResponsesWsTransformResult::passthrough(enforce_model( @@ -44,7 +44,7 @@ impl ResponsesWebSocketProviderConfig for OpenAiResponsesApiConfig { fn transform_ws_response( &self, - event: &ResponsesWsEvent, + event: &ResponsesStreamEvent, _model: &str, ) -> Result { Ok(ResponsesWsTransformResult::passthrough(event.clone())) @@ -230,7 +230,7 @@ mod tests { #[test] fn openai_config_is_native_and_enforces_model() { - let event: ResponsesWsEvent = + let event: ResponsesStreamEvent = serde_json::from_value(serde_json::json!({"type":"response.create"})) .expect("valid event"); let result = OPENAI_RESPONSES_WS_CONFIG From 0ebec123fd465676f2fb62a1f8355631dec8b438 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 13 Oct 2026 22:35:23 +0000 Subject: [PATCH 019/325] chore(typing): clear 670 basedpyright Any errors by validating provider payloads at the boundary (#45388) * chore(typing): validate provider payloads at the boundary to clear 464 basedpyright Any errors Replaces untyped reads of decoded JSON in 48 provider, guardrail and proxy modules with one pydantic validation at the point the payload enters the module, so every value read afterwards carries a real type. reportAny 6111 -> 5663 and reportExplicitAny 1416 -> 1400. No basedpyright, ruff-strict or type-discipline rule is higher repo-wide or in any single file, and no cast, ignore or suppression marker was added. * chore(typing): replace 47 Any annotations that every caller already satisfies with object Each annotation here is a parameter or return type whose value is only stored, compared, logged or forwarded, so object describes it exactly and no function body changes. A site was kept only if no basedpyright, ruff-strict or type-discipline rule rose in any file that references it. reportExplicitAny 1400 -> 1353 and reportAny 5663 -> 5655. * chore(typing): type upstream responses and loose annotations across providers, proxy and integrations Second pass of the Any reduction. Decoded upstream JSON is validated once at the provider boundary and read through typed views afterwards, annotations that every caller already satisfies move from Any to a concrete type, and the nine search transformations validate each result item on its own so error text no longer carries the item index. No cast, no type: ignore, no suppression marker and no budget file edit * ci(unit): run the brave, google_pse, nlp_cloud and opensandbox tests in the providers shard --- .github/workflows/test-unit.yml | 4 + .../providers/bedrock_agentcore/handler.py | 2 +- litellm/batches/batch_utils.py | 2 +- .../harness/handlers/deepagents_handler.py | 2 +- litellm/integrations/argilla.py | 2 +- .../arize/arize_phoenix_client.py | 4 +- .../clickhouse/clickhouse_spend_logger.py | 2 +- .../integrations/cloudzero/cz_stream_api.py | 2 +- litellm/integrations/custom_guardrail.py | 2 +- .../integrations/dotprompt/prompt_manager.py | 4 +- litellm/integrations/focus/focus_logger.py | 2 +- .../integrations/vantage/vantage_logger.py | 4 +- litellm/llms/a2a/chat/transformation.py | 6 +- litellm/llms/a2a/common_utils.py | 2 +- .../aiml/image_generation/transformation.py | 33 +- .../anthropic/count_tokens/token_counter.py | 8 +- .../llms/anthropic/files/transformation.py | 9 +- .../pass_through/adapters/transformation.py | 2 +- .../anthropic/count_tokens/handler.py | 9 +- .../anthropic/count_tokens/token_counter.py | 20 +- .../image_edit/flux2_transformation.py | 2 +- .../base_llm/managed_resources/isolation.py | 5 +- .../amazon_qwen2_transformation.py | 21 +- .../amazon_qwen3_transformation.py | 20 +- .../base_invoke_transformation.py | 6 +- litellm/llms/brave/search/transformation.py | 28 +- litellm/llms/cohere/embed/transformation.py | 11 +- litellm/llms/cometapi/embed/transformation.py | 10 +- .../llms/dashscope/embed/transformation.py | 2 +- .../llms/databricks/chat/transformation.py | 4 +- .../llms/dataforseo/search/transformation.py | 24 +- .../llms/deepagents/harness/transformation.py | 17 +- .../llms/duckduckgo/search/transformation.py | 23 +- litellm/llms/exa_ai/search/transformation.py | 31 +- .../llms/fal_ai/image_edit/transformation.py | 9 +- .../image_generation/bria_transformation.py | 9 +- .../imagen4_transformation.py | 14 +- .../fal_ai/image_generation/transformation.py | 10 +- litellm/llms/fastcrw/search/transformation.py | 24 +- .../llms/firecrawl/search/transformation.py | 63 ++-- .../llms/fireworks_ai/chat/transformation.py | 16 +- .../gemini/image_generation/transformation.py | 29 +- .../llms/google_pse/search/transformation.py | 24 +- .../llms/infinity/embedding/transformation.py | 10 +- litellm/llms/langgraph/chat/transformation.py | 4 +- litellm/llms/lemonade/chat/transformation.py | 21 +- litellm/llms/linkup/search/transformation.py | 40 ++- .../mistral/audio_speech/transformation.py | 9 +- litellm/llms/nlp_cloud/chat/transformation.py | 3 +- litellm/llms/oci/chat/transformation.py | 2 +- litellm/llms/ollama/common_utils.py | 6 +- .../llms/ollama/completion/transformation.py | 5 +- .../openai/responses/count_tokens/handler.py | 4 +- .../responses/count_tokens/token_counter.py | 14 +- .../responses/count_tokens/transformation.py | 2 +- .../llms/openrouter/chat/transformation.py | 16 +- .../opensandbox/sandbox/transformation.py | 9 +- .../llms/ovhcloud/embedding/transformation.py | 10 +- .../llms/perplexity/chat/transformation.py | 4 +- .../perplexity/embedding/transformation.py | 17 +- .../llms/perplexity/search/transformation.py | 24 +- litellm/llms/replicate/chat/handler.py | 8 +- litellm/llms/replicate/chat/transformation.py | 7 +- .../sambanova/embedding/transformation.py | 10 +- litellm/llms/sap/chat/transformation.py | 3 +- .../llms/searchapi/search/transformation.py | 24 +- litellm/llms/searxng/search/transformation.py | 24 +- litellm/llms/serper/search/transformation.py | 24 +- litellm/llms/tavily/search/transformation.py | 24 +- .../llms/vertex_ai/files/transformation.py | 2 +- litellm/llms/vertex_ai/fine_tuning/handler.py | 5 +- .../vertex_imagen_transformation.py | 17 +- .../vertex_gemini_transformation.py | 23 +- .../vertex_imagen_transformation.py | 21 +- .../text_to_speech/text_to_speech_handler.py | 13 +- .../text_to_speech/transformation.py | 31 +- .../vertex_embeddings/embedding_handler.py | 12 +- .../llms/voyage/embedding/transformation.py | 10 +- .../embedding/transformation_contextual.py | 9 +- .../embedding/transformation_multimodal.py | 9 +- .../llms/watsonx/completion/transformation.py | 4 +- litellm/llms/xai/common_utils.py | 9 +- .../guardrail_translation/handler.py | 6 +- .../mcp_server/mcp_server_manager.py | 2 +- .../mcp_server/openapi_to_mcp_generator.py | 2 +- .../_experimental/mcp_server/tool_registry.py | 6 +- .../proxy/agent_endpoints/a2a_endpoints.py | 2 +- .../proxy/anthropic_endpoints/endpoints.py | 2 +- litellm/proxy/auth/ip_address_utils.py | 2 +- litellm/proxy/auth/model_checks.py | 6 +- litellm/proxy/auth/trusted_proxy_utils.py | 6 +- litellm/proxy/caching_routes.py | 4 +- litellm/proxy/client/cli/commands/auth.py | 26 +- litellm/proxy/client/cli/commands/http.py | 45 ++- litellm/proxy/client/cli/commands/keys.py | 27 +- litellm/proxy/client/cli/commands/models.py | 2 +- litellm/proxy/client/cli/main.py | 13 +- litellm/proxy/common_request_processing.py | 14 +- .../proxy/common_utils/http_parsing_utils.py | 2 +- .../common_utils/openapi_schema_compat.py | 4 +- litellm/proxy/guardrails/_content_utils.py | 4 +- .../guardrail_hooks/dynamoai/dynamoai.py | 6 +- .../generic_guardrail_api.py | 4 +- .../ibm_guardrails/ibm_detector.py | 4 +- .../llm_shield_proxy/llm_shield_proxy.py | 11 +- .../guardrails/guardrail_hooks/presidio.py | 1 + .../semantic_guard/route_loader.py | 6 +- .../shared_health_check_manager.py | 4 +- litellm/proxy/litellm_pre_call_utils.py | 2 +- litellm/proxy/mcp_tools.py | 7 +- .../file_content_streaming_handler.py | 2 +- .../anthropic_passthrough_logging_handler.py | 7 +- litellm/proxy/plugin_routes.py | 4 +- .../policy_engine/attachment_registry.py | 2 +- litellm/proxy/policy_engine/init_policies.py | 8 +- .../proxy/policy_engine/policy_registry.py | 15 +- litellm/proxy/proxy_server.py | 8 +- litellm/proxy/read_model_list.py | 4 +- .../response_polling/background_streaming.py | 17 +- .../proxy/response_polling/polling_handler.py | 25 +- litellm/proxy/types_utils/utils.py | 12 +- litellm/proxy/vector_store_endpoints/utils.py | 6 +- litellm/proxy/video_endpoints/utils.py | 2 +- litellm/rag/rag_query.py | 2 +- litellm/router.py | 4 +- .../secret_managers/google_secret_manager.py | 14 +- litellm/utils.py | 4 +- .../test_bedrock_agentcore_a2a.py | 50 +++ .../integrations/focus/test_focus_logger.py | 92 ++++++ tests/unit/integrations/test_argilla.py | 110 +++++++ tests/unit/integrations/vantage/__init__.py | 0 .../vantage/test_vantage_logger.py | 92 ++++++ .../a2a/chat/test_a2a_chat_transformation.py | 105 +++++++ ...st_aiml_image_generation_transformation.py | 133 +++++++- .../test_anthropic_files_transformation.py | 151 +++++++++ .../unit/llms/azure_ai/anthropic/__init__.py | 0 .../anthropic/count_tokens/__init__.py | 0 .../count_tokens/test_token_counter.py | 188 ++++++++++++ .../test_amazon_qwen2_transformation.py | 100 +++++- .../test_amazon_qwen3_transformation.py | 85 ++++- .../test_base_invoke_transformation.py | 21 ++ tests/unit/llms/brave/__init__.py | 0 tests/unit/llms/brave/search/__init__.py | 0 .../test_brave_search_transformation.py | 123 ++++++++ .../llms/cohere/embed/test_transformation.py | 139 +++++++++ tests/unit/llms/cometapi/embed/__init__.py | 0 .../test_cometapi_embed_transformation.py | 89 ++++++ .../test_dataforseo_search_transformation.py | 118 +++++++ .../test_duckduckgo_search_transformation.py | 133 ++++++++ .../llms/exa_ai/search/test_transformation.py | 109 +++++++ .../test_fal_ai_image_edit_transformation.py | 68 ++++ .../test_bria_transformation.py | 93 ++++++ ..._fal_ai_image_generation_transformation.py | 97 ++++++ .../test_imagen4_transformation.py | 112 +++++++ .../fastcrw/search/test_transformation.py | 102 ++++++ .../test_firecrawl_search_transformation.py | 160 ++++++++++ .../test_fireworks_ai_chat_transformation.py | 118 ++++++- ..._gemini_image_generation_transformation.py | 116 +++++++ tests/unit/llms/google_pse/__init__.py | 0 tests/unit/llms/google_pse/search/__init__.py | 0 .../test_google_pse_search_transformation.py | 110 +++++++ .../unit/llms/infinity/embedding/__init__.py | 0 .../test_infinity_embedding_transformation.py | 89 ++++++ .../chat/test_langgraph_transformation.py | 36 +++ tests/unit/llms/lemonade/chat/__init__.py | 0 .../chat/test_lemonade_chat_transformation.py | 172 +++++++++++ .../test_linkup_search_transformation.py | 100 ++++++ ...est_mistral_audio_speech_transformation.py | 112 +++++++ tests/unit/llms/nlp_cloud/__init__.py | 0 tests/unit/llms/nlp_cloud/chat/__init__.py | 0 .../test_nlp_cloud_chat_transformation.py | 47 +++ .../test_ollama_completion_transformation.py | 53 ++++ .../count_tokens/test_token_counter.py | 187 +++++++++++ .../test_openrouter_chat_transformation.py | 85 ++++- tests/unit/llms/opensandbox/__init__.py | 0 .../unit/llms/opensandbox/sandbox/__init__.py | 0 .../sandbox/test_transformation.py | 153 +++++++++ ...test_ovhcloud_embeddings_transformation.py | 89 ++++++ .../test_perplexity_chat_transformation.py | 167 ++++++++++ ...est_perplexity_embedding_transformation.py | 100 ++++++ .../test_perplexity_search_transformation.py | 116 +++++++ .../unit/llms/replicate/chat/test_handler.py | 206 +++++++++++++ .../replicate/chat/test_transformation.py | 103 ++++++- .../unit/llms/sambanova/embedding/__init__.py | 0 ...test_sambanova_embedding_transformation.py | 89 ++++++ .../test_searchapi_search_transformation.py | 94 +++++- .../test_searxng_search_transformation.py | 105 +++++++ .../test_serper_search_transformation.py | 83 +++++ .../test_tavily_search_transformation.py | 91 ++++++ .../vertex_ai/fine_tuning/test_handler.py | 106 +++++++ ...est_vertex_ai_image_edit_transformation.py | 186 +++++++++++ ...rtex_ai_image_generation_transformation.py | 290 +++++++++++++++++- .../test_text_to_speech_handler.py | 155 ++++++++++ .../text_to_speech/test_transformation.py | 168 ++++++++++ .../vertex_ai/vertex_embeddings/__init__.py | 0 .../test_embedding_handler.py | 289 +++++++++++++++++ .../test_transformation_contextual.py | 89 ++++++ .../test_transformation_multimodal.py | 91 ++++++ .../test_voyage_embedding_transformation.py | 91 ++++++ .../unit/llms/watsonx/completion/__init__.py | 0 .../test_watsonx_completion_transformation.py | 60 ++++ tests/unit/llms/xai/test_xai_common_utils.py | 124 ++++++++ .../agent_endpoints/test_a2a_endpoints.py | 59 ++++ .../proxy/client/cli/test_auth_commands.py | 91 ++++++ .../proxy/client/cli/test_global_options.py | 103 +++++++ .../proxy/client/cli/test_keys_commands.py | 91 +++++- tests/unit/proxy/client/test_http_commands.py | 119 +++++++ .../guardrail_hooks/dynamoai/test_dynamoai.py | 85 +++++ .../ibm_guardrails/__init__.py | 0 .../ibm_guardrails/test_ibm_detector.py | 124 ++++++++ ...t_anthropic_passthrough_logging_handler.py | 46 +++ .../policy_engine/test_policy_registry.py | 131 ++++++++ .../test_response_polling_handler.py | 45 +++ tests/unit/proxy/test_caching_routes.py | 63 ++++ tests/unit/proxy/test_plugin_routes.py | 90 ++++++ tests/unit/proxy/types_utils/test_utils.py | 250 +++++++++++++++ .../test_vector_store_utils.py | 94 ++++++ .../test_google_secret_manager.py | 243 +++++++++++++++ 218 files changed, 9258 insertions(+), 484 deletions(-) create mode 100644 tests/unit/integrations/focus/test_focus_logger.py create mode 100644 tests/unit/integrations/test_argilla.py create mode 100644 tests/unit/integrations/vantage/__init__.py create mode 100644 tests/unit/integrations/vantage/test_vantage_logger.py create mode 100644 tests/unit/llms/azure_ai/anthropic/__init__.py create mode 100644 tests/unit/llms/azure_ai/anthropic/count_tokens/__init__.py create mode 100644 tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py create mode 100644 tests/unit/llms/brave/__init__.py create mode 100644 tests/unit/llms/brave/search/__init__.py create mode 100644 tests/unit/llms/brave/search/test_brave_search_transformation.py create mode 100644 tests/unit/llms/cometapi/embed/__init__.py create mode 100644 tests/unit/llms/cometapi/embed/test_cometapi_embed_transformation.py create mode 100644 tests/unit/llms/fal_ai/image_generation/test_bria_transformation.py create mode 100644 tests/unit/llms/fal_ai/image_generation/test_fal_ai_image_generation_transformation.py create mode 100644 tests/unit/llms/fal_ai/image_generation/test_imagen4_transformation.py create mode 100644 tests/unit/llms/firecrawl/search/test_firecrawl_search_transformation.py create mode 100644 tests/unit/llms/google_pse/__init__.py create mode 100644 tests/unit/llms/google_pse/search/__init__.py create mode 100644 tests/unit/llms/google_pse/search/test_google_pse_search_transformation.py create mode 100644 tests/unit/llms/infinity/embedding/__init__.py create mode 100644 tests/unit/llms/infinity/embedding/test_infinity_embedding_transformation.py create mode 100644 tests/unit/llms/lemonade/chat/__init__.py create mode 100644 tests/unit/llms/lemonade/chat/test_lemonade_chat_transformation.py create mode 100644 tests/unit/llms/nlp_cloud/__init__.py create mode 100644 tests/unit/llms/nlp_cloud/chat/__init__.py create mode 100644 tests/unit/llms/nlp_cloud/chat/test_nlp_cloud_chat_transformation.py create mode 100644 tests/unit/llms/openai/responses/count_tokens/test_token_counter.py create mode 100644 tests/unit/llms/opensandbox/__init__.py create mode 100644 tests/unit/llms/opensandbox/sandbox/__init__.py create mode 100644 tests/unit/llms/opensandbox/sandbox/test_transformation.py create mode 100644 tests/unit/llms/replicate/chat/test_handler.py create mode 100644 tests/unit/llms/sambanova/embedding/__init__.py create mode 100644 tests/unit/llms/sambanova/embedding/test_sambanova_embedding_transformation.py create mode 100644 tests/unit/llms/vertex_ai/text_to_speech/test_text_to_speech_handler.py create mode 100644 tests/unit/llms/vertex_ai/vertex_embeddings/__init__.py create mode 100644 tests/unit/llms/vertex_ai/vertex_embeddings/test_embedding_handler.py create mode 100644 tests/unit/llms/voyage/embedding/test_transformation_multimodal.py create mode 100644 tests/unit/llms/voyage/embedding/test_voyage_embedding_transformation.py create mode 100644 tests/unit/llms/watsonx/completion/__init__.py create mode 100644 tests/unit/llms/watsonx/completion/test_watsonx_completion_transformation.py create mode 100644 tests/unit/llms/xai/test_xai_common_utils.py create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py create mode 100644 tests/unit/proxy/types_utils/test_utils.py create mode 100644 tests/unit/proxy/vector_store_endpoints/test_vector_store_utils.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index e9236ee7787..0b448569095 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -233,6 +233,7 @@ jobs: tests/unit/llms/base_llm tests/unit/llms/baseten tests/unit/llms/black_forest_labs + tests/unit/llms/brave tests/unit/llms/bytez tests/unit/llms/cerebras tests/unit/llms/chat @@ -268,6 +269,7 @@ jobs: tests/unit/llms/gdc tests/unit/llms/gemini tests/unit/llms/gigachat + tests/unit/llms/google_pse tests/unit/llms/github_copilot tests/unit/llms/gradient_ai tests/unit/llms/groq @@ -300,6 +302,7 @@ jobs: tests/unit/llms/nebius tests/unit/llms/neosantara tests/unit/llms/nimble + tests/unit/llms/nlp_cloud tests/unit/llms/novita tests/unit/llms/nscale tests/unit/llms/nvidia_nim @@ -310,6 +313,7 @@ jobs: tests/unit/llms/oobabooga tests/unit/llms/openai_like tests/unit/llms/opencode + tests/unit/llms/opensandbox tests/unit/llms/openrouter tests/unit/llms/ovhcloud tests/unit/llms/parallel_ai diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index aecc887fbfe..e36d455a20f 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -117,7 +117,7 @@ class BedrockAgentCoreA2AHandler: response.raise_for_status() # Check content type — AgentCore may return JSON instead of SSE - content_type: Final = response.headers.get("content-type", "").lower() + content_type: Final = ", ".join(response.headers.get_list("content-type")).lower() if "application/json" in content_type: # Single JSON response fallback (not SSE) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 6573f13be75..416391d194b 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -768,7 +768,7 @@ def _get_response_from_batch_job_output_file( def _batch_response_was_successful( - batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai" + batch_job_output_file: Mapping[str, object], custom_llm_provider: str = "openai" ) -> bool: """ Check if the batch job response was successful diff --git a/litellm/harness/handlers/deepagents_handler.py b/litellm/harness/handlers/deepagents_handler.py index e6ca7c161c7..5215562a1c5 100644 --- a/litellm/harness/handlers/deepagents_handler.py +++ b/litellm/harness/handlers/deepagents_handler.py @@ -158,7 +158,7 @@ class DeepAgentsHandler(BaseHarnessHandler): agent, deps = self._require_agent() snapshot = await agent.aget_state(self._run_config(ctx, None)) messages = (snapshot.values or MappingProxyType({})).get("messages") or () - converted: list[dict[str, Any]] = deps.convert_to_openai_messages( # mutable-ok: LangChain returns a list + converted: list[dict[str, object]] = deps.convert_to_openai_messages( # mutable-ok: LangChain returns a list messages ) return converted diff --git a/litellm/integrations/argilla.py b/litellm/integrations/argilla.py index 664ef8efda1..f873f274ab6 100644 --- a/litellm/integrations/argilla.py +++ b/litellm/integrations/argilla.py @@ -50,7 +50,7 @@ class ArgillaLogger(CustomBatchLogger): if litellm.argilla_transformation_object is None: raise Exception("'litellm.argilla_transformation_object' is required, to log your payload to Argilla.") self.validate_argilla_transformation_object(litellm.argilla_transformation_object) - self.argilla_transformation_object = litellm.argilla_transformation_object + self.argilla_transformation_object: Mapping[str, str] = litellm.argilla_transformation_object self.default_credentials = self.get_credentials_from_env( argilla_api_key=argilla_api_key, argilla_dataset_name=argilla_dataset_name, diff --git a/litellm/integrations/arize/arize_phoenix_client.py b/litellm/integrations/arize/arize_phoenix_client.py index 18d35fee34b..b63cacede03 100644 --- a/litellm/integrations/arize/arize_phoenix_client.py +++ b/litellm/integrations/arize/arize_phoenix_client.py @@ -3,7 +3,7 @@ Arize Phoenix API client for fetching prompt versions from Arize Phoenix. """ import urllib.parse -from typing import Any, Final +from typing import Final from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -53,7 +53,7 @@ class ArizePhoenixClient: # Initialize HTTPHandler self.http_handler = HTTPHandler(disable_default_headers=True) - def get_prompt_version(self, prompt_version_id: str) -> dict[str, Any] | None: + def get_prompt_version(self, prompt_version_id: str) -> dict[str, object] | None: """ Fetch a prompt version from Arize Phoenix. diff --git a/litellm/integrations/clickhouse/clickhouse_spend_logger.py b/litellm/integrations/clickhouse/clickhouse_spend_logger.py index 585ea6c2afc..21c5e828152 100644 --- a/litellm/integrations/clickhouse/clickhouse_spend_logger.py +++ b/litellm/integrations/clickhouse/clickhouse_spend_logger.py @@ -80,7 +80,7 @@ def _json(value: object) -> str: return value if isinstance(value, str) else json.dumps(value, default=str) -def _json_mapping(value: Mapping[str, Any]) -> str: +def _json_mapping(value: Mapping[str, object]) -> str: return _json(dict(value)) diff --git a/litellm/integrations/cloudzero/cz_stream_api.py b/litellm/integrations/cloudzero/cz_stream_api.py index 81d08e6a4f4..1e0413e8576 100644 --- a/litellm/integrations/cloudzero/cz_stream_api.py +++ b/litellm/integrations/cloudzero/cz_stream_api.py @@ -209,7 +209,7 @@ class CloudZeroStreamer: return payload - def _convert_cbf_to_api_format(self, row: dict[str, object]) -> dict[str, Any] | None: + def _convert_cbf_to_api_format(self, row: dict[str, object]) -> dict[str, object] | None: """Convert CBF row to CloudZero API format - keeping CBF field names as CloudZero expects them.""" try: # CloudZero expects CBF format field names directly, not converted names diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 45b20211ac5..09ca4857789 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -511,7 +511,7 @@ class CustomGuardrail(CustomLogger): sticky_session_routing=self.sticky_session_routing, ) - def _get_session_id_from_request_data(self, request_data: dict[str, Any]) -> str | None: + def _get_session_id_from_request_data(self, request_data: dict[str, object]) -> str | None: """Extract session_id from request data.""" return get_session_id_from_request_data(request_data) diff --git a/litellm/integrations/dotprompt/prompt_manager.py b/litellm/integrations/dotprompt/prompt_manager.py index 47c19b12b65..0eda266740c 100644 --- a/litellm/integrations/dotprompt/prompt_manager.py +++ b/litellm/integrations/dotprompt/prompt_manager.py @@ -70,7 +70,7 @@ class PromptManager: self, prompt_id: str | None = None, prompt_directory: str | None = None, - prompt_data: dict[str, dict[str, Any]] | None = None, + prompt_data: dict[str, dict[str, object]] | None = None, prompt_file: str | None = None, ): self.prompt_directory = Path(prompt_directory) if prompt_directory else None @@ -372,6 +372,6 @@ class PromptManager: } return result - def load_prompts_from_json_data(self, prompt_data: dict[str, dict[str, Any]]) -> None: + def load_prompts_from_json_data(self, prompt_data: dict[str, dict[str, object]]) -> None: """Load additional prompts from JSON data (merges with existing prompts).""" self._load_prompts_from_json(prompt_data) diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index dce51b8190b..449c7a2f9a3 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -115,7 +115,7 @@ class FocusLogger(CustomLogger): pod_lock_manager: PodLockManager | None = None if proxy_logging_obj is not None: - writer: Final = getattr(proxy_logging_obj, "db_spend_update_writer", None) + writer: Final[object] = getattr(proxy_logging_obj, "db_spend_update_writer", None) if writer is not None: pod_lock_manager = getattr(writer, "pod_lock_manager", None) diff --git a/litellm/integrations/vantage/vantage_logger.py b/litellm/integrations/vantage/vantage_logger.py index 48a492fdd72..79beac78fcc 100644 --- a/litellm/integrations/vantage/vantage_logger.py +++ b/litellm/integrations/vantage/vantage_logger.py @@ -16,6 +16,8 @@ from litellm.integrations.focus.focus_logger import FocusLogger if TYPE_CHECKING: from apscheduler.schedulers.asyncio import AsyncIOScheduler + + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager else: AsyncIOScheduler = Any @@ -91,7 +93,7 @@ class VantageLogger(FocusLogger): """ from litellm.proxy.proxy_server import proxy_logging_obj - pod_lock_manager = None + pod_lock_manager: PodLockManager | None = None if proxy_logging_obj is not None: writer: Final[object] = getattr(proxy_logging_obj, "db_spend_update_writer", None) if writer is not None: diff --git a/litellm/llms/a2a/chat/transformation.py b/litellm/llms/a2a/chat/transformation.py index 6171b341ca4..be3b5fea601 100644 --- a/litellm/llms/a2a/chat/transformation.py +++ b/litellm/llms/a2a/chat/transformation.py @@ -7,6 +7,7 @@ from collections.abc import AsyncIterator, Iterator, Mapping from typing import TYPE_CHECKING, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.azure_ai.common_utils import AZURE_ENTRA_LITELLM_PARAM_KEYS, get_azure_ai_agent_entra_token from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator @@ -30,6 +31,7 @@ if TYPE_CHECKING: _REGISTRY_PARAMS_KEPT_OUT_OF_OPTIONAL_PARAMS: Final = ( frozenset({"api_key", "api_base", "headers", "model"}) | AZURE_ENTRA_LITELLM_PARAM_KEYS ) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) def _card_declares_no_streaming(agent_card_params: Mapping[str, object]) -> bool: @@ -325,7 +327,7 @@ class A2AConfig(BaseConfig): # Check for JSON-RPC error if "error" in response_json: - error: Final = response_json["error"] + error: Final = _JSON_OBJECT.validate_python(response_json["error"]) raise A2AError( status_code=raw_response.status_code, message=f"A2A error: {error.get('message', 'Unknown error')}", @@ -333,7 +335,7 @@ class A2AConfig(BaseConfig): ) # Extract text from A2A response - text: Final = extract_text_from_a2a_response(response_json) + text: Final = extract_text_from_a2a_response(_JSON_OBJECT.validate_python(response_json)) # Populate model response model_response.choices = [ diff --git a/litellm/llms/a2a/common_utils.py b/litellm/llms/a2a/common_utils.py index 030c5bc222e..1818b406823 100644 --- a/litellm/llms/a2a/common_utils.py +++ b/litellm/llms/a2a/common_utils.py @@ -22,7 +22,7 @@ class A2AError(BaseLLMException): self, status_code: int, message: str, - headers: dict[str, Any] = {}, + headers: dict[str, str] = {}, ): super().__init__( status_code=status_code, diff --git a/litellm/llms/aiml/image_generation/transformation.py b/litellm/llms/aiml/image_generation/transformation.py index b585d35a2ac..931038689cc 100644 --- a/litellm/llms/aiml/image_generation/transformation.py +++ b/litellm/llms/aiml/image_generation/transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, @@ -21,6 +23,10 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object]]] = TypeAdapter( + Mapping[str, object] | str | Sequence[object], config=ConfigDict(hide_input_in_errors=True) +) OPENAI_STYLE_IMAGE_MODEL_PREFIXES: Final[tuple[str, ...]] = ("openai/",) @@ -198,29 +204,33 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): if "data" in response_data and isinstance(response_data["data"], list): # Handle OpenAI-like format: {"data": [{"url": "...", "width": 1024, "height": 768, "content_type": "image/jpeg"}]} - for image in response_data["data"]: + for raw_image in _JSON_CONTAINER.validate_python(response_data["data"]): + image = raw_image if isinstance(raw_image, str) else _JSON_CONTAINER.validate_python(raw_image) if "url" in image: + url_image = _JSON_OBJECT.validate_python(image) model_response.data.append( ImageObject( b64_json=None, - url=image["url"], - revised_prompt=image.get("revised_prompt"), + url=url_image["url"], + revised_prompt=url_image.get("revised_prompt"), ) ) elif "b64_json" in image or "image_base64" in image: + b64_image = _JSON_OBJECT.validate_python(image) model_response.data.append( ImageObject( - b64_json=image.get("b64_json") or image.get("image_base64"), + b64_json=b64_image.get("b64_json") or b64_image.get("image_base64"), url=None, - revised_prompt=image.get("revised_prompt"), + revised_prompt=b64_image.get("revised_prompt"), ) ) elif "output" in response_data and "choices" in response_data["output"]: - for choice in response_data["output"]["choices"]: + for raw_choice in _JSON_CONTAINER.validate_python(response_data["output"]["choices"]): + choice = raw_choice if isinstance(raw_choice, str) else _JSON_CONTAINER.validate_python(raw_choice) if "image_base64" in choice: model_response.data.append( ImageObject( - b64_json=choice["image_base64"], + b64_json=_JSON_OBJECT.validate_python(choice)["image_base64"], url=None, ) ) @@ -228,23 +238,24 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): model_response.data.append( ImageObject( b64_json=None, - url=choice["url"], + url=_JSON_OBJECT.validate_python(choice)["url"], ) ) elif "images" in response_data: # Handle alternative format: {"images": [{"url": "...", "width": 1024, "height": 768, "content_type": "image/jpeg"}]} - for image in response_data["images"]: + for raw_image in _JSON_CONTAINER.validate_python(response_data["images"]): + image = raw_image if isinstance(raw_image, str) else _JSON_CONTAINER.validate_python(raw_image) if "url" in image: model_response.data.append( ImageObject( b64_json=None, - url=image["url"], + url=_JSON_OBJECT.validate_python(image)["url"], ) ) elif "image_base64" in image: model_response.data.append( ImageObject( - b64_json=image["image_base64"], + b64_json=_JSON_OBJECT.validate_python(image)["image_base64"], url=None, ) ) diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 401c09b580b..a45d6c1184f 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -4,6 +4,8 @@ Anthropic Token Counter implementation using the CountTokens API. from typing import Any, Final +from pydantic import JsonValue + from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler @@ -26,12 +28,12 @@ class AnthropicTokenCounter(BaseTokenCounter): async def count_tokens( self, model_to_use: str, - messages: list[dict[str, Any]] | None, + messages: list[dict[str, JsonValue]] | None, contents: list[dict[str, object]] | None, deployment: dict[str, Any] | None = None, request_model: str = "", - tools: list[dict[str, Any]] | None = None, - system: Any | None = None, + tools: list[dict[str, JsonValue]] | None = None, + system: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using Anthropic's CountTokens API. diff --git a/litellm/llms/anthropic/files/transformation.py b/litellm/llms/anthropic/files/transformation.py index 082902e8897..8433235dde4 100644 --- a/litellm/llms/anthropic/files/transformation.py +++ b/litellm/llms/anthropic/files/transformation.py @@ -14,7 +14,7 @@ Anthropic Files API endpoints: import calendar import time -from collections.abc import Mapping +from collections.abc import Iterable, Mapping from typing import Final, cast import httpx @@ -49,6 +49,7 @@ ANTHROPIC_FILES_BETA_HEADER: Final = "files-api-2025-04-14" ANTHROPIC_MESSAGE_BATCH_ID_PREFIX: Final = "msgbatch_" _JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_ITEMS: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) class AnthropicFilesConfig(BaseFilesConfig): @@ -284,9 +285,9 @@ class AnthropicFilesConfig(BaseFilesConfig): "last_id": "..." } """ - response_json: Final = raw_response.json() - files_data: Final = response_json.get("data", []) - return [self._parse_anthropic_file(f) for f in files_data] + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) + files_data: Final = _JSON_ITEMS.validate_python(response_json.get("data", [])) + return [self._parse_anthropic_file(_JSON_OBJECT.validate_python(f)) for f in files_data] def transform_file_content_request( self, diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py index 2ee23852894..d2a0849ffef 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -1347,7 +1347,7 @@ class LiteLLMAnthropicMessagesAdapter: self, choices: list[Choices], tool_name_mapping: dict[str, str] | None = None, - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: new_content: Final[list[dict[str, Any]]] = [] for choice, compaction_blocks in ( (choice, _compaction_blocks(_optional_attr(choice.message, "provider_specific_fields"))) diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 3270fb3534a..0e0b68df449 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -4,9 +4,10 @@ Azure AI Anthropic CountTokens API handler. Uses httpx for HTTP requests with Azure authentication. """ -from typing import Any, Final +from typing import Final import httpx +from pydantic import JsonValue import litellm from litellm._logging import verbose_logger @@ -28,14 +29,14 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): async def handle_count_tokens_request( self, model: str, - messages: list[dict[str, Any]], + messages: list[dict[str, JsonValue]], api_key: str, api_base: str, litellm_params: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, - tools: list[dict[str, Any]] | None = None, + tools: list[dict[str, JsonValue]] | None = None, system: object = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Handle a CountTokens request using httpx with Azure authentication. diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index c31d2c427bb..1d6d59c0592 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -5,6 +5,8 @@ Azure AI Anthropic Token Counter implementation using the CountTokens API. import os from typing import Any, Final +from pydantic import JsonValue + from litellm._logging import verbose_logger from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, @@ -28,11 +30,11 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): async def count_tokens( self, model_to_use: str, - messages: list[dict[str, object]] | None, + messages: list[dict[str, JsonValue]] | None, contents: list[dict[str, object]] | None, deployment: dict[str, Any] | None = None, request_model: str = "", - tools: list[dict[str, object]] | None = None, + tools: list[dict[str, JsonValue]] | None = None, system: object | None = None, ) -> TokenCountResponse | None: """ @@ -86,12 +88,14 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): ) if result is not None: - return TokenCountResponse( - total_tokens=result.get("input_tokens", 0), - request_model=request_model, - model_used=model_to_use, - tokenizer_type="azure_ai_anthropic_api", - original_response=result, + return TokenCountResponse.model_validate( + { + "total_tokens": result.get("input_tokens", 0), + "request_model": request_model, + "model_used": model_to_use, + "tokenizer_type": "azure_ai_anthropic_api", + "original_response": result, + } ) except AnthropicError as e: verbose_logger.warning( diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index f91a87ba0f4..9bd476efdff 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -113,7 +113,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): for index, reference_image in enumerate(images, start=1) } ) - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "prompt": prompt, "model": model, **reference_images, diff --git a/litellm/llms/base_llm/managed_resources/isolation.py b/litellm/llms/base_llm/managed_resources/isolation.py index 547bed21247..09ca23a8ded 100644 --- a/litellm/llms/base_llm/managed_resources/isolation.py +++ b/litellm/llms/base_llm/managed_resources/isolation.py @@ -10,6 +10,7 @@ they created. Callers with no admin role and no identifying ids at all are denied so an empty user_id can never select an unscoped query. """ +from collections.abc import Mapping, Sequence from typing import Any, Final from litellm.types.proxy.auth.user_api_key_auth import ( @@ -46,7 +47,7 @@ def resolve_resource_owner_id( return None -def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, Any]: +def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, object]: """Build the OpenAI-style paginated list response shape used by managed file/batch/vector-store listings. ``first_id`` and ``last_id`` are sourced from each item's ``.id`` attribute.""" @@ -61,7 +62,7 @@ def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, Any]: def build_owner_filter( user_api_key_dict: UserAPIKeyAuth, -) -> dict[str, Any] | None: +) -> dict[str, str | Sequence[Mapping[str, str]]] | None: """Return a Prisma `where` fragment that scopes a managed-resource listing to records the caller is allowed to see. diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py index 67364ccfda0..9ee5b264935 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py @@ -7,9 +7,11 @@ The main difference is in the response format: Qwen2 uses "text" field while Qwe Qwen2 + Invoke API Tutorial: https://docs.aws.amazon.com/bedrock/latest/userguide/invoke-imported-model.html """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformation import ( AmazonQwen3Config, @@ -23,6 +25,9 @@ from litellm.types.utils import ModelResponse, Usage if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_TEXT: Final = TypeAdapter(str, config=ConfigDict(strict=True, hide_input_in_errors=True)) + class AmazonQwen2Config(AmazonQwen3Config): """ @@ -54,10 +59,10 @@ class AmazonQwen2Config(AmazonQwen3Config): Qwen2 uses "text" field, but we also support "generation" field for compatibility. """ try: - response_data: Final = raw_response.json() + response_data: Final = _JSON_OBJECT.validate_python(raw_response.json()) # Extract the generated text - Qwen2 uses "text" field, but also support "generation" for compatibility - generated_text = response_data.get("generation", "") or response_data.get("text", "") + generated_text = _TEXT.validate_python(response_data.get("generation", "") or response_data.get("text", "")) # Clean up the response (remove assistant start token if present) generated_text = generated_text.removeprefix("<|im_start|>assistant\n") @@ -71,14 +76,16 @@ class AmazonQwen2Config(AmazonQwen3Config): # Set usage information if available in response if "usage" in response_data: - usage_data: Final = response_data["usage"] + usage_data: Final = _JSON_OBJECT.validate_python(response_data["usage"]) setattr( model_response, "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.model_validate( + { + "prompt_tokens": usage_data.get("prompt_tokens", 0), + "completion_tokens": usage_data.get("completion_tokens", 0), + "total_tokens": usage_data.get("total_tokens", 0), + } ), ) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py index 04248a785dd..1e31f0292c0 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py @@ -10,6 +10,7 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( @@ -22,6 +23,9 @@ from litellm.types.utils import ModelResponse, Usage if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_TEXT: Final = TypeAdapter(str, config=ConfigDict(strict=True, hide_input_in_errors=True)) + class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): """ @@ -179,10 +183,10 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): Transform Qwen3 Bedrock response to OpenAI format """ try: - response_data: Final = raw_response.json() + response_data: Final = _JSON_OBJECT.validate_python(raw_response.json()) # Extract the generated text - Qwen3 uses "generation" field - generated_text = response_data.get("generation", "") + generated_text = _TEXT.validate_python(response_data.get("generation", "")) # Clean up the response (remove assistant start token if present) generated_text = generated_text.removeprefix("<|im_start|>assistant\n") @@ -196,14 +200,16 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): # Set usage information if available in response if "usage" in response_data: - usage_data: Final = response_data["usage"] + usage_data: Final = _JSON_OBJECT.validate_python(response_data["usage"]) setattr( model_response, "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.model_validate( + { + "prompt_tokens": usage_data.get("prompt_tokens", 0), + "completion_tokens": usage_data.get("completion_tokens", 0), + "total_tokens": usage_data.get("total_tokens", 0), + } ), ) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index af833d93084..2356efe0f48 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -51,6 +51,10 @@ _GUARDRAIL_CONFIG_EXPECTED_FORMAT: Final = ( "{'guardrailIdentifier': str, 'guardrailVersion': str, 'trace': 'enabled'|'disabled'|'enabled_full'}" ) +_INVOKE_PROVIDERS: Final = TypeAdapter(tuple[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, ...]).validate_python( + get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL) +) + def _bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]": try: @@ -568,7 +572,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): if provider is not None: return provider - for provider in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): + for provider in _INVOKE_PROVIDERS: if provider in model: return provider return None diff --git a/litellm/llms/brave/search/transformation.py b/litellm/llms/brave/search/transformation.py index 1f78f8d154c..459bc3193cb 100644 --- a/litellm/llms/brave/search/transformation.py +++ b/litellm/llms/brave/search/transformation.py @@ -6,15 +6,20 @@ Documentation: https://api-dashboard.search.brave.com/app/documentation/web-sear from __future__ import annotations import re +from collections.abc import Iterable, Mapping from datetime import datetime, timezone from typing import Final, Literal, TypedDict import httpx from dateutil import parser +from pydantic import ConfigDict, TypeAdapter _ISO_YMD: Final = re.compile(r"^\s*\d{4}[-/]\d{1,2}[-/]\d{1,2}\s*$") _UNIX_TIMESTAMP: Final = re.compile(r"^\s*-?\d+(\.\d+)?\s*$") BRAVE_SECTIONS: Final = ["web", "discussions", "faqs", "faq", "news", "videos"] +_JSON_VALUE: Final = TypeAdapter(object) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( @@ -26,7 +31,7 @@ from litellm.secret_managers.main import get_secret_str def to_yyyy_mm_dd( - s: str | float | None, + s: object, *, dayfirst: bool = False, yearfirst: bool = False, @@ -246,7 +251,7 @@ class BraveSearchConfig(BaseSearchConfig): """ Transform Brave Search API response to LiteLLM unified SearchResponse format. """ - response_json: Final = raw_response.json() + response_body: Final = _JSON_VALUE.validate_python(raw_response.json()) # Transform results to SearchResult objects results: Final[list[SearchResult]] = [] @@ -254,9 +259,11 @@ class BraveSearchConfig(BaseSearchConfig): query_params: Final = raw_response.request.url.params if raw_response.request else {} sections_to_process: Final = self._sections_from_params(dict(query_params)) max_results: Final = max(1, min(int(query_params.get("count", 20)), 20)) + response_json: Final = _JSON_OBJECT.validate_python(response_body) for section in sections_to_process: - for result in response_json.get(section, {}).get("results", []): + section_json = _JSON_OBJECT.validate_python(response_json.get(section, {})) + for raw_result in _JSON_VALUES.validate_python(section_json.get("results", [])): # Because the `max_results`/`count` parameters do not affect # the number of "discussion", "faq", "news", or "videos" # results, we need to manually limit the number of results @@ -264,18 +271,21 @@ class BraveSearchConfig(BaseSearchConfig): if len(results) >= max_results: break + result = _JSON_OBJECT.validate_python(raw_result) title = result.get("title", "") url = result.get("url", "") snippet = result.get("description", "") date = to_yyyy_mm_dd(result.get("page_age") or result.get("age")) last_updated = to_yyyy_mm_dd(result.get("fetched_content_timestamp", "")) - search_result = SearchResult( - title=title, - url=url, - snippet=snippet, - date=date, - last_updated=last_updated, + search_result = SearchResult.model_validate( + { + "title": title, + "url": url, + "snippet": snippet, + "date": date, + "last_updated": last_updated, + } ) results.append(search_result) diff --git a/litellm/llms/cohere/embed/transformation.py b/litellm/llms/cohere/embed/transformation.py index bac899c4142..73b58254a3d 100644 --- a/litellm/llms/cohere/embed/transformation.py +++ b/litellm/llms/cohere/embed/transformation.py @@ -10,10 +10,11 @@ Convers Docs - https://docs.cohere.com/v2/reference/embed """ -from collections.abc import Sized +from collections.abc import Iterable, Mapping, Sized from typing import Final, Protocol, cast import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm import COHERE_DEFAULT_EMBEDDING_INPUT_TYPE @@ -30,6 +31,9 @@ from litellm.utils import is_base64_encoded from ..common_utils import CohereError +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_EMBEDDINGS_BY_TYPE: Final = TypeAdapter(Mapping[str, Iterable[object]], config=ConfigDict(hide_input_in_errors=True)) + class _SupportsEncode(Protocol): """Tokenizer handle: the embedding usage path only encodes text to measure its token length.""" @@ -198,7 +202,8 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig): 'usage' } """ - embeddings: Final = response_json["embeddings"] + payload: Final = _JSON_OBJECT.validate_python(response_json) + embeddings: Final = _EMBEDDINGS_BY_TYPE.validate_python(payload["embeddings"]) output_data: Final = [] for k, embedding_list in embeddings.items(): for idx, embedding in enumerate(embedding_list): @@ -213,7 +218,7 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig): setattr( model_response, "usage", - self._calculate_usage(input, encoding, response_json.get("meta", {})), + self._calculate_usage(input, encoding, dict(_JSON_OBJECT.validate_python(payload.get("meta", {})))), ) return model_response diff --git a/litellm/llms/cometapi/embed/transformation.py b/litellm/llms/cometapi/embed/transformation.py index adcecd5c6c9..60f781be261 100644 --- a/litellm/llms/cometapi/embed/transformation.py +++ b/litellm/llms/cometapi/embed/transformation.py @@ -2,9 +2,11 @@ CometAPI Embedding API support - OpenAI compatible """ +from collections.abc import Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -15,6 +17,9 @@ from litellm.types.utils import EmbeddingResponse, Usage from ..common_utils import CometAPIException +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_INT: Final = TypeAdapter(int) + class CometAPIEmbeddingConfig(BaseEmbeddingConfig): """ @@ -136,9 +141,10 @@ class CometAPIEmbeddingConfig(BaseEmbeddingConfig): model_response.data = raw_response_json.get("data") model_response.object = raw_response_json.get("object") + usage_payload: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {})) usage: Final = Usage( - prompt_tokens=raw_response_json.get("usage", {}).get("prompt_tokens", 0), - total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0), + prompt_tokens=_INT.validate_python(usage_payload.get("prompt_tokens", 0) or 0), + total_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0), ) model_response.usage = usage diff --git a/litellm/llms/dashscope/embed/transformation.py b/litellm/llms/dashscope/embed/transformation.py index 63ee984a65c..e5a987ca1ad 100644 --- a/litellm/llms/dashscope/embed/transformation.py +++ b/litellm/llms/dashscope/embed/transformation.py @@ -147,7 +147,7 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): ) if "error" in response_json: - error: Final = response_json["error"] + error: Final[object] = response_json["error"] message: Final = error.get("message", str(error)) if isinstance(error, dict) else str(error) raise DashScopeError( status_code=raw_response.status_code, diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 94a74d703df..4b6f178fa1a 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -475,7 +475,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): # Move message-level cache_control into a content block when content is a string. if "cache_control" in _message and isinstance(_message.get("content"), str): _message = self._move_cache_control_into_string_content_block(_message) - _sanitize_empty_content(cast(dict[str, Any], _message)) + _sanitize_empty_content(cast(dict[str, object], _message)) if _is_bare_assistant_message(_message): continue new_messages.append(_message) @@ -504,7 +504,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): """ content: Final = message.get("content") # Create new message with cache_control moved into content block - transformed_message: Final = cast(dict[str, Any], message.copy()) + transformed_message: Final = cast(dict[str, object], message.copy()) cache_control: Final = transformed_message.pop("cache_control") transformed_message["content"] = [ { diff --git a/litellm/llms/dataforseo/search/transformation.py b/litellm/llms/dataforseo/search/transformation.py index fcd4ae70645..5cf7735325d 100644 --- a/litellm/llms/dataforseo/search/transformation.py +++ b/litellm/llms/dataforseo/search/transformation.py @@ -4,9 +4,11 @@ Calls DataForSEO SERP API to search the web. DataForSEO API Reference: https://docs.dataforseo.com/v3/serp/google/organic/live/advanced/?bash """ +from collections.abc import Iterable, Mapping from typing import Final, Literal import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.constants import DEFAULT_DATAFORSEO_LOCATION_CODE from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -17,6 +19,9 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.secret_managers.main import get_secret_str +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) + class DataForSEOSearchConfig(BaseSearchConfig): """ @@ -201,18 +206,21 @@ class DataForSEOSearchConfig(BaseSearchConfig): if task.get("status_code") == 20000 and "result" in task: # Result is an array, take first element if len(task["result"]) > 0: - result: Final = task["result"][0] + result: Final = _JSON_OBJECT.validate_python(task["result"][0]) # Items contain the actual search results - for item in result.get("items", []): + for raw_item in _JSON_VALUES.validate_python(result.get("items", [])): + item = _JSON_OBJECT.validate_python(raw_item) # Only process organic search results if item.get("type") == "organic": - search_result = SearchResult( - title=item.get("title", ""), - url=item.get("url", ""), - snippet=item.get("description", ""), - date=None, # DataForSEO doesn't provide date in standard response - last_updated=None, + search_result = SearchResult.model_validate( + { + "title": item.get("title", ""), + "url": item.get("url", ""), + "snippet": item.get("description", ""), + "date": None, # DataForSEO doesn't provide date in standard response + "last_updated": None, + } ) results.append(search_result) diff --git a/litellm/llms/deepagents/harness/transformation.py b/litellm/llms/deepagents/harness/transformation.py index 1f8728230bc..cdac6944019 100644 --- a/litellm/llms/deepagents/harness/transformation.py +++ b/litellm/llms/deepagents/harness/transformation.py @@ -14,7 +14,7 @@ import json from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol from litellm.harness.options import DeepAgentsOptions from litellm.harness.types import ( @@ -73,7 +73,7 @@ _APPROVAL_DECISIONS: Final = ("approve", "reject") def chat_model_kwargs( ctx: SessionContext, -) -> dict[str, Any]: # mutable-ok: ChatLiteLLM constructor kwargs, splatted as **kwargs +) -> dict[str, object]: # mutable-ok: ChatLiteLLM constructor kwargs, splatted as **kwargs """ChatLiteLLM constructor kwargs for gateway or SDK mode.""" if not ctx.model: raise ValueError("Harness.DEEPAGENTS needs model=") @@ -107,7 +107,7 @@ def blocked_tools(permissions: str, disable_tools: Sequence[str]) -> frozenset[s def interrupt_config( permissions: str, blocked: frozenset[str] -) -> dict[str, Any] | None: # mutable-ok: deepagents create_deep_agent(interrupt_on=) takes a dict +) -> dict[str, dict[str, list[str]]] | None: # mutable-ok: deepagents create_deep_agent(interrupt_on=) takes a dict """interrupt_on for permissions='ask': approve/reject every mutating built-in.""" if permissions != "ask": return None @@ -221,9 +221,14 @@ def _message_events(message: object, skip_tools: frozenset[str]) -> tuple[Event, return () +class InterruptLike(Protocol): + @property + def id(self) -> str: ... + + def interrupts_in( update: object, -) -> list[Any]: # mutable-ok: returns a list; existing callers/tests compare it to list literals +) -> list[InterruptLike]: # mutable-ok: returns a list; existing callers/tests compare it to list literals if not isinstance(update, Mapping): return [] # mutable-ok: returns a list; existing callers/tests compare it to list literals found = update.get("__interrupt__") @@ -260,7 +265,7 @@ def approval_requests( return list(kept) # mutable-ok: list return; callers/tests compare to lists -def decision(allowed: bool, reason: str) -> dict[str, Any]: # mutable-ok: LangGraph resume payload (HITL decision dict) +def decision(allowed: bool, reason: str) -> dict[str, str]: # mutable-ok: LangGraph resume payload (HITL decision dict) if allowed: return {"type": "approve"} # mutable-ok: LangGraph resume payload (HITL decision dict) return { # mutable-ok: LangGraph HITL decision @@ -273,7 +278,7 @@ def decision(allowed: bool, reason: str) -> dict[str, Any]: # mutable-ok: LangG class TurnState: """Mutable state across the stream passes of one turn.""" - interrupts: tuple[Any, ...] = () + interrupts: tuple[InterruptLike, ...] = () class DeepAgentsHarnessConfig(BaseHarnessConfig): diff --git a/litellm/llms/duckduckgo/search/transformation.py b/litellm/llms/duckduckgo/search/transformation.py index 93e70f65bcd..1816a0d24e4 100644 --- a/litellm/llms/duckduckgo/search/transformation.py +++ b/litellm/llms/duckduckgo/search/transformation.py @@ -4,10 +4,12 @@ Calls DuckDuckGo's Instant Answer API to search the web. DuckDuckGo API Reference: https://duckduckgo.com/api """ +from collections.abc import Iterable, Mapping from typing import Final, Literal, TypedDict from urllib.parse import urlencode import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( @@ -17,6 +19,9 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.secret_managers.main import get_secret_str +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) + class _DuckDuckGoSearchRequestRequired(TypedDict): """Required fields for DuckDuckGo Search API request.""" @@ -159,7 +164,7 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): Returns: SearchResponse with standardized format """ - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) # Extract max_results from the request URL params query_params: Final = raw_response.request.url.params if raw_response.request else {} @@ -178,17 +183,19 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): # Check if there's an Abstract with URL if response_json.get("AbstractURL") and response_json.get("AbstractText"): - abstract_result: Final = SearchResult( - title=response_json.get("Heading", ""), - url=response_json.get("AbstractURL", ""), - snippet=response_json.get("AbstractText", ""), - date=None, - last_updated=None, + abstract_result: Final = SearchResult.model_validate( + { + "title": response_json.get("Heading", ""), + "url": response_json.get("AbstractURL", ""), + "snippet": response_json.get("AbstractText", ""), + "date": None, + "last_updated": None, + } ) results.append(abstract_result) # Process RelatedTopics - related_topics: Final = response_json.get("RelatedTopics", []) + related_topics: Final = _JSON_VALUES.validate_python(response_json.get("RelatedTopics", [])) for topic in related_topics: # Stop if we've reached max_results if max_results is not None and len(results) >= max_results: diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index cd904e493f2..c3c29e62f57 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -4,9 +4,11 @@ Calls Exa AI's /search endpoint to search the web. Exa AI API Reference: https://docs.exa.ai/reference/search """ +from collections.abc import Iterable, Mapping from typing import Final, TypedDict import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( @@ -16,6 +18,10 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.secret_managers.main import get_secret_str +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) +_STRS: Final = TypeAdapter(Iterable[str], config=ConfigDict(hide_input_in_errors=True)) + class _ExaAISearchRequestRequired(TypedDict): """Required fields for Exa AI Search API request.""" @@ -177,20 +183,23 @@ class ExaAISearchConfig(BaseSearchConfig): Returns: SearchResponse with standardized format """ - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) # Transform results to SearchResult objects results: Final = [] - for result in response_json.get("results", []): - search_result = SearchResult( - title=result.get("title", ""), - url=result.get("url", ""), - snippet=result.get("text") - or "\n\n".join(result.get("highlights") or []) - or result.get("summary") - or "", - date=result.get("publishedDate"), # ISO 8601 datetime string - last_updated=None, # Exa AI doesn't provide last_updated in response + for raw_result in _JSON_VALUES.validate_python(response_json.get("results", [])): + result = _JSON_OBJECT.validate_python(raw_result) + search_result = SearchResult.model_validate( + { + "title": result.get("title", ""), + "url": result.get("url", ""), + "snippet": result.get("text") + or "\n\n".join(_STRS.validate_python(result.get("highlights") or [])) + or result.get("summary") + or "", + "date": result.get("publishedDate"), # ISO 8601 datetime string + "last_updated": None, # Exa AI doesn't provide last_updated in response + } ) results.append(search_result) diff --git a/litellm/llms/fal_ai/image_edit/transformation.py b/litellm/llms/fal_ai/image_edit/transformation.py index d77769b130d..ea4e8cac5bf 100644 --- a/litellm/llms/fal_ai/image_edit/transformation.py +++ b/litellm/llms/fal_ai/image_edit/transformation.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable import httpx from httpx._types import RequestFiles +from pydantic import ConfigDict, TypeAdapter from litellm.images.utils import ImageEditRequestUtils from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig @@ -23,6 +24,8 @@ from litellm.types.utils import FileTypes, ImageResponse if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +_JSON_VALUE: Final = TypeAdapter(object) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) DEFAULT_BASE_URL: Final[str] = "https://fal.run" EDIT_SUFFIX: Final[str] = "/edit" SUPPORTED_OPENAI_PARAMS: Final[tuple[str, ...]] = ("background", "mask", "n", "quality", "size") @@ -163,7 +166,7 @@ class FalAIImageEditConfig(BaseImageEditConfig): logging_obj: "LiteLLMLoggingObj", ) -> ImageResponse: try: - response_json: Final = raw_response.json() + response_json: Final = _JSON_VALUE.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error parsing Fal AI image edit response: {e}", @@ -171,5 +174,7 @@ class FalAIImageEditConfig(BaseImageEditConfig): headers=raw_response.headers, ) model_response: Final = ImageResponse() - model_response.data = list(fal_images_to_image_objects(response_json.get("images", ()))) + model_response.data = list( + fal_images_to_image_objects(_JSON_OBJECT.validate_python(response_json).get("images", ())) + ) return model_response diff --git a/litellm/llms/fal_ai/image_generation/bria_transformation.py b/litellm/llms/fal_ai/image_generation/bria_transformation.py index 53da48e62ca..df8019747ef 100644 --- a/litellm/llms/fal_ai/image_generation/bria_transformation.py +++ b/litellm/llms/fal_ai/image_generation/bria_transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams from litellm.types.utils import ImageObject, ImageResponse @@ -15,6 +17,9 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_VALUE: Final = TypeAdapter(object) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class FalAIBriaConfig(FalAIBaseConfig): """ @@ -206,7 +211,7 @@ class FalAIBriaConfig(FalAIBaseConfig): } """ try: - response_data: Final = raw_response.json() + response_data: Final = _JSON_VALUE.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", @@ -218,7 +223,7 @@ class FalAIBriaConfig(FalAIBaseConfig): model_response.data = [] # Handle Bria response format - uses "image" (singular) not "images" - image_data: Final = response_data.get("image") + image_data: Final = _JSON_OBJECT.validate_python(response_data).get("image") if image_data and isinstance(image_data, dict): model_response.data.append( ImageObject( diff --git a/litellm/llms/fal_ai/image_generation/imagen4_transformation.py b/litellm/llms/fal_ai/image_generation/imagen4_transformation.py index c5a5f0112ca..09e96926767 100644 --- a/litellm/llms/fal_ai/image_generation/imagen4_transformation.py +++ b/litellm/llms/fal_ai/image_generation/imagen4_transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, cast import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams @@ -16,6 +18,9 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_VALUE: Final = TypeAdapter(object) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class FalAIImagen4Config(FalAIBaseConfig): """ @@ -203,7 +208,7 @@ class FalAIImagen4Config(FalAIBaseConfig): } """ try: - response_data: Final = raw_response.json() + response_data: Final = _JSON_VALUE.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", @@ -215,7 +220,8 @@ class FalAIImagen4Config(FalAIBaseConfig): model_response.data = [] # Handle Imagen4 response format - images: Final = response_data.get("images", []) + response_object: Final = _JSON_OBJECT.validate_python(response_data) + images: Final = response_object.get("images", []) if isinstance(images, list): for image_data in images: if isinstance(image_data, dict): @@ -239,7 +245,7 @@ class FalAIImagen4Config(FalAIBaseConfig): hidden_params: Final = cast( # cast-ok: preserve mapping operations on dynamic response metadata dict[str, object], getattr(model_response, HIDDEN_PARAMS_ATTR) ) - if "seed" in response_data: - hidden_params["seed"] = response_data["seed"] + if "seed" in response_object: + hidden_params["seed"] = response_object["seed"] return model_response diff --git a/litellm/llms/fal_ai/image_generation/transformation.py b/litellm/llms/fal_ai/image_generation/transformation.py index 8f081c7228d..0bb5dada071 100644 --- a/litellm/llms/fal_ai/image_generation/transformation.py +++ b/litellm/llms/fal_ai/image_generation/transformation.py @@ -2,7 +2,7 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx -from pydantic import TypeAdapter +from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm.llms.base_llm.image_generation.transformation import ( @@ -31,6 +31,8 @@ class FalImageProviderSpecificFields(TypedDict, total=False): _FAL_IMAGE_DATA: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_JSON_VALUE: Final = TypeAdapter(object) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) def fal_images_to_image_objects(images: object) -> tuple[ImageObject, ...]: @@ -124,7 +126,7 @@ class FalAIBaseConfig(BaseImageGenerationConfig): Transform the image generation response to the litellm image response """ try: - response_data: Final = raw_response.json() + response_data: Final = _JSON_VALUE.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", @@ -134,7 +136,9 @@ class FalAIBaseConfig(BaseImageGenerationConfig): if not model_response.data: model_response.data = [] - model_response.data.extend(fal_images_to_image_objects(response_data.get("images", ()))) + model_response.data.extend( + fal_images_to_image_objects(_JSON_OBJECT.validate_python(response_data).get("images", ())) + ) return model_response diff --git a/litellm/llms/fastcrw/search/transformation.py b/litellm/llms/fastcrw/search/transformation.py index 2ded3208f2e..0aef446c23b 100644 --- a/litellm/llms/fastcrw/search/transformation.py +++ b/litellm/llms/fastcrw/search/transformation.py @@ -8,9 +8,11 @@ or cloud). The search response uses the Firecrawl-compatible envelope fastCRW API Reference: https://fastcrw.com/docs/rest-api """ +from collections.abc import Iterable, Mapping from typing import Final, TypedDict import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( @@ -20,6 +22,9 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.secret_managers.main import get_secret_str +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) + class _FastCRWSearchRequestRequired(TypedDict): """Required fields for fastCRW Search API request.""" @@ -157,21 +162,24 @@ class FastCRWSearchConfig(BaseSearchConfig): Returns: SearchResponse with standardized format """ - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) results: Final = [] data: Final = response_json.get("data", []) if isinstance(data, list): - for result in data: + for raw_result in _JSON_VALUES.validate_python(data): + result = _JSON_OBJECT.validate_python(raw_result) snippet = result.get("markdown") or result.get("description", "") - search_result = SearchResult( - title=result.get("title", ""), - url=result.get("url", ""), - snippet=snippet, - date=None, - last_updated=None, + search_result = SearchResult.model_validate( + { + "title": result.get("title", ""), + "url": result.get("url", ""), + "snippet": snippet, + "date": None, + "last_updated": None, + } ) results.append(search_result) diff --git a/litellm/llms/firecrawl/search/transformation.py b/litellm/llms/firecrawl/search/transformation.py index 552388ce47c..736904c4e36 100644 --- a/litellm/llms/firecrawl/search/transformation.py +++ b/litellm/llms/firecrawl/search/transformation.py @@ -4,9 +4,11 @@ Calls Firecrawl's /search endpoint to search the web. Firecrawl API Reference: https://docs.firecrawl.dev/api-reference/endpoint/search """ +from collections.abc import Iterable, Mapping from typing import Final, TypedDict import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( @@ -16,6 +18,9 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.secret_managers.main import get_secret_str +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) + class _FirecrawlSearchRequestRequired(TypedDict): """Required fields for Firecrawl Search API request.""" @@ -170,7 +175,7 @@ class FirecrawlSearchConfig(BaseSearchConfig): Returns: SearchResponse with standardized format """ - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) # Transform results to SearchResult objects results: Final = [] @@ -179,44 +184,54 @@ class FirecrawlSearchConfig(BaseSearchConfig): if isinstance(data, list): # Self-hosted Firecrawl (v1) format: data is a flat list of results - for result in data: + for raw_result in _JSON_VALUES.validate_python(data): + result = _JSON_OBJECT.validate_python(raw_result) snippet = result.get("markdown") or result.get("description", "") - search_result = SearchResult( - title=result.get("title", ""), - url=result.get("url", ""), - snippet=snippet, - date=None, - last_updated=None, + search_result = SearchResult.model_validate( + { + "title": result.get("title", ""), + "url": result.get("url", ""), + "snippet": snippet, + "date": None, + "last_updated": None, + } ) results.append(search_result) elif isinstance(data, dict): # Firecrawl Cloud (v2) format: data is a dict with web/news keys - web_results: Final = data.get("web", []) + sections: Final = _JSON_OBJECT.validate_python(data) + web_results: Final = _JSON_VALUES.validate_python(sections.get("web", [])) - for result in web_results: + for raw_result in web_results: + result = _JSON_OBJECT.validate_python(raw_result) # Use markdown if available, otherwise fall back to description snippet = result.get("markdown") or result.get("description", "") - search_result = SearchResult( - title=result.get("title", ""), - url=result.get("url", ""), - snippet=snippet, - date=None, - last_updated=None, + search_result = SearchResult.model_validate( + { + "title": result.get("title", ""), + "url": result.get("url", ""), + "snippet": snippet, + "date": None, + "last_updated": None, + } ) results.append(search_result) # Process news results if available (they have date field) - news_results: Final = data.get("news", []) - for result in news_results: + news_results: Final = _JSON_VALUES.validate_python(sections.get("news", [])) + for raw_result in news_results: + result = _JSON_OBJECT.validate_python(raw_result) snippet = result.get("markdown") or result.get("snippet", "") - search_result = SearchResult( - title=result.get("title", ""), - url=result.get("url", ""), - snippet=snippet, - date=result.get("date"), # News results include date - last_updated=None, + search_result = SearchResult.model_validate( + { + "title": result.get("title", ""), + "url": result.get("url", ""), + "snippet": snippet, + "date": result.get("date"), # News results include date + "last_updated": None, + } ) results.append(search_result) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 2c240544ee8..bd7c0d11b7d 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -1,8 +1,9 @@ import json -from collections.abc import AsyncIterator, Iterator, Mapping +from collections.abc import AsyncIterator, Iterable, Iterator, Mapping from typing import TYPE_CHECKING, Final, Literal, cast import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_logger @@ -51,6 +52,10 @@ from ..common_utils import ( if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer +_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_STR: Final = TypeAdapter(str) + def _map_reasoning_effort(value: object) -> object: effort: Final[object] = cast(Mapping[str, object], value).get("effort") if isinstance(value, Mapping) else value @@ -757,7 +762,8 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): additional_headers: Final = get_response_headers(raw_response_headers) - response: Final = ModelResponse(**completion_response) + completion_fields: Final = _JSON_DICT.validate_python(completion_response) + response: Final = ModelResponse(**completion_fields) if response.model is not None: response.model = "fireworks_ai/" + response.model @@ -771,7 +777,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): response.hidden_params = { "additional_headers": additional_headers, - **_extract_fireworks_hidden_params(completion_response), + **_extract_fireworks_hidden_params(completion_fields), } return response @@ -832,9 +838,9 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): f"Failed to fetch models from Fireworks AI. Status code: {response.status_code}, Response: {response.json()}" ) - models: Final = response.json()["models"] + models: Final = _JSON_OBJECTS.validate_python(_JSON_DICT.validate_python(response.json())["models"]) - return ["fireworks_ai/" + model["name"] for model in models] + return ["fireworks_ai/" + _STR.validate_python(model["name"]) for model in models] @staticmethod def get_api_key(api_key: str | None = None) -> str | None: diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 4d13d0d2e6b..fe82a9cae81 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -1,4 +1,4 @@ -from collections.abc import Iterable, Mapping +from collections.abc import Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx @@ -35,6 +35,10 @@ else: _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) _JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object]]] = TypeAdapter( + Mapping[str, object] | str | Sequence[object], config=ConfigDict(hide_input_in_errors=True) +) class GoogleImageGenConfig(BaseImageGenerationConfig): @@ -199,19 +203,22 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): # Handle different response formats based on model if is_gemini_image_model(model): # Gemini Flash Image Preview models return in candidates format - candidates: Final = response_data.get("candidates", []) + response_object: Final = _JSON_OBJECT.validate_python(response_data) + candidates: Final = _JSON_CONTAINER.validate_python(response_object.get("candidates", [])) for candidate in candidates: - content = candidate.get("content", {}) - parts = content.get("parts", []) - for part in parts: + content = _JSON_OBJECT.validate_python(_JSON_OBJECT.validate_python(candidate).get("content", {})) + parts = _JSON_CONTAINER.validate_python(content.get("parts", [])) + for raw_part in parts: + part = raw_part if isinstance(raw_part, str) else _JSON_CONTAINER.validate_python(raw_part) # Look for inlineData with image if "inlineData" in part: - inline_data = part["inlineData"] + part_object = _JSON_OBJECT.validate_python(part) + inline_data = _JSON_CONTAINER.validate_python(part_object["inlineData"]) if "data" in inline_data: - thought_sig = part.get("thoughtSignature") + thought_sig = part_object.get("thoughtSignature") model_response.data.append( ImageObject( - b64_json=inline_data["data"], + b64_json=_JSON_OBJECT.validate_python(inline_data)["data"], url=None, provider_specific_fields=( {"thought_signature": thought_sig} if thought_sig else None @@ -224,12 +231,14 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): model_response.usage = transform_gemini_image_usage( _JSON_DICT.validate_python(response_data["usageMetadata"]) ) - web_search_requests: Final = get_gemini_image_web_search_requests(response_data) + web_search_requests: Final = get_gemini_image_web_search_requests(response_object) if web_search_requests and model_response.usage is not None: setattr(model_response.usage, "web_search_requests", web_search_requests) else: # Original Imagen format - predictions with generated images - predictions: Final = _JSON_OBJECTS.validate_python(response_data.get("predictions", [])) + predictions: Final = _JSON_OBJECTS.validate_python( + _JSON_OBJECT.validate_python(response_data).get("predictions", []) + ) for prediction in predictions: # Google AI returns base64 encoded images in the prediction model_response.data.append( diff --git a/litellm/llms/google_pse/search/transformation.py b/litellm/llms/google_pse/search/transformation.py index 0128db9a823..1ae465277c5 100644 --- a/litellm/llms/google_pse/search/transformation.py +++ b/litellm/llms/google_pse/search/transformation.py @@ -4,9 +4,11 @@ Calls Google Programmable Search Engine (PSE) API to search the web. Google PSE API Reference: https://developers.google.com/custom-search/v1/reference/rest/v1/cse/list """ +from collections.abc import Iterable, Mapping from typing import Final, Literal, TypedDict import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( @@ -16,6 +18,9 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.secret_managers.main import get_secret_str +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) + class _GooglePSESearchRequestRequired(TypedDict): """Required fields for Google PSE Search API request.""" @@ -239,17 +244,20 @@ class GooglePSESearchConfig(BaseSearchConfig): Returns: SearchResponse with standardized format """ - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) # Transform results to SearchResult objects results: Final = [] - for item in response_json.get("items", []): - search_result = SearchResult( - title=item.get("title", ""), - url=item.get("link", ""), - snippet=item.get("snippet", ""), - date=None, # Google PSE doesn't provide date in standard response - last_updated=None, # Google PSE doesn't provide last_updated in response + for raw_item in _JSON_VALUES.validate_python(response_json.get("items", [])): + item = _JSON_OBJECT.validate_python(raw_item) + search_result = SearchResult.model_validate( + { + "title": item.get("title", ""), + "url": item.get("link", ""), + "snippet": item.get("snippet", ""), + "date": None, # Google PSE doesn't provide date in standard response + "last_updated": None, # Google PSE doesn't provide last_updated in response + } ) results.append(search_result) diff --git a/litellm/llms/infinity/embedding/transformation.py b/litellm/llms/infinity/embedding/transformation.py index 26130e4ca6e..f289918cb3c 100644 --- a/litellm/llms/infinity/embedding/transformation.py +++ b/litellm/llms/infinity/embedding/transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -11,6 +13,9 @@ from litellm.types.utils import EmbeddingResponse, Usage from ..common_utils import InfinityError +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_INT: Final = TypeAdapter(int) + class InfinityEmbeddingConfig(BaseEmbeddingConfig): """ @@ -124,9 +129,10 @@ class InfinityEmbeddingConfig(BaseEmbeddingConfig): model_response.data = raw_response_json.get("data") model_response.object = raw_response_json.get("object") + usage_payload: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {})) usage: Final = Usage( - prompt_tokens=raw_response_json.get("usage", {}).get("prompt_tokens", 0), - total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0), + prompt_tokens=_INT.validate_python(usage_payload.get("prompt_tokens", 0) or 0), + total_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0), ) model_response.usage = usage return model_response diff --git a/litellm/llms/langgraph/chat/transformation.py b/litellm/llms/langgraph/chat/transformation.py index f1f4054099a..4f93f435957 100644 --- a/litellm/llms/langgraph/chat/transformation.py +++ b/litellm/llms/langgraph/chat/transformation.py @@ -10,7 +10,7 @@ Non-streaming endpoint: POST /runs/wait import json from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Final, Optional, Union import httpx @@ -369,7 +369,7 @@ class LangGraphConfig(BaseConfig): from litellm.utils import CustomStreamWrapper if client is None or not isinstance(client, AsyncHTTPHandler): - client = get_async_httpx_client(llm_provider=cast(Any, "langgraph"), params={}) + client = get_async_httpx_client(llm_provider="langgraph", params={}) verbose_logger.debug("Making async streaming request to: %s", api_base) diff --git a/litellm/llms/lemonade/chat/transformation.py b/litellm/llms/lemonade/chat/transformation.py index c4079dfc73b..82712af9af3 100644 --- a/litellm/llms/lemonade/chat/transformation.py +++ b/litellm/llms/lemonade/chat/transformation.py @@ -2,10 +2,12 @@ Translate from OpenAI's `/v1/chat/completions` to Lemonade's `/v1/chat/completions` """ +from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final from urllib.parse import quote import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_logger @@ -21,6 +23,10 @@ from ...openai_like.chat.transformation import OpenAILikeChatConfig if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer +_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True)) +_STR: Final = TypeAdapter(str) + class LemonadeChatConfig(OpenAILikeChatConfig): _DEFAULT_API_KEY = "lemonade" @@ -55,7 +61,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig): response_format: dict | None = None, tools: list | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = dict(locals()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @@ -104,8 +110,8 @@ class LemonadeChatConfig(OpenAILikeChatConfig): f"Failed to fetch models from Lemonade. Status code: {response.status_code}, Response: {response.text}" ) - model_list: Final = response.json().get("data", []) - return ["lemonade/" + model["id"] for model in model_list] + model_list: Final = _JSON_OBJECTS.validate_python(_JSON_DICT.validate_python(response.json()).get("data", [])) + return ["lemonade/" + _STR.validate_python(model["id"]) for model in model_list] @staticmethod def _get_positive_int(value: Any) -> int | None: @@ -188,10 +194,11 @@ class LemonadeChatConfig(OpenAILikeChatConfig): verbose_logger.debug("LemonadeError: Could not get model info.") return self._get_default_model_info(model) - max_input_tokens: Final = self._get_context_window(model_info) - max_output_tokens: Final = self._get_positive_int(model_info.get("max_output_tokens")) - max_tokens: Final = self._get_positive_int(model_info.get("max_tokens")) - provider_specific_entry: Final = self._get_provider_specific_entry(model_info) + model_info_fields: Final = _JSON_DICT.validate_python(model_info) + max_input_tokens: Final = self._get_context_window(model_info_fields) + max_output_tokens: Final = self._get_positive_int(model_info_fields.get("max_output_tokens")) + max_tokens: Final = self._get_positive_int(model_info_fields.get("max_tokens")) + provider_specific_entry: Final = self._get_provider_specific_entry(model_info_fields) model_info_response: Final = self._get_default_model_info(model) model_info_response.update( diff --git a/litellm/llms/linkup/search/transformation.py b/litellm/llms/linkup/search/transformation.py index 8ecd5a24c08..f20a1f5285f 100644 --- a/litellm/llms/linkup/search/transformation.py +++ b/litellm/llms/linkup/search/transformation.py @@ -4,9 +4,11 @@ Calls Linkup's /search endpoint to search the web. Linkup API Reference: https://docs.linkup.so/pages/documentation/api-reference/endpoint/post-search """ +from collections.abc import Iterable, Mapping from typing import Final, Literal, TypedDict import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( @@ -16,6 +18,9 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.secret_managers.main import get_secret_str +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) + class _LinkupSearchRequestRequired(TypedDict): """Required fields for Linkup Search API request.""" @@ -165,35 +170,40 @@ class LinkupSearchConfig(BaseSearchConfig): Returns: SearchResponse with standardized format """ - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) # Transform results to SearchResult objects results: Final = [] # Process results array - raw_results: Final = response_json.get("results", []) + raw_results: Final = _JSON_VALUES.validate_python(response_json.get("results", [])) - for result in raw_results: + for raw_result in raw_results: + result = _JSON_OBJECT.validate_python(raw_result) # Handle both text and image result types result_type = result.get("type", "text") if result_type == "text": - search_result = SearchResult( - title=result.get("name", ""), - url=result.get("url", ""), - snippet=result.get("content", ""), - date=None, - last_updated=None, + search_result = SearchResult.model_validate( + { + "title": result.get("name", ""), + "url": result.get("url", ""), + "snippet": result.get("content", ""), + "date": None, + "last_updated": None, + } ) results.append(search_result) elif result_type == "image": # For image results, use the URL as both title and snippet if name not provided - search_result = SearchResult( - title=result.get("name", result.get("url", "")), - url=result.get("url", ""), - snippet=result.get("content", ""), - date=None, - last_updated=None, + search_result = SearchResult.model_validate( + { + "title": result.get("name", result.get("url", "")), + "url": result.get("url", ""), + "snippet": result.get("content", ""), + "date": None, + "last_updated": None, + } ) results.append(search_result) diff --git a/litellm/llms/mistral/audio_speech/transformation.py b/litellm/llms/mistral/audio_speech/transformation.py index 04a7e3b9341..5a8960d0898 100644 --- a/litellm/llms/mistral/audio_speech/transformation.py +++ b/litellm/llms/mistral/audio_speech/transformation.py @@ -11,6 +11,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final import httpx +from pydantic import ConfigDict, TypeAdapter, ValidationError from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.text_to_speech.transformation import ( @@ -23,6 +24,8 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.llms.openai import HttpxBinaryResponseContent +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class MistralTextToSpeechException(BaseLLMException): pass @@ -143,7 +146,7 @@ class MistralTextToSpeechConfig(BaseTextToSpeechConfig): return request_data def _requested_content_type(self, request: httpx.Request) -> str: - request_body: Final = json.loads(request.content or b"{}") + request_body: Final = _JSON_OBJECT.validate_python(json.loads(request.content or b"{}")) requested_format: Final = request_body.get("response_format") if not isinstance(requested_format, str): return "audio/mpeg" @@ -158,7 +161,9 @@ class MistralTextToSpeechConfig(BaseTextToSpeechConfig): from litellm.types.llms.openai import HttpxBinaryResponseContent try: - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) + except ValidationError: + raise except (json.JSONDecodeError, ValueError): raise MistralTextToSpeechException( status_code=raw_response.status_code, diff --git a/litellm/llms/nlp_cloud/chat/transformation.py b/litellm/llms/nlp_cloud/chat/transformation.py index 2ff48894619..0cca299de93 100644 --- a/litellm/llms/nlp_cloud/chat/transformation.py +++ b/litellm/llms/nlp_cloud/chat/transformation.py @@ -1,5 +1,6 @@ import json import time +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -79,7 +80,7 @@ class NLPCloudConfig(BaseConfig): num_beams: int | None = None, num_return_sequences: int | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = dict(locals()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 14ef6225666..6b038872b8d 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -759,7 +759,7 @@ class OCIStreamWrapper(CustomStreamWrapper): self.sent_last_chunk = True return self.model_response_creator(chunk={"choices": parsed.choices}) - def chunk_creator(self, chunk: Any) -> ModelResponseStream | None: + def chunk_creator(self, chunk: object) -> ModelResponseStream | None: if not isinstance(chunk, str): raise ValueError(f"Chunk is not a string: {chunk}") if not chunk.startswith("data:"): diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index 409e763dde1..af8d492316e 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -1,6 +1,6 @@ import base64 import io -from typing import Any, Final +from typing import Final import httpx @@ -174,7 +174,7 @@ class OllamaModelInfo(BaseLLMModelInfo): model: str, api_base: str | None = None, api_key: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: from litellm import module_level_client model = self._strip_ollama_model_prefix(model) @@ -223,7 +223,7 @@ class OllamaModelInfo(BaseLLMModelInfo): model: str, api_base: str | None = None, api_key: str | None = None, - ) -> dict[str, Any] | None: + ) -> dict[str, object] | None: if self._is_static_ollama_model(model): return None return self.get_runtime_model_info(model=model, api_base=api_base, api_key=api_key) diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index dcae8524c4f..09946251227 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -97,6 +97,7 @@ class _OllamaGenerateResponse(TypedDict): _OLLAMA_GENERATE_RESPONSE: Final = TypeAdapter(_OllamaGenerateResponse) +_JSON_VALUE: Final = TypeAdapter(object) _JSON_CODE_FENCE: Final = re.compile(r"^```(?:json)?\s*(.*?)\s*(?:```)?$", re.DOTALL) _JSON_OBJECT_START: Final = re.compile(r"(?:```(?:json)?\s*)?\{") _JSON_OBJECT_PARTIAL_START: Final = re.compile(r"`{0,3}|```(?:j|js|jso|json)?\s*") @@ -334,7 +335,9 @@ class OllamaConfig(BaseConfig): model_response.choices[0].finish_reason = "stop" else: try: - response_content: Final[object] = json.loads(_strip_json_code_fence(response_text)) + response_content: Final = _JSON_VALUE.validate_python( + json.loads(_strip_json_code_fence(response_text)) + ) # Check if this is a function call format with name/arguments structure if ( diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index 0d7a553d633..02a05b060b8 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -6,7 +6,7 @@ Uses httpx for HTTP requests to OpenAI's /v1/responses/input_tokens endpoint. import json from collections.abc import Sequence -from typing import Any, Final +from typing import Final import httpx @@ -33,7 +33,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): timeout: float | httpx.Timeout | None = None, tools: list[dict[str, object]] | None = None, instructions: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Handle a token counting request to OpenAI's Responses API. diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index c05d943b7cf..ab3ddbd179d 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -81,12 +81,14 @@ class OpenAITokenCounter(BaseTokenCounter): ) if result is not None: - return TokenCountResponse( - total_tokens=result.get("input_tokens", 0), - request_model=request_model, - model_used=model_to_use, - tokenizer_type="openai_api", - original_response=result, + return TokenCountResponse.model_validate( + { + "total_tokens": result.get("input_tokens", 0), + "request_model": request_model, + "model_used": model_to_use, + "tokenizer_type": "openai_api", + "original_response": result, + } ) except OpenAIError as e: verbose_logger.warning("OpenAI CountTokens API error: status=%s, message=%s", e.status_code, e.message) diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index 9f596505f91..45d8ab0b4c4 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -118,7 +118,7 @@ class OpenAICountTokensConfig: self, model: str, input: str | Sequence[object], - tools: list[dict[str, Any]] | None = None, + tools: list[dict[str, object]] | None = None, instructions: str | None = None, ) -> dict[str, object]: """ diff --git a/litellm/llms/openrouter/chat/transformation.py b/litellm/llms/openrouter/chat/transformation.py index c63863bbca4..9aef15f15a3 100644 --- a/litellm/llms/openrouter/chat/transformation.py +++ b/litellm/llms/openrouter/chat/transformation.py @@ -6,11 +6,12 @@ Calls done in OpenAI/openai.py as OpenRouter is openai-compatible. Docs: https://openrouter.ai/docs/parameters """ -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Iterator, Mapping from enum import Enum -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Final, cast import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator @@ -26,6 +27,8 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + class CacheControlSupportedModels(str, Enum): """Models that support cache_control in content blocks.""" @@ -203,9 +206,9 @@ class OpenrouterConfig(OpenAIGPTConfig): # Extract cost from OpenRouter response body # OpenRouter returns cost information in the usage object when usage.include=true try: - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) if "usage" in response_json and response_json["usage"]: - response_cost: Final = response_json["usage"].get("cost") + response_cost: Final = _JSON_OBJECT.validate_python(response_json["usage"]).get("cost") if response_cost is not None: # Store cost in hidden params for the cost calculator to use if not hasattr(model_response, HIDDEN_PARAMS_ATTR): @@ -218,7 +221,8 @@ class OpenrouterConfig(OpenAIGPTConfig): additional_headers: Final = cast( # cast-ok: preserve mapping operations on response metadata dict[str, object], hidden_params["additional_headers"] ) - additional_headers["llm_provider-x-litellm-response-cost"] = float(response_cost) + if isinstance(response_cost, str | int | float): + additional_headers["llm_provider-x-litellm-response-cost"] = float(response_cost) except Exception: # If we can't extract cost, continue without it - don't fail the response pass @@ -237,7 +241,7 @@ class OpenrouterConfig(OpenAIGPTConfig): streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, sync_stream: bool, json_mode: bool | None = False, - ) -> Any: + ) -> "OpenRouterChatCompletionStreamingHandler": return OpenRouterChatCompletionStreamingHandler( streaming_response=streaming_response, sync_stream=sync_stream, diff --git a/litellm/llms/opensandbox/sandbox/transformation.py b/litellm/llms/opensandbox/sandbox/transformation.py index d2c558241f0..13c9d805638 100644 --- a/litellm/llms/opensandbox/sandbox/transformation.py +++ b/litellm/llms/opensandbox/sandbox/transformation.py @@ -1,9 +1,11 @@ import asyncio import json import time +from collections.abc import Mapping from typing import Final, cast import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.constants import ( OPEN_SANDBOX_API_BASE_ENV_VAR, @@ -37,6 +39,9 @@ DEFAULT_READY_TIMEOUT: Final = OPEN_SANDBOX_READY_TIMEOUT DEFAULT_POLL_INTERVAL: Final = OPEN_SANDBOX_POLL_INTERVAL MAX_OUTPUT_BYTES: Final = SANDBOX_MAX_OUTPUT_BYTES +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_VALUE: Final = TypeAdapter(object) + class OpenSandboxSandboxConfig(BaseSandboxConfig): def _http(self, client: AsyncHTTPHandler | None) -> AsyncHTTPHandler: @@ -92,7 +97,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): headers=self._lifecycle_headers(key), json=body, ) - data: Final = response.json() + data: Final = _JSON_OBJECT.validate_python(response.json()) sandbox_id: Final = str(data["id"]) if self._sandbox_state(data) != "Running": @@ -256,7 +261,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): url=f"{api_base}/sandboxes/{sandbox_id}", headers=headers, ) - data = response.json() + data = _JSON_VALUE.validate_python(response.json()) state = self._sandbox_state(data) if state == "Running": return diff --git a/litellm/llms/ovhcloud/embedding/transformation.py b/litellm/llms/ovhcloud/embedding/transformation.py index b78a81d05bf..cdf095e6ad1 100644 --- a/litellm/llms/ovhcloud/embedding/transformation.py +++ b/litellm/llms/ovhcloud/embedding/transformation.py @@ -3,9 +3,11 @@ This is OpenAI compatible - no transformation is applied """ +from collections.abc import Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -16,6 +18,9 @@ from litellm.types.utils import EmbeddingResponse, Usage from ..utils import OVHCloudException +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_INT: Final = TypeAdapter(int) + class OVHCloudEmbeddingConfig(BaseEmbeddingConfig): def __init__(self) -> None: @@ -107,9 +112,10 @@ class OVHCloudEmbeddingConfig(BaseEmbeddingConfig): model_response.data = raw_response_json.get("data") model_response.object = raw_response_json.get("object") + usage_payload: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {})) usage: Final = Usage( - prompt_tokens=raw_response_json.get("usage", {}).get("prompt_tokens", 0), - total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0), + prompt_tokens=_INT.validate_python(usage_payload.get("prompt_tokens", 0) or 0), + total_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0), ) model_response.usage = usage diff --git a/litellm/llms/perplexity/chat/transformation.py b/litellm/llms/perplexity/chat/transformation.py index 5b85e1e8ae6..017a74cd51e 100644 --- a/litellm/llms/perplexity/chat/transformation.py +++ b/litellm/llms/perplexity/chat/transformation.py @@ -7,6 +7,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_logger @@ -20,6 +21,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer _PERPLEXITY_ATTRIBUTION_HEADERS: Final[Mapping[str, str]] = MappingProxyType({"X-Pplx-Integration": "litellm"}) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) class PerplexityChatConfig(OpenAIGPTConfig): @@ -110,7 +112,7 @@ class PerplexityChatConfig(OpenAIGPTConfig): # Extract and enhance usage with Perplexity-specific fields try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) self._enhance_usage_with_perplexity_fields(model_response, raw_response_json) self._add_citations_as_annotations(model_response, raw_response_json) except Exception as e: diff --git a/litellm/llms/perplexity/embedding/transformation.py b/litellm/llms/perplexity/embedding/transformation.py index c93206db2bb..8e9e69e1099 100644 --- a/litellm/llms/perplexity/embedding/transformation.py +++ b/litellm/llms/perplexity/embedding/transformation.py @@ -13,9 +13,11 @@ This module decodes them into float arrays for OpenAI-compatible responses. import base64 import struct +from collections.abc import Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -24,6 +26,9 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues from litellm.types.utils import EmbeddingResponse, Usage +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_INT: Final = TypeAdapter(int) + class PerplexityEmbeddingError(BaseLLMException): def __init__( @@ -156,15 +161,17 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig): raw_data: Final = raw_response_json.get("data", []) decoded_data: Final[list[dict[str, object]]] = [] for item in raw_data: - decoded_item = dict(item) - decoded_item["embedding"] = self._decode_base64_embedding(item.get("embedding")) + decoded_item = dict(_JSON_OBJECT.validate_python(item)) + decoded_item["embedding"] = self._decode_base64_embedding(decoded_item.get("embedding")) decoded_data.append(decoded_item) model_response.data = decoded_data - usage_data: Final = raw_response_json.get("usage", {}) + usage_data: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {})) usage: Final = Usage( - prompt_tokens=usage_data.get("prompt_tokens", 0) or usage_data.get("total_tokens", 0), - total_tokens=usage_data.get("total_tokens", 0), + prompt_tokens=_INT.validate_python( + usage_data.get("prompt_tokens", 0) or usage_data.get("total_tokens", 0) or 0 + ), + total_tokens=_INT.validate_python(usage_data.get("total_tokens", 0) or 0), ) model_response.usage = usage return model_response diff --git a/litellm/llms/perplexity/search/transformation.py b/litellm/llms/perplexity/search/transformation.py index 98299468839..5cfbdf207bb 100644 --- a/litellm/llms/perplexity/search/transformation.py +++ b/litellm/llms/perplexity/search/transformation.py @@ -2,9 +2,11 @@ Calls Perplexity's /search endpoint to search the web. """ +from collections.abc import Iterable, Mapping from typing import Final, TypedDict import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( @@ -14,6 +16,9 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.secret_managers.main import get_secret_str +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) + class _PerplexitySearchRequestRequired(TypedDict): """Required fields for Perplexity Search API request.""" @@ -146,17 +151,20 @@ class PerplexitySearchConfig(BaseSearchConfig): Returns: SearchResponse with standardized format """ - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) # Transform results to SearchResult objects results: Final = [] - for result in response_json.get("results", []): - search_result = SearchResult( - title=result.get("title", ""), - url=result.get("url", ""), - snippet=result.get("snippet", ""), - date=result.get("date"), - last_updated=result.get("last_updated"), + for raw_result in _JSON_VALUES.validate_python(response_json.get("results", [])): + result = _JSON_OBJECT.validate_python(raw_result) + search_result = SearchResult.model_validate( + { + "title": result.get("title", ""), + "url": result.get("url", ""), + "snippet": result.get("snippet", ""), + "date": result.get("date"), + "last_updated": result.get("last_updated"), + } ) results.append(search_result) diff --git a/litellm/llms/replicate/chat/handler.py b/litellm/llms/replicate/chat/handler.py index e4233ca725d..ad267847e81 100644 --- a/litellm/llms/replicate/chat/handler.py +++ b/litellm/llms/replicate/chat/handler.py @@ -1,7 +1,7 @@ import asyncio import json import time -from collections.abc import Callable +from collections.abc import Callable, Iterable from typing import Final from pydantic import ConfigDict, TypeAdapter @@ -23,6 +23,8 @@ from .transformation import ReplicateConfig replicate_config: Final = ReplicateConfig() _JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) +_POLLED_PREDICTION: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) +_TEXTS: Final = TypeAdapter(Iterable[str], config=ConfigDict(strict=True, hide_input_in_errors=True)) # Function to handle prediction response (streaming) @@ -38,11 +40,11 @@ def handle_prediction_response_streaming( print_verbose(f"replicate: polling endpoint: {prediction_url}") response = http_client.get(prediction_url, headers=headers) if response.status_code == 200: - response_data = response.json() + response_data = _POLLED_PREDICTION.validate_python(response.json()) status = response_data["status"] if "output" in response_data: try: - output_string = "".join(response_data["output"]) + output_string = "".join(_TEXTS.validate_python(response_data["output"])) except Exception: raise ReplicateError( status_code=422, diff --git a/litellm/llms/replicate/chat/transformation.py b/litellm/llms/replicate/chat/transformation.py index e8dd8816997..42e9eb2c3a6 100644 --- a/litellm/llms/replicate/chat/transformation.py +++ b/litellm/llms/replicate/chat/transformation.py @@ -1,4 +1,4 @@ -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -29,6 +29,7 @@ else: LoggingClass = Any _TEXTS: Final = TypeAdapter(Iterable[str], config=ConfigDict(strict=True, hide_input_in_errors=True)) +_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) class ReplicateConfig(BaseConfig): @@ -79,7 +80,7 @@ class ReplicateConfig(BaseConfig): seed: int | None = None, debug: bool | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = dict(locals()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @@ -250,7 +251,7 @@ class ReplicateConfig(BaseConfig): original_response=raw_response.text, additional_args={"complete_input_dict": request_data}, ) - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _JSON_DICT.validate_python(raw_response.json()) if raw_response_json.get("status") != "succeeded": raise ReplicateError( status_code=422, diff --git a/litellm/llms/sambanova/embedding/transformation.py b/litellm/llms/sambanova/embedding/transformation.py index 5483f73cf71..cf39c0a092f 100644 --- a/litellm/llms/sambanova/embedding/transformation.py +++ b/litellm/llms/sambanova/embedding/transformation.py @@ -3,9 +3,11 @@ This is OpenAI compatible - no transformation is applied """ +from collections.abc import Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -16,6 +18,9 @@ from litellm.types.utils import EmbeddingResponse, Usage from ..common_utils import SambaNovaError +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_INT: Final = TypeAdapter(int) + class SambaNovaEmbeddingConfig(BaseEmbeddingConfig): def __init__(self) -> None: @@ -124,9 +129,10 @@ class SambaNovaEmbeddingConfig(BaseEmbeddingConfig): model_response.data = raw_response_json.get("data") model_response.object = raw_response_json.get("object") + usage_payload: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {})) usage: Final = Usage( - prompt_tokens=raw_response_json.get("usage", {}).get("prompt_tokens", 0), - total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0), + prompt_tokens=_INT.validate_python(usage_payload.get("prompt_tokens", 0) or 0), + total_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0), ) model_response.usage = usage diff --git a/litellm/llms/sap/chat/transformation.py b/litellm/llms/sap/chat/transformation.py index 342a4ff7025..34e8acd766a 100755 --- a/litellm/llms/sap/chat/transformation.py +++ b/litellm/llms/sap/chat/transformation.py @@ -7,6 +7,7 @@ from functools import cached_property from typing import TYPE_CHECKING, Any, Final, Union import httpx +from pydantic import BaseModel import litellm from litellm.types.llms.openai import AllMessageValues @@ -54,7 +55,7 @@ _SAP_MODEL_PARAMS_EXCLUDED_KEYS: Final[frozenset[str]] = frozenset( def validate_dict( - data: dict, model: type + data: dict, model: type[BaseModel] ) -> dict: # mutable-ok: pydantic validation boundary; both input and output are untyped wire dicts return model(**data).model_dump(by_alias=True, exclude_unset=True) diff --git a/litellm/llms/searchapi/search/transformation.py b/litellm/llms/searchapi/search/transformation.py index ab4a68cef7f..088b807d02a 100644 --- a/litellm/llms/searchapi/search/transformation.py +++ b/litellm/llms/searchapi/search/transformation.py @@ -4,10 +4,12 @@ Calls SearchAPI.io's Google Search API endpoint. SearchAPI.io API Reference: https://www.searchapi.io/docs/google """ +from collections.abc import Iterable, Mapping from typing import Final, Literal, TypedDict, cast from urllib.parse import urlencode import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( @@ -17,6 +19,9 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.secret_managers.main import get_secret_str +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) + class _SearchAPIRequestRequired(TypedDict): """Required fields for SearchAPI.io request.""" @@ -213,24 +218,27 @@ class SearchAPIConfig(BaseSearchConfig): - organic_results[].snippet → SearchResult.snippet - organic_results[].date → SearchResult.date """ - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) # Transform results to SearchResult objects results: Final[list[SearchResult]] = [] # Process organic results - for result in response_json.get("organic_results", []): + for raw_result in _JSON_VALUES.validate_python(response_json.get("organic_results", [])): + result = _JSON_OBJECT.validate_python(raw_result) title = result.get("title", "") url = result.get("link", "") snippet = result.get("snippet", "") date = result.get("date") # SearchAPI.io provides date in some results - search_result = SearchResult( - title=title, - url=url, - snippet=snippet, - date=date, - last_updated=None, # SearchAPI.io doesn't provide last_updated + search_result = SearchResult.model_validate( + { + "title": title, + "url": url, + "snippet": snippet, + "date": date, + "last_updated": None, # SearchAPI.io doesn't provide last_updated + } ) results.append(search_result) diff --git a/litellm/llms/searxng/search/transformation.py b/litellm/llms/searxng/search/transformation.py index 4afb8b32e4b..265a3db72f5 100644 --- a/litellm/llms/searxng/search/transformation.py +++ b/litellm/llms/searxng/search/transformation.py @@ -4,9 +4,11 @@ Calls SearXNG's /search endpoint to search the web. SearXNG API Reference: https://docs.searxng.org/dev/search_api.html """ +from collections.abc import Iterable, Mapping from typing import Final, TypedDict import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( @@ -16,6 +18,9 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.secret_managers.main import get_secret_str +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) + class _SearXNGSearchRequestRequired(TypedDict): """Required fields for SearXNG Search API request.""" @@ -204,22 +209,25 @@ class SearXNGSearchConfig(BaseSearchConfig): Returns: SearchResponse with standardized format """ - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) # Transform results to SearchResult objects # Note: SearXNG doesn't natively support limiting results via API params # It returns ~20 results per page by default results: Final = [] - for result in response_json.get("results", []): + for raw_result in _JSON_VALUES.validate_python(response_json.get("results", [])): + result = _JSON_OBJECT.validate_python(raw_result) # Get date from either publishedDate or pubdate field date = result.get("publishedDate") or result.get("pubdate") - search_result = SearchResult( - title=result.get("title", ""), - url=result.get("url", ""), - snippet=result.get("content", ""), # SearXNG uses "content" for snippet - date=date, - last_updated=None, # SearXNG doesn't provide last_updated in response + search_result = SearchResult.model_validate( + { + "title": result.get("title", ""), + "url": result.get("url", ""), + "snippet": result.get("content", ""), # SearXNG uses "content" for snippet + "date": date, + "last_updated": None, # SearXNG doesn't provide last_updated in response + } ) results.append(search_result) diff --git a/litellm/llms/serper/search/transformation.py b/litellm/llms/serper/search/transformation.py index 54d5a1ff7fa..2539610ea59 100644 --- a/litellm/llms/serper/search/transformation.py +++ b/litellm/llms/serper/search/transformation.py @@ -4,9 +4,11 @@ Calls Serper's /search endpoint to search Google. Serper API Reference: https://serper.dev """ +from collections.abc import Iterable, Mapping from typing import Final, TypedDict import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( @@ -16,6 +18,9 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.secret_managers.main import get_secret_str +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) + class _SerperSearchRequestRequired(TypedDict): """Required fields for Serper Search API request.""" @@ -156,16 +161,19 @@ class SerperSearchConfig(BaseSearchConfig): Returns: SearchResponse with standardized format """ - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) results: Final = [] - for result in response_json.get("organic", []): - search_result = SearchResult( - title=result.get("title", ""), - url=result.get("link", ""), - snippet=result.get("snippet", ""), - date=result.get("date"), - last_updated=None, + for raw_result in _JSON_VALUES.validate_python(response_json.get("organic", [])): + result = _JSON_OBJECT.validate_python(raw_result) + search_result = SearchResult.model_validate( + { + "title": result.get("title", ""), + "url": result.get("link", ""), + "snippet": result.get("snippet", ""), + "date": result.get("date"), + "last_updated": None, + } ) results.append(search_result) diff --git a/litellm/llms/tavily/search/transformation.py b/litellm/llms/tavily/search/transformation.py index 158346051f6..8643322e1a8 100644 --- a/litellm/llms/tavily/search/transformation.py +++ b/litellm/llms/tavily/search/transformation.py @@ -4,9 +4,11 @@ Calls Tavily's /search endpoint to search the web. Tavily API Reference: https://docs.tavily.com/documentation/api-reference/endpoint/search """ +from collections.abc import Iterable, Mapping from typing import Final, TypedDict import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( @@ -16,6 +18,9 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.secret_managers.main import get_secret_str +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) + class _TavilySearchRequestRequired(TypedDict): """Required fields for Tavily Search API request.""" @@ -174,17 +179,20 @@ class TavilySearchConfig(BaseSearchConfig): Returns: SearchResponse with standardized format """ - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) # Transform results to SearchResult objects results: Final = [] - for result in response_json.get("results", []): - search_result = SearchResult( - title=result.get("title", ""), - url=result.get("url", ""), - snippet=result.get("content", ""), # Tavily uses "content" instead of "snippet" - date=None, # Tavily doesn't provide date in response - last_updated=None, # Tavily doesn't provide last_updated in response + for raw_result in _JSON_VALUES.validate_python(response_json.get("results", [])): + result = _JSON_OBJECT.validate_python(raw_result) + search_result = SearchResult.model_validate( + { + "title": result.get("title", ""), + "url": result.get("url", ""), + "snippet": result.get("content", ""), # Tavily uses "content" instead of "snippet" + "date": None, # Tavily doesn't provide date in response + "last_updated": None, # Tavily doesn't provide last_updated in response + } ) results.append(search_result) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 09ec2ba6ee4..7f3ef479ab5 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -672,7 +672,7 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( def _openai_batch_jsonl_entry_to_vertex_rows( openai_entry: dict[str, Any], - map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, object]], + map_openai_to_vertex_params: Callable[[dict[str, object]], dict[str, object]], ) -> tuple[Mapping[str, object], ...]: """ Transforms a single OpenAI JSONL batch entry into the Vertex rows it maps to. diff --git a/litellm/llms/vertex_ai/fine_tuning/handler.py b/litellm/llms/vertex_ai/fine_tuning/handler.py index c79b6ffce43..e1acea53192 100644 --- a/litellm/llms/vertex_ai/fine_tuning/handler.py +++ b/litellm/llms/vertex_ai/fine_tuning/handler.py @@ -5,6 +5,7 @@ from datetime import datetime from typing import Final, Literal import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger @@ -23,6 +24,8 @@ from litellm.types.llms.vertex_ai import ( ) from litellm.types.utils import LiteLLMFineTuningJob +_JSON_VALUE: Final = TypeAdapter(object) + class VertexFineTuningAPI(VertexLLM): """ @@ -341,4 +344,4 @@ class VertexFineTuningAPI(VertexLLM): f"Error creating fine tuning job. Status code: {response.status_code}. Response: {response.text}" ) - return response.json() + return _JSON_VALUE.validate_python(response.json()) diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index fddc075bfc6..f916d494aad 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -1,13 +1,14 @@ import base64 import json import os -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from io import BufferedRandom, BufferedReader, BytesIO from pathlib import Path from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -26,6 +27,11 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object]]] = TypeAdapter( + Mapping[str, object] | str | Sequence[object], config=ConfigDict(hide_input_in_errors=True) +) + class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): """ @@ -208,15 +214,18 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): headers=raw_response.headers, ) - predictions: Final = response_json.get("predictions", []) + predictions: Final = _JSON_CONTAINER.validate_python(response_json.get("predictions", [])) data_list: Final[list[ImageObject]] = [] - for prediction in predictions: + for raw_prediction in predictions: + prediction = ( + raw_prediction if isinstance(raw_prediction, str) else _JSON_CONTAINER.validate_python(raw_prediction) + ) # Imagen returns images as bytesBase64Encoded if "bytesBase64Encoded" in prediction: data_list.append( ImageObject( - b64_json=prediction["bytesBase64Encoded"], + b64_json=_JSON_OBJECT.validate_python(prediction)["bytesBase64Encoded"], url=None, ) ) diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index 9f3126fd4ff..916e25a027e 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -1,5 +1,5 @@ import os -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx @@ -35,6 +35,9 @@ else: _JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object]]] = TypeAdapter( + Mapping[str, object] | str | Sequence[object], config=ConfigDict(hide_input_in_errors=True) +) class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): @@ -308,25 +311,27 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): model_response.data = [] # Gemini image generation models return in candidates format - candidates: Final = response_data.get("candidates", []) + response_object: Final = _JSON_OBJECT.validate_python(response_data) + candidates: Final = _JSON_CONTAINER.validate_python(response_object.get("candidates", [])) for candidate in candidates: - content = candidate.get("content", {}) - parts = content.get("parts", []) - for part in parts: + content = _JSON_OBJECT.validate_python(_JSON_OBJECT.validate_python(candidate).get("content", {})) + parts = _JSON_CONTAINER.validate_python(content.get("parts", [])) + for raw_part in parts: + part = raw_part if isinstance(raw_part, str) else _JSON_CONTAINER.validate_python(raw_part) # Look for inlineData with image if "inlineData" in part: - inline_data = part["inlineData"] + part_object = _JSON_OBJECT.validate_python(part) + inline_data = _JSON_CONTAINER.validate_python(part_object["inlineData"]) if "data" in inline_data: - thought_sig = part.get("thoughtSignature") + thought_sig = part_object.get("thoughtSignature") model_response.data.append( ImageObject( - b64_json=inline_data["data"], + b64_json=_JSON_OBJECT.validate_python(inline_data)["data"], url=None, provider_specific_fields=({"thought_signature": thought_sig} if thought_sig else None), ) ) - response_object: Final = _JSON_OBJECT.validate_python(response_data) if usage_metadata := response_object.get("usageMetadata", None): model_response.usage = self._transform_image_usage(_JSON_DICT.validate_python(usage_metadata)) diff --git a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py index ae6e08611ae..927c7b74e67 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py @@ -1,7 +1,9 @@ import os +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.llms.base_llm.image_generation.transformation import ( @@ -27,6 +29,12 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_JSON_VALUE: Final = TypeAdapter(object) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object]]] = TypeAdapter( + Mapping[str, object] | str | Sequence[object], config=ConfigDict(hide_input_in_errors=True) +) + class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): """ @@ -221,7 +229,7 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): Transform Imagen image generation response to litellm ImageResponse format """ try: - response_data: Final = raw_response.json() + response_data: Final = _JSON_VALUE.validate_python(raw_response.json()) except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", @@ -233,13 +241,18 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): model_response.data = [] # Imagen format - predictions with generated images - predictions: Final = response_data.get("predictions", []) - for prediction in predictions: + predictions: Final = _JSON_CONTAINER.validate_python( + _JSON_OBJECT.validate_python(response_data).get("predictions", []) + ) + for raw_prediction in predictions: + prediction = ( + raw_prediction if isinstance(raw_prediction, str) else _JSON_CONTAINER.validate_python(raw_prediction) + ) # Imagen returns images as bytesBase64Encoded if "bytesBase64Encoded" in prediction: model_response.data.append( ImageObject( - b64_json=prediction["bytesBase64Encoded"], + b64_json=_JSON_OBJECT.validate_python(prediction)["bytesBase64Encoded"], url=None, ) ) diff --git a/litellm/llms/vertex_ai/text_to_speech/text_to_speech_handler.py b/litellm/llms/vertex_ai/text_to_speech/text_to_speech_handler.py index d053f564dee..655ae215f76 100644 --- a/litellm/llms/vertex_ai/text_to_speech/text_to_speech_handler.py +++ b/litellm/llms/vertex_ai/text_to_speech/text_to_speech_handler.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from typing_extensions import TypedDict import litellm @@ -12,6 +14,9 @@ from litellm.llms.openai.openai import HttpxBinaryResponseContent from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_STR: Final = TypeAdapter(str, config=ConfigDict(hide_input_in_errors=True)) + class VertexInput(TypedDict, total=False): text: str | None @@ -150,12 +155,12 @@ class VertexTextToSpeechAPI(VertexLLM): if response.status_code != 200: raise Exception(f"Request failed with status code {response.status_code}, {response.text}") ############ Process the response ############ - _json_response: Final = response.json() + _json_response: Final = _JSON_OBJECT.validate_python(response.json()) response_content: Final = _json_response["audioContent"] # Decode base64 to get binary content - binary_data: Final = base64.b64decode(response_content) + binary_data: Final = base64.b64decode(_STR.validate_python(response_content)) # Create an httpx.Response object response = httpx.Response( @@ -187,12 +192,12 @@ class VertexTextToSpeechAPI(VertexLLM): if response.status_code != 200: raise Exception(f"Request did not return a 200 status code: {response.status_code}, {response.text}") - _json_response: Final = response.json() + _json_response: Final = _JSON_OBJECT.validate_python(response.json()) response_content: Final = _json_response["audioContent"] # Decode base64 to get binary content - binary_data: Final = base64.b64decode(response_content) + binary_data: Final = base64.b64decode(_STR.validate_python(response_content)) # Create an httpx.Response object response = httpx.Response( diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index a78145fdcf4..2fae2a9de1c 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -47,6 +47,7 @@ else: _LyriaVoice: TypeAlias = str | dict | None _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_ARRAY: Final = TypeAdapter(tuple[object, ...], config=ConfigDict(hide_input_in_errors=True)) _STR: Final = TypeAdapter(str, config=ConfigDict(hide_input_in_errors=True)) @@ -650,27 +651,37 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): ) -> "HttpxBinaryResponseContent": from litellm.types.llms.openai import HttpxBinaryResponseContent - response_json: Final = raw_response.json() + response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) base_model: Final = model.removeprefix("vertex_ai/") model_info: Final = self._get_model_info(model=model) - audio_data: str | None = None # rebind-ok: response parsing discovers audio data in provider-specific shapes - mime_type: str | None = None # rebind-ok: response parsing discovers the MIME type beside the audio payload + audio_data: object = None # rebind-ok: response parsing discovers audio data in provider-specific shapes + mime_type: object = None # rebind-ok: response parsing discovers the MIME type beside the audio payload if model_info["vertex_ai_audio_api"] == "lyria_predict": - predictions: Final = response_json.get("predictions") or () + predictions: Final = _JSON_ARRAY.validate_python(response_json.get("predictions") or ()) if predictions: - audio_data = predictions[0].get("audioContent") or predictions[0].get("bytesBase64Encoded") - mime_type = predictions[0].get("mimeType") # rebind-ok: predict response supplies its audio MIME type + prediction: Final = _JSON_OBJECT.validate_python(predictions[0]) + audio_data = prediction.get("audioContent") or prediction.get("bytesBase64Encoded") + mime_type = prediction.get("mimeType") # rebind-ok: predict response supplies its audio MIME type else: - for step in response_json.get("steps") or response_json.get("outputs") or (): - content_items = step.get("content") or () if step.get("type") == "model_output" else (step,) + steps: Final = _JSON_ARRAY.validate_python(response_json.get("steps") or response_json.get("outputs") or ()) + for step in map(_JSON_OBJECT.validate_python, steps): + content_items = ( + map(_JSON_OBJECT.validate_python, _JSON_ARRAY.validate_python(step.get("content") or ())) + if step.get("type") == "model_output" + else (step,) + ) for content in content_items: if content.get("type") == "audio" and content.get("data"): audio_data = content["data"] mime_type = content.get("mime_type") if audio_data is None: raise ValueError(f"No generated audio found in Vertex AI {base_model} response") - binary_data: Final = base64.b64decode(audio_data) - media_type: Final = mime_type or speech_media_type_from_audio_bytes(binary_data) or DEFAULT_SPEECH_MEDIA_TYPE + binary_data: Final = base64.b64decode(_STR.validate_python(audio_data)) + media_type: Final = ( + (_STR.validate_python(mime_type) if mime_type else None) + or speech_media_type_from_audio_bytes(binary_data) + or DEFAULT_SPEECH_MEDIA_TYPE + ) return HttpxBinaryResponseContent( httpx.Response( status_code=raw_response.status_code, diff --git a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py index 9626c879dd0..f1afaa7be59 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py @@ -1,6 +1,7 @@ from typing import Final, Literal import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject @@ -17,6 +18,9 @@ from litellm.types.utils import EmbeddingResponse from .types import * +_JSON_VALUE: Final = TypeAdapter(object) +_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) + class VertexEmbedding(VertexBase): def __init__(self) -> None: @@ -123,12 +127,12 @@ class VertexEmbedding(VertexBase): except httpx.TimeoutException: raise VertexAIError(status_code=408, message="Timeout error occurred.") - _json_response: Final = response.json() + _json_response: Final = _JSON_VALUE.validate_python(response.json()) ## LOGGING POST-CALL logging_obj.post_call(input=input, api_key=None, original_response=_json_response) model_response = litellm.vertexAITextEmbeddingConfig.transform_vertex_response_to_openai( - response=_json_response, model=model, model_response=model_response + response=_JSON_DICT.validate_python(_json_response), model=model, model_response=model_response ) return model_response @@ -215,12 +219,12 @@ class VertexEmbedding(VertexBase): except httpx.TimeoutException: raise VertexAIError(status_code=408, message="Timeout error occurred.") - _json_response: Final = response.json() + _json_response: Final = _JSON_VALUE.validate_python(response.json()) ## LOGGING POST-CALL logging_obj.post_call(input=input, api_key=None, original_response=_json_response) model_response = litellm.vertexAITextEmbeddingConfig.transform_vertex_response_to_openai( - response=_json_response, model=model, model_response=model_response + response=_JSON_DICT.validate_python(_json_response), model=model, model_response=model_response ) return model_response diff --git a/litellm/llms/voyage/embedding/transformation.py b/litellm/llms/voyage/embedding/transformation.py index a9232ad27e2..53220d59f51 100644 --- a/litellm/llms/voyage/embedding/transformation.py +++ b/litellm/llms/voyage/embedding/transformation.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -9,6 +11,9 @@ from litellm.llms.voyage.common_utils import get_default_base_url, get_voyage_ap from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues from litellm.types.utils import EmbeddingResponse, Usage +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_INT: Final = TypeAdapter(int) + class VoyageError(BaseLLMException): def __init__( @@ -123,9 +128,10 @@ class VoyageEmbeddingConfig(BaseEmbeddingConfig): model_response.data = raw_response_json.get("data") model_response.object = raw_response_json.get("object") + usage_payload: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {})) usage: Final = Usage( - prompt_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0), - total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0), + prompt_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0), + total_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0), ) model_response.usage = usage return model_response diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index f5e7fdb48f5..231efb596d8 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -7,6 +7,7 @@ from collections.abc import Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -15,6 +16,9 @@ from litellm.llms.voyage.common_utils import get_default_base_url, get_voyage_ap from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues from litellm.types.utils import EmbeddingResponse, Usage +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_INT: Final = TypeAdapter(int) + class VoyageError(BaseLLMException): def __init__( @@ -149,9 +153,10 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): model_response.data = raw_response_json.get("data") model_response.object = raw_response_json.get("object") + usage_payload: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {})) usage: Final = Usage( - prompt_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0), - total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0), + prompt_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0), + total_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0), ) model_response.usage = usage return model_response diff --git a/litellm/llms/voyage/embedding/transformation_multimodal.py b/litellm/llms/voyage/embedding/transformation_multimodal.py index 035765b6691..550c3e87eeb 100644 --- a/litellm/llms/voyage/embedding/transformation_multimodal.py +++ b/litellm/llms/voyage/embedding/transformation_multimodal.py @@ -6,9 +6,11 @@ containing content blocks, unlike standard Voyage embeddings which use /v1/embeddings and a string/list `input` field. """ +from collections.abc import Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -17,6 +19,9 @@ from litellm.llms.voyage.common_utils import get_default_base_url, get_voyage_ap from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues from litellm.types.utils import EmbeddingResponse, Usage +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_INT: Final = TypeAdapter(int) + class VoyageMultimodalEmbeddingError(BaseLLMException): def __init__( @@ -155,8 +160,8 @@ class VoyageMultimodalEmbeddingConfig(BaseEmbeddingConfig): model_response.data = raw_response_json.get("data") model_response.object = raw_response_json.get("object") - usage_payload: Final = raw_response_json.get("usage", {}) - total_tokens: Final = usage_payload.get("total_tokens", 0) + usage_payload: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {})) + total_tokens: Final = _INT.validate_python(usage_payload.get("total_tokens", 0) or 0) model_response.usage = Usage( prompt_tokens=total_tokens, total_tokens=total_tokens, diff --git a/litellm/llms/watsonx/completion/transformation.py b/litellm/llms/watsonx/completion/transformation.py index e01a0ff85e5..5c0416bd6ca 100644 --- a/litellm/llms/watsonx/completion/transformation.py +++ b/litellm/llms/watsonx/completion/transformation.py @@ -1,5 +1,5 @@ import time -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Iterator, Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, Final @@ -102,7 +102,7 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): stream: bool | None = None, **kwargs, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = dict(locals()) for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) diff --git a/litellm/llms/xai/common_utils.py b/litellm/llms/xai/common_utils.py index cf76a851a86..4843caeac59 100644 --- a/litellm/llms/xai/common_utils.py +++ b/litellm/llms/xai/common_utils.py @@ -1,6 +1,8 @@ +from collections.abc import Iterable, Mapping from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter import litellm from litellm.llms.base_llm.base_utils import BaseLLMModelInfo @@ -9,6 +11,9 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ProviderSpecificModelInfo USD_TICKS_PER_DOLLAR: Final = 10_000_000_000 +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) +_MODEL_ID: Final = TypeAdapter(str, config=ConfigDict(strict=True, hide_input_in_errors=True)) def xai_reported_cost_in_usd(cost_in_usd_ticks: object) -> float | None: @@ -97,11 +102,11 @@ class XAIModelInfo(BaseLLMModelInfo): f"Failed to fetch models from XAI. Status code: {response.status_code}, Response: {response.text}" ) - models: Final = response.json()["data"] + models: Final = _JSON_VALUES.validate_python(_JSON_OBJECT.validate_python(response.json())["data"]) litellm_model_names: Final = [] for model in models: - stripped_model_name = model["id"] + stripped_model_name = _MODEL_ID.validate_python(_JSON_OBJECT.validate_python(model)["id"]) litellm_model_name = "xai/" + stripped_model_name litellm_model_names.append(litellm_model_name) return litellm_model_names diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py index 74f85d5fa04..d1a8084f4b2 100644 --- a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py +++ b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py @@ -142,7 +142,7 @@ class MCPGuardrailTranslationHandler(BaseTranslation): data: dict[str, Any], guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: mcp_tool_name: Final = data.get("mcp_tool_name") or data.get("name") mcp_arguments: Final[object] = data.get("mcp_arguments") or data.get("arguments") mcp_tool_description: Final = data.get("mcp_tool_description") or data.get("description") @@ -228,9 +228,9 @@ class MCPGuardrailTranslationHandler(BaseTranslation): response: "CallToolResult", guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, - user_api_key_dict: Any | None = None, + user_api_key_dict: object | None = None, request_data: dict | None = None, - ) -> Any: + ) -> object: """Scan the text content of an MCP tool result and write masked text back. The content list is rewritten in place (only the entries the guardrail diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 6beea5bfd71..a7375b4f17c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -6023,7 +6023,7 @@ class MCPServerManager: server=server, ) - hook_result: Final[dict[str, Any]] = {} + hook_result: Final[dict[str, object]] = {} if proxy_logging_obj is None: return hook_result diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 8ef573aa489..d0133739787 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -157,7 +157,7 @@ def _sanitize_path_parameter_value(param_value: object, param_name: str) -> str: return quote(value_str, safe="") -def load_openapi_spec(filepath: str) -> dict[str, Any]: +def load_openapi_spec(filepath: str) -> dict[str, object]: """ Sync wrapper. For URL specs, use the shared/custom MCP httpx client. """ diff --git a/litellm/proxy/_experimental/mcp_server/tool_registry.py b/litellm/proxy/_experimental/mcp_server/tool_registry.py index dd25f09cc57..1f1d952f96c 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_registry.py +++ b/litellm/proxy/_experimental/mcp_server/tool_registry.py @@ -3,7 +3,7 @@ import json from collections.abc import Callable, Generator, Mapping from contextlib import contextmanager from contextvars import ContextVar -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger from litellm.proxy.types_utils.utils import get_instance_fn @@ -56,7 +56,7 @@ class MCPToolRegistry: self, name: str, description: str, - input_schema: dict[str, Any], + input_schema: dict[str, object], handler: Callable, *, server_id: str | None = None, @@ -117,7 +117,7 @@ class MCPToolRegistry: def load_tools_from_config( self, - mcp_tools_config: dict[str, Any] | None = None, + mcp_tools_config: dict[str, object] | None = None, config_file_path: str | None = None, ) -> None: """ diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 8c9fdfcf67d..304f282d1f8 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -303,7 +303,7 @@ async def _a2a_sse_event_source( error_body: Final = await resp.aread() error_event: Mapping[str, object] | None = None try: - parsed: Final = json.loads(error_body) + parsed: Final = _DECODED_JSON.validate_python(json.loads(error_body)) if isinstance(parsed, dict) and "error" in parsed: error_event = normalize_a2a_jsonrpc_response(parsed, request_id=request_id) except Exception: diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index f2c5d066714..d4727cc8d8d 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -86,7 +86,7 @@ def _anthropic_error_json_response(exc: ProxyException, request: Request) -> JSO return JSONResponse(status_code=status_code, content=content, headers=exc.headers) -def _strip_total_tokens_from_anthropic_response(response: Any) -> None: +def _strip_total_tokens_from_anthropic_response(response: object) -> None: """Remove the OpenAI-flavored `usage.total_tokens` field that LiteLLM injects into Anthropic /v1/messages responses. diff --git a/litellm/proxy/auth/ip_address_utils.py b/litellm/proxy/auth/ip_address_utils.py index 5d32cc7eb4c..f1ca2c70f3b 100644 --- a/litellm/proxy/auth/ip_address_utils.py +++ b/litellm/proxy/auth/ip_address_utils.py @@ -199,7 +199,7 @@ class IPAddressUtils: @staticmethod def is_request_https( request: Request, - general_settings: Mapping[str, Any] | None = None, + general_settings: Mapping[str, object] | None = None, ) -> bool: """ Whether this request's PUBLIC-facing origin is HTTPS, for deciding diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 3101f19ced2..b70cdb08b96 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -2,7 +2,7 @@ ## Common checks for /v1/models and `/model/info` import copy from collections.abc import Sequence -from typing import Any, Final +from typing import Final import litellm from litellm._logging import verbose_proxy_logger @@ -336,8 +336,8 @@ def get_known_models_from_wildcard(wildcard_model: str, litellm_params: LiteLLM_ def expand_wildcard_deployments_for_model_info( - deployments: list[dict[str, Any]], -) -> list[dict[str, Any]]: + deployments: list[dict[str, object]], +) -> list[dict[str, object]]: """Expand wildcard deployments into one row per known provider model. PR #30025 changed /model/info to read from llm_router.model_list (correct, diff --git a/litellm/proxy/auth/trusted_proxy_utils.py b/litellm/proxy/auth/trusted_proxy_utils.py index 55cf3e98444..04c09abd5dd 100644 --- a/litellm/proxy/auth/trusted_proxy_utils.py +++ b/litellm/proxy/auth/trusted_proxy_utils.py @@ -1,4 +1,4 @@ -from typing import Any, Final +from typing import Final from fastapi import Request @@ -22,7 +22,7 @@ def _get_proxy_general_settings() -> dict[str, object]: def get_trusted_proxy_cidrs( - general_settings: dict[str, Any] | None = None, + general_settings: dict[str, object] | None = None, ) -> list[str]: """Operator-configured trusted reverse-proxy CIDRs, normalized to strings. @@ -48,7 +48,7 @@ def _get_direct_client_ip(request: Request) -> str | None: def require_trusted_proxy_request( *, request: Request, - general_settings: dict[str, Any] | None = None, + general_settings: dict[str, object] | None = None, feature_name: str, setting_name: str = TRUSTED_PROXY_RANGES_KEY, ) -> None: diff --git a/litellm/proxy/caching_routes.py b/litellm/proxy/caching_routes.py index c621ec6a33a..597080be3b3 100644 --- a/litellm/proxy/caching_routes.py +++ b/litellm/proxy/caching_routes.py @@ -40,7 +40,9 @@ def _extract_cache_params() -> dict[str, object]: return {} try: cache_params: Final = vars(litellm.cache.cache) - cleaned_params: Final = HealthCheckCacheParams(**cache_params).model_dump() if cache_params else {} + cleaned_params: Final = ( + HealthCheckCacheParams.model_validate(dict(**cache_params)).model_dump() if cache_params else {} + ) return masker.mask_dict(cleaned_params) except (AttributeError, TypeError) as e: verbose_proxy_logger.debug("Error extracting cache params: %s", e) diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 684449acd0c..9e2e6966ed7 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -3,7 +3,7 @@ import sys import time import webbrowser from collections.abc import Callable, Mapping, Sequence -from typing import Any, Final, TypeVar +from typing import Any, Final, Protocol, TypeVar from urllib.parse import urlencode import click @@ -96,6 +96,24 @@ class CliContextObj(TypedDict): api_key_from_token_file: ReadOnly[NotRequired[bool]] +class _CliContextObjView(Protocol): + @property + def obj(self) -> CliContextObj: ... + + +class _OptionalCliContextObjView(Protocol): + @property + def obj(self) -> CliContextObj | None: ... + + +def _cli_context_obj(view: _CliContextObjView) -> CliContextObj: + return view.obj + + +def _optional_cli_context_obj(view: _OptionalCliContextObjView) -> CliContextObj | None: + return view.obj + + class CliPollData(TypedDict, total=False): status: str key: str @@ -197,7 +215,7 @@ def keychain_unreadable_notice(vault: SecretVault) -> str: def context_secret_vault(ctx: click.Context) -> SecretVault: """Where this invocation reads and writes secret material; injectable through ctx.obj for tests""" - ctx_obj: Final[CliContextObj | None] = ctx.obj + ctx_obj: Final = _optional_cli_context_obj(ctx) if ctx_obj is None: return SYSTEM_KEYRING return ctx_obj.get("secret_vault") or SYSTEM_KEYRING @@ -880,7 +898,7 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None: """Login to LiteLLM proxy using SSO authentication""" from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER - ctx_obj: Final[CliContextObj] = ctx.obj + ctx_obj: Final = _cli_context_obj(ctx) base_url: Final = ctx_obj["base_url"] if config_claude: settings_path: Final = claude_settings_path(os.environ) @@ -1030,7 +1048,7 @@ def print_token(ctx: click.Context): # explicitly pointed us at a server, trust whichever one `lite login` # actually issued this token for -- that's the whole point of not # needing a wrapper command. - ctx_obj: Final[CliContextObj] = ctx.obj + ctx_obj: Final = _cli_context_obj(ctx) issued_for_this_server: Final = token_data.get("base_url") == ctx_obj.get("base_url", "").rstrip("/") if ctx_obj.get("base_url_explicit") and not issued_for_this_server: click.echo("Not authenticated for this server. Run 'lite login'.", err=True) diff --git a/litellm/proxy/client/cli/commands/http.py b/litellm/proxy/client/cli/commands/http.py index 722c3afac23..529564c685c 100644 --- a/litellm/proxy/client/cli/commands/http.py +++ b/litellm/proxy/client/cli/commands/http.py @@ -1,13 +1,32 @@ import json as json_lib -from typing import Final +from typing import Final, Protocol import click import requests import rich +from typing_extensions import ReadOnly, TypedDict from ...http_client import HTTPClient +class _CliContext(TypedDict): + base_url: ReadOnly[str] + api_key: ReadOnly[str | None] + + +class _CliContextView(Protocol): + @property + def obj(self) -> _CliContext: ... + + +class _JsonBodyView(TypedDict): + body: ReadOnly[object] + + +def _cli_context(view: _CliContextView) -> _CliContext: + return view.obj + + @click.group() def http(): """Make HTTP requests to the LiteLLM proxy server""" @@ -79,21 +98,23 @@ def request( # If not JSON, use as raw data request_data = data - client: Final = HTTPClient(ctx.obj["base_url"], ctx.obj["api_key"]) + client: Final = HTTPClient(_cli_context(ctx)["base_url"], _cli_context(ctx)["api_key"]) try: - response: Final = client.request( - method=method, - uri=uri, - data=request_data, - json=json_data, - headers=headers, - ) - rich.print_json(data=response) + response: Final[_JsonBodyView] = { + "body": client.request( + method=method, + uri=uri, + data=request_data, + json=json_data, + headers=headers, + ) + } + rich.print_json(data=response["body"]) except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) try: - error_body: Final = e.response.json() - rich.print_json(data=error_body) + error_body: Final[_JsonBodyView] = {"body": e.response.json()} + rich.print_json(data=error_body["body"]) except json_lib.JSONDecodeError: click.echo(e.response.text, err=True) raise click.Abort() diff --git a/litellm/proxy/client/cli/commands/keys.py b/litellm/proxy/client/cli/commands/keys.py index 0ab3d2480d9..558c9139f5b 100644 --- a/litellm/proxy/client/cli/commands/keys.py +++ b/litellm/proxy/client/cli/commands/keys.py @@ -2,7 +2,7 @@ import builtins import json from collections.abc import Mapping, Sequence from datetime import datetime -from typing import Any, Final, Literal +from typing import Any, Final, Literal, Protocol import click import requests @@ -20,8 +20,9 @@ class _CliContext(TypedDict): api_key: ReadOnly[str | None] -class _CliContextView(TypedDict): - obj: ReadOnly[_CliContext] +class _CliContextView(Protocol): + @property + def obj(self) -> _CliContext: ... class _KeyRowsView(TypedDict): @@ -32,6 +33,10 @@ class _JsonBodyView(TypedDict): body: ReadOnly[object] +def _cli_context(view: _CliContextView) -> _CliContext: + return view.obj + + @click.group() def keys(): """Manage API keys for the LiteLLM proxy server""" @@ -74,8 +79,8 @@ def list( return_full_object: bool, ): """List all API keys""" - context: Final[_CliContextView] = {"obj": ctx.obj} - client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"]) + context: Final = _cli_context(ctx) + client: Final = KeysManagementClient(context["base_url"], context["api_key"]) response: Final = client.list( page=page, size=size, @@ -140,8 +145,8 @@ def generate( config: str | None, ): """Generate a new API key""" - context: Final[_CliContextView] = {"obj": ctx.obj} - client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"]) + context: Final = _cli_context(ctx) + client: Final = KeysManagementClient(context["base_url"], context["api_key"]) try: models_list: Final = [m.strip() for m in models.split(",")] if models else None aliases_dict: Final = json.loads(aliases) if aliases else None @@ -177,8 +182,8 @@ def generate( @click.pass_context def delete(ctx: click.Context, keys: str | None, key_aliases: str | None): """Delete API keys by key or alias""" - context: Final[_CliContextView] = {"obj": ctx.obj} - client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"]) + context: Final = _cli_context(ctx) + client: Final = KeysManagementClient(context["base_url"], context["api_key"]) keys_list: Final = [k.strip() for k in keys.split(",")] if keys else None aliases_list: Final = [a.strip() for a in key_aliases.split(",")] if key_aliases else None try: @@ -377,8 +382,8 @@ def import_keys( # Create clients for both source and destination source_client: Final = KeysManagementClient(source_base_url, source_api_key) - context: Final[_CliContextView] = {"obj": ctx.obj} - dest_client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"]) + context: Final = _cli_context(ctx) + dest_client: Final = KeysManagementClient(context["base_url"], context["api_key"]) try: # Get all keys from source instance with pagination diff --git a/litellm/proxy/client/cli/commands/models.py b/litellm/proxy/client/cli/commands/models.py index f2b38c6eab4..ac7d5ee45ea 100644 --- a/litellm/proxy/client/cli/commands/models.py +++ b/litellm/proxy/client/cli/commands/models.py @@ -378,7 +378,7 @@ def _print_summary_table(provider_counts): rich.print(summary_table) -def get_model_list_from_yaml_file(yaml_file: str) -> list[dict[str, Any]]: +def get_model_list_from_yaml_file(yaml_file: str) -> list[dict[str, object]]: """Load and validate the model list from a YAML file.""" with open(yaml_file, "r") as f: data: Final = yaml.safe_load(f) diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index 682dd61ac42..cbbf4247624 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -1,7 +1,7 @@ # stdlib imports # third party imports -from typing import Final +from typing import Final, Protocol import click @@ -37,6 +37,15 @@ from .commands.users import users from .interface import interactive_shell +class _CliContextObjView(Protocol): + @property + def obj(self) -> CliContextObj: ... + + +def _cli_context_obj(view: _CliContextObjView) -> CliContextObj: + return view.obj + + def print_version(base_url: str, api_key: str | None): """Print CLI and server version info.""" click.echo(f"LiteLLM Proxy CLI Version: {litellm_version}") @@ -135,7 +144,7 @@ def cli(ctx: click.Context, show_version: bool, base_url: str | None, api_key: s @click.pass_context def version(ctx: click.Context): """Show the LiteLLM Proxy CLI and server version.""" - ctx_obj: Final[CliContextObj] = ctx.obj + ctx_obj: Final = _cli_context_obj(ctx) print_version(ctx_obj.get("base_url"), ctx_obj.get("api_key")) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6b87960fcc0..00729c12cc1 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -27,7 +27,7 @@ import httpx import orjson from fastapi import HTTPException, Request, status from fastapi.responses import JSONResponse, Response, StreamingResponse -from pydantic import BaseModel, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from starlette.types import Receive, Scope, Send from typing_extensions import Never @@ -556,6 +556,7 @@ def _stream_usage_tracking_updates( } +@runtime_checkable class _UpstreamHttpResponse(Protocol): @property def status_code(self) -> int: ... @@ -570,6 +571,11 @@ def _as_upstream_response(response: _UpstreamHttpResponse) -> _UpstreamHttpRespo return response +_UPSTREAM_RESPONSE: Final = TypeAdapter( + _UpstreamHttpResponse, config=ConfigDict(arbitrary_types_allowed=True, hide_input_in_errors=True) +) + + class _ReadsHeaderValues(Protocol): def get(self, key: str, default: str = "") -> str: ... @@ -2562,7 +2568,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def _response_cost_from_logging_obj( *, - response: Any, + response: object, logging_obj: LiteLLMLoggingObj, ) -> float | str: """ @@ -2618,7 +2624,7 @@ class ProxyBaseLLMRequestProcessing: is_streaming_request: bool | None = False, contents: list[object] | None = None, skip_pre_call_logic: bool = False, - ) -> Any: + ) -> object: """Run the request, sending SSE keepalives while the upstream is still silent. Everything below this point, the upstream call included, happens before the @@ -3261,7 +3267,7 @@ class ProxyBaseLLMRequestProcessing: if isinstance(result, Response): return result - upstream: Final = _as_upstream_response(result) + upstream: Final = _UPSTREAM_RESPONSE.validate_python(result) content: Final = await upstream.aread() return Response( content=content, diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 80d80224617..81ac9eba665 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -545,7 +545,7 @@ async def convert_upload_files_to_file_data( return data -async def get_request_body(request: Request) -> dict[str, Any]: +async def get_request_body(request: Request) -> dict[str, object]: """ Read the request body and parse it as JSON. """ diff --git a/litellm/proxy/common_utils/openapi_schema_compat.py b/litellm/proxy/common_utils/openapi_schema_compat.py index 903e6c1cb24..1826daa2762 100644 --- a/litellm/proxy/common_utils/openapi_schema_compat.py +++ b/litellm/proxy/common_utils/openapi_schema_compat.py @@ -5,7 +5,7 @@ FastAPI 0.120+ has stricter schema generation that fails on certain types like o This module provides a compatibility layer to handle these cases gracefully. """ -from typing import Any, Final +from typing import Final from litellm._logging import verbose_proxy_logger @@ -16,7 +16,7 @@ def get_openapi_schema_with_compat( version: str, description: str, routes: list, -) -> dict[str, Any]: +) -> dict[str, object]: """ Generate OpenAPI schema with compatibility handling for FastAPI 0.120+. diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 2a2ef5217b8..64c935ce0d0 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -146,7 +146,7 @@ def iter_message_text(data: Mapping[str, object]) -> Iterator[str]: yield from _iter_text_parts_in_content(message.get("content")) -def walk_user_text(data: dict[str, Any], visit: Callable[[str], str]) -> int: +def walk_user_text(data: dict[str, object], visit: Callable[[str], str]) -> int: """Rewrite every text fragment in place via ``visit``. Mutates ``data["messages"]`` and ``data["input"]``. Returns the number @@ -286,7 +286,7 @@ def has_non_string_content(data: Mapping[str, object]) -> bool: return False -def build_inspection_messages(data: dict[str, Any]) -> list[dict[str, str]]: +def build_inspection_messages(data: dict[str, object]) -> list[dict[str, str]]: """Synthesize a chat-style messages list for posting to a guardrail API. Each returned message has a plain-string ``content`` — multimodal text diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py index 3a8bd54c587..50a14375107 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py @@ -8,7 +8,7 @@ import os from collections.abc import AsyncGenerator, AsyncIterable from datetime import datetime -from typing import Any, Final +from typing import Final import httpx @@ -86,7 +86,7 @@ class DynamoAIGuardrails(CustomGuardrail): async def _call_dynamoai_guardrails( self, - messages: list[dict[str, Any]], + messages: list[dict[str, object]], event_type: GuardrailEventHooks, text_type: str = "input", request_data: dict | None = None, @@ -405,7 +405,7 @@ class DynamoAIGuardrails(CustomGuardrail): # to avoid sending empty content to DynamoAI (e.g., during tool calls) if isinstance(response, litellm.ModelResponse): has_text_content = False - dynamoai_messages: Final[list[dict[str, str]]] = [] + dynamoai_messages: Final[list[dict[str, object]]] = [] for choice in response.choices: if isinstance(choice, litellm.Choices): diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index c745a6c91b3..25df7a20e14 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -8,7 +8,7 @@ import fnmatch import os from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional import httpx @@ -194,7 +194,7 @@ class GenericGuardrailAPI(CustomGuardrail): def __init__( self, - headers: dict[str, Any] | None = None, + headers: dict[str, str] | None = None, api_base: str | None = None, api_key: str | None = None, additional_provider_specific_params: Mapping[str, object] | None = None, diff --git a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py index 5da64de329b..380562aacac 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py @@ -6,7 +6,7 @@ # +-------------------------------------------------------------+ import os -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Mapping from datetime import datetime from typing import Any, Final @@ -40,7 +40,7 @@ class IBMGuardrailDetector(CustomGuardrail): base_url: str | None = None, detector_id: str | None = None, is_detector_server: bool = True, - detector_params: dict[str, Any] | None = None, + detector_params: Mapping[str, object] | None = None, extra_headers: dict[str, str] | None = None, score_threshold: float | None = None, block_on_detection: bool = True, diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py index 6ea4317a174..c970ac263e7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py @@ -371,8 +371,10 @@ class LLMShieldProxyGuardrail(CustomGuardrail): return response pending: Final[SlotSink] = [] + choice: object + tool_call: object for choice in choices: - message = getattr(choice, "message", None) + message: object = getattr(choice, "message", None) if message is None: text = read_field(choice, "text") if isinstance(text, str) and text: @@ -382,11 +384,11 @@ class LLMShieldProxyGuardrail(CustomGuardrail): if isinstance(content, str) and content: pending.append((content, functools.partial(setattr, message, "content"))) for tool_call in getattr(message, "tool_calls", None) or (): - function = getattr(tool_call, "function", None) + function: object = getattr(tool_call, "function", None) arguments = getattr(function, "arguments", None) if function is not None else None if isinstance(arguments, str) and arguments: pending.append((arguments, functools.partial(setattr, function, "arguments"))) - legacy = getattr(message, "function_call", None) + legacy: object = getattr(message, "function_call", None) legacy_arguments = getattr(legacy, "arguments", None) if legacy is not None else None if isinstance(legacy_arguments, str) and legacy_arguments: pending.append((legacy_arguments, functools.partial(setattr, legacy, "arguments"))) @@ -448,6 +450,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): function_call_output holds `output` -- so the two directions stay symmetric. """ slots: Final[SlotSink] = [] + item: object for item in getattr(response, "output", None) or (): collect_response_item(item, slots) return tuple(slots) @@ -501,6 +504,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): continue restored_chunk = detached(chunk) last_chunk = restored_chunk + choice: object for choice in getattr(restored_chunk, "choices", None) or (): await self._restore_choice(choice, carries, session_id) yield restored_chunk @@ -532,6 +536,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): await self._restore_content_window(delta, (index, None), carries, session_id, is_final) + tool_call: object for tool_call in getattr(delta, "tool_calls", None) or (): await self._restore_tool_call_window(tool_call, index, carries, session_id) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index a60059fcc4d..0c1be621923 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1375,6 +1375,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # 2. Process tool calls tool_calls = getattr(message, "tool_calls", None) if tool_calls: + tool_call: object for tool_call in tool_calls: function = getattr(tool_call, "function", None) if function and hasattr(function, "arguments"): diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py index b2139779925..1a571047977 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py +++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py @@ -7,7 +7,7 @@ then builds a SemanticRouter for prompt matching. import os from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final import yaml @@ -28,7 +28,7 @@ class SemanticGuardRouteLoader: """Loads route definitions from YAML templates and custom configs, builds SemanticRouter.""" @staticmethod - def load_builtin_template(template_name: str) -> dict[str, Any]: + def load_builtin_template(template_name: str) -> dict[str, object]: """Load a built-in route template YAML by name.""" file_path: Final = os.path.join(ROUTE_TEMPLATES_DIR, f"{template_name}.yaml") if not os.path.exists(file_path): @@ -50,7 +50,7 @@ class SemanticGuardRouteLoader: return sorted(templates) @staticmethod - def load_custom_routes_file(file_path: str) -> list[dict[str, Any]]: + def load_custom_routes_file(file_path: str) -> list[dict[str, object]]: """Load custom routes from a YAML file.""" if not os.path.exists(file_path): raise ValueError(f"SemanticGuard: custom routes file not found: {file_path}") diff --git a/litellm/proxy/health_check_utils/shared_health_check_manager.py b/litellm/proxy/health_check_utils/shared_health_check_manager.py index 764d570fbea..62c973196b9 100644 --- a/litellm/proxy/health_check_utils/shared_health_check_manager.py +++ b/litellm/proxy/health_check_utils/shared_health_check_manager.py @@ -192,12 +192,12 @@ class SharedHealthCheckManager: @with_service_target(HEALTH_CHECKS_TARGET) async def perform_shared_health_check( self, - model_list: list[dict[str, Any]], + model_list: list[Mapping[str, object]], details: bool = True, max_concurrency: int | None = None, health_check_skip_disabled_background_models: bool = False, router: "Router | None" = None, - ) -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, Any]]: + ) -> tuple[list[Mapping[str, object]], list[Mapping[str, object]], dict[str, object]]: """ Perform health check with shared state coordination. diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index ba319be01cb..5ea4e85a730 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1550,7 +1550,7 @@ class LiteLLMProxyRequestSetup: headers: dict, request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, - general_settings: dict[str, Any] | None = None, + general_settings: dict[str, object] | None = None, ) -> LitellmDataForBackendLLMCall: """ - Adds user from headers diff --git a/litellm/proxy/mcp_tools.py b/litellm/proxy/mcp_tools.py index eaf9e4b2665..239bb4b2fcf 100644 --- a/litellm/proxy/mcp_tools.py +++ b/litellm/proxy/mcp_tools.py @@ -1,7 +1,4 @@ -from typing import Any - - -def get_current_time(params: dict[str, Any] | None = None) -> str: +def get_current_time(params: dict[str, object] | None = None) -> str: """ Get the current time (hardcoded sample implementation) @@ -18,7 +15,7 @@ def get_current_time(params: dict[str, Any] | None = None) -> str: return "10:30:45 AM" -def get_current_date(params: dict[str, Any] | None = None) -> str: +def get_current_date(params: dict[str, object] | None = None) -> str: """ Get the current date (hardcoded sample implementation) diff --git a/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py b/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py index 3d408accba3..457af931f56 100644 --- a/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py +++ b/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py @@ -96,7 +96,7 @@ class FileContentStreamingHandler: *, custom_llm_provider: str, file_id: str, - data: dict[str, Any], + data: dict[str, object], proxy_logging_obj: "ProxyLogging", user_api_key_dict: "UserAPIKeyAuth", version: str, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index d787c16a534..23ef89b3c5c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -35,6 +35,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( PassthroughStandardLoggingPayload, ) from litellm.types.utils import ( + CompletionTokensDetailsWrapper, LiteLLMBatch, ModelResponse, TextCompletionResponse, @@ -306,13 +307,13 @@ class AnthropicPassthroughLoggingHandler: if not line.startswith("data:"): continue try: - data = json.loads(line[len("data:") :].strip()) + data = _JSON_OBJECT.validate_python(json.loads(line[len("data:") :].strip())) except (json.JSONDecodeError, ValueError): continue if not isinstance(data, dict): continue if data.get("type") == "message_start": - value = (data.get("message") or {}).get(field) + value = _JSON_OBJECT.validate_python(data.get("message") or {}).get(field) if isinstance(value, str) and value: return value return None @@ -398,7 +399,7 @@ class AnthropicPassthroughLoggingHandler: # Anthropic costing reads completion_tokens_details.text_tokens, so the # stale message_start placeholder there must be corrected too or spend # stays undercounted even after completion_tokens is fixed. - details: Final = getattr(usage, "completion_tokens_details", None) + details: Final[CompletionTokensDetailsWrapper | None] = getattr(usage, "completion_tokens_details", None) if details is not None and getattr(details, "text_tokens", None) is not None: details.text_tokens = recovered_output_tokens return usage diff --git a/litellm/proxy/plugin_routes.py b/litellm/proxy/plugin_routes.py index eb6fe7dd177..1dfe5cbc6bb 100644 --- a/litellm/proxy/plugin_routes.py +++ b/litellm/proxy/plugin_routes.py @@ -296,8 +296,8 @@ async def plugin_proxy( # Forward caller identity so the plugin can enforce its own access control. # The plugin MUST NOT trust these as credentials — they are informational. # The plugin_key above is the only authentication mechanism. - user_id: Final = getattr(user_api_key_dict, "user_id", None) - user_role: Final = getattr(user_api_key_dict, "user_role", None) + user_id: Final[object] = getattr(user_api_key_dict, "user_id", None) + user_role: Final[object] = getattr(user_api_key_dict, "user_role", None) if user_id: forward_headers["x-litellm-user-id"] = str(user_id) if user_role: diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index d81471b3c1a..c751091ff9b 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -80,7 +80,7 @@ class AttachmentRegistry: self._config_attachments: tuple[PolicyAttachment, ...] = () self._initialized: bool = False - def load_attachments(self, attachments_config: list[dict[str, Any]]) -> None: + def load_attachments(self, attachments_config: list[dict[str, object]]) -> None: """ Load attachments from a configuration list. diff --git a/litellm/proxy/policy_engine/init_policies.py b/litellm/proxy/policy_engine/init_policies.py index 061a852b701..b6e84109222 100644 --- a/litellm/proxy/policy_engine/init_policies.py +++ b/litellm/proxy/policy_engine/init_policies.py @@ -98,7 +98,7 @@ def _print_policies_on_startup( async def init_policies( policies_config: dict[str, Any], - policy_attachments_config: list[dict[str, Any]] | None = None, + policy_attachments_config: list[dict[str, object]] | None = None, prisma_client: Optional["PrismaClient"] = None, validate_db: bool = True, fail_on_error: bool = True, @@ -184,8 +184,8 @@ async def init_policies( def init_policies_sync( - policies_config: dict[str, Any], - policy_attachments_config: list[dict[str, Any]] | None = None, + policies_config: dict[str, object], + policy_attachments_config: list[dict[str, object]] | None = None, fail_on_error: bool = True, ) -> None: """ @@ -218,7 +218,7 @@ def init_policies_sync( ) -def get_policies_summary() -> dict[str, Any]: +def get_policies_summary() -> dict[str, object]: """ Get a summary of loaded policies for debugging/display. diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index cfc1f481988..c5f7e14e61d 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -22,6 +22,8 @@ from typing import ( cast, # noqa: TID251 # prisma types the condition/pipeline Json columns as str, but reads return decoded values ) +from typing_extensions import ReadOnly + from litellm._logging import verbose_proxy_logger from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import PolicyRepository @@ -160,6 +162,15 @@ def _row_to_policy_db_response(row: _PolicyRow) -> PolicyDBResponse: ) +class _AttributeView(TypedDict): + value: ReadOnly[object] + + +def _attribute_of(source: object, name: str) -> object: + attribute: Final[_AttributeView] = {"value": getattr(source, name)} + return attribute["value"] + + class PolicyRegistry: """ In-memory registry for storing and managing policies. @@ -1020,8 +1031,8 @@ class PolicyRegistry: ) field_diffs: Final[dict[str, dict[str, object]]] = {} for field in compare_fields: - val_a = getattr(resp_a, field) - val_b = getattr(resp_b, field) + val_a = _attribute_of(resp_a, field) + val_b = _attribute_of(resp_b, field) if val_a != val_b: field_diffs[field] = {"version_a": val_a, "version_b": val_b} diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3a831062b0a..bdaa4af4e26 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15117,7 +15117,7 @@ async def _apply_search_filter_to_models( size: int = 50, sort_by: str | None = None, model_name: str | None = None, -) -> tuple[list[dict[str, Any]], int | None]: +) -> tuple[list[dict[str, object]], int | None]: """ Apply search filter to models, querying database for additional matching models. @@ -15267,10 +15267,10 @@ def _normalize_datetime_for_sorting(dt: object) -> datetime | None: def _sort_models( - all_models: list[dict[str, Any]], + all_models: list[dict[str, object]], sort_by: str | None, sort_order: str = "asc", -) -> list[dict[str, Any]]: +) -> list[dict[str, object]]: """ Sort models by the specified field and order. @@ -15560,7 +15560,7 @@ async def _filter_models_by_team_id( prisma_client: PrismaClient, llm_router: Router, user_api_key_dict: UserAPIKeyAuth | None = None, -) -> list[dict[str, Any]]: +) -> list[dict[str, object]]: """ Filter models by team ID. Returns models where: - team_id matches the model's BYOK team_id, OR diff --git a/litellm/proxy/read_model_list.py b/litellm/proxy/read_model_list.py index a1830e7f2bc..92bfa8657cb 100644 --- a/litellm/proxy/read_model_list.py +++ b/litellm/proxy/read_model_list.py @@ -17,10 +17,10 @@ secret manager, which is what makes that resolution work. Returns the resolved from __future__ import annotations import asyncio -from typing import Any, Final +from typing import Final -def read_model_list(config_path: str) -> list[dict[str, Any]]: +def read_model_list(config_path: str) -> list[dict[str, object]]: """Load ``config_path`` via the proxy's own reader and return its resolved ``model_list``.""" from litellm.proxy.proxy_server import ProxyConfig diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index b13042dfb6c..5f6a7432afb 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -10,11 +10,10 @@ https://platform.openai.com/docs/api-reference/responses-streaming import asyncio import json -from collections.abc import Callable, Mapping, Sequence -from typing import TYPE_CHECKING, Final, TypeAlias +from collections.abc import AsyncIterable, Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, runtime_checkable from fastapi import Request, Response -from fastapi.responses import StreamingResponse from starlette.types import Message from typing_extensions import ReadOnly, TypedDict @@ -75,6 +74,12 @@ class _StreamEventParser: parse: Callable[[str], _StreamEvent] = staticmethod(json.loads) +@runtime_checkable +class _StreamsBody(Protocol): + @property + def body_iterator(self) -> AsyncIterable[str | bytes | memoryview]: ... + + def _sse_frame_data(frame: str) -> str | None: return next((line[6:].strip() for line in frame.splitlines() if line.startswith("data: ")), None) @@ -141,7 +146,7 @@ async def background_streaming_task( # Make streaming request. # Pre-call checks (rate limits, guardrails, budget) were already run # before polling ID creation, so skip them here to avoid double-counting. - response: Final[StreamingResponse] = await processor.base_process_llm_request( + response: Final = await processor.base_process_llm_request( request=detach_request_from_client(request), fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -216,13 +221,13 @@ async def background_streaming_task( last_update_time = current_time # Handle StreamingResponse - if not hasattr(response, "body_iterator"): + if not isinstance(response, _StreamsBody): verbose_proxy_logger.warning( "background_streaming_task: response for %s has no body_iterator; this may indicate a misconfiguration or provider error", polling_id, ) - if hasattr(response, "body_iterator"): + if isinstance(response, _StreamsBody): async for chunk in response.body_iterator: # Parse chunk if isinstance(chunk, bytes): diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py index 4ae227162d2..3f82168d958 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -3,8 +3,9 @@ Response Polling Handler for Background Responses with Cache """ import json +from collections.abc import Mapping from datetime import datetime, timezone -from typing import Any, Final +from typing import Final from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger @@ -44,7 +45,7 @@ class ResponsePollingHandler: async def create_initial_state( self, polling_id: str, - request_data: dict[str, Any], + request_data: Mapping[str, object], ) -> ResponsesAPIResponse: """ Create initial state in Redis for a polling request @@ -62,14 +63,16 @@ class ResponsePollingHandler: created_timestamp: Final = int(datetime.now(timezone.utc).timestamp()) # Create OpenAI-compliant response object - response: Final = ResponsesAPIResponse( - id=polling_id, - object="response", - status="queued", # OpenAI native status - created_at=created_timestamp, - output=[], - metadata=request_data.get("metadata", {}), - usage=None, + response: Final = ResponsesAPIResponse.model_validate( + { + "id": polling_id, + "object": "response", + "status": "queued", # OpenAI native status + "created_at": created_timestamp, + "output": [], + "metadata": request_data.get("metadata", {}), + "usage": None, + } ) cache_key: Final = self.get_cache_key(polling_id) @@ -218,7 +221,7 @@ class ResponsePollingHandler: ) @with_service_target(_RESPONSE_POLLING_TARGET) - async def get_state(self, polling_id: str) -> dict[str, Any] | None: + async def get_state(self, polling_id: str) -> Mapping[str, object] | None: """Get current polling state from Redis""" if not self.redis_cache: return None diff --git a/litellm/proxy/types_utils/utils.py b/litellm/proxy/types_utils/utils.py index c3fd237c0d8..62b87dbbf36 100644 --- a/litellm/proxy/types_utils/utils.py +++ b/litellm/proxy/types_utils/utils.py @@ -5,6 +5,12 @@ import os from collections.abc import Callable from typing import Any, Final, Literal, get_type_hints +from typing_extensions import ReadOnly, TypedDict + + +class _AttributeView(TypedDict): + value: ReadOnly[object] + def get_instance_fn(value: str, config_file_path: str | None = None) -> Any: module_name = value @@ -52,7 +58,8 @@ def get_instance_fn(value: str, config_file_path: str | None = None) -> Any: module = importlib.import_module(module_name) # Get the instance from the module - instance: Final[object] = getattr(module, instance_name) + attribute: Final[_AttributeView] = {"value": getattr(module, instance_name)} + instance: Final[object] = attribute["value"] return instance except ImportError as e: @@ -167,7 +174,8 @@ def _load_instance_from_remote_storage(remote_url: str, config_file_path: str | spec.loader.exec_module(module) # Get the instance - instance: Final[object] = getattr(module, instance_name) + attribute: Final[_AttributeView] = {"value": getattr(module, instance_name)} + instance: Final[object] = attribute["value"] # Clean up the temporary file try: diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 8363aaee99a..b97c3a949d5 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -5,6 +5,7 @@ from types import MappingProxyType from typing import Final, Literal from fastapi import HTTPException, Request +from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -21,6 +22,8 @@ from litellm.types.utils import LlmProviders from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.utils import ProviderConfigManager +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) + def _normalize_litellm_params( vector_store: LiteLLM_ManagedVectorStore, @@ -29,8 +32,7 @@ def _normalize_litellm_params( if isinstance(litellm_params, str): normalized: Final = LiteLLM_ManagedVectorStore(**dict(vector_store)) try: - parsed: Final = json.loads(litellm_params) - normalized["litellm_params"] = parsed if isinstance(parsed, dict) else {} + normalized["litellm_params"] = _JSON_OBJECT.validate_python(json.loads(litellm_params)) except (TypeError, ValueError): normalized["litellm_params"] = {} return normalized diff --git a/litellm/proxy/video_endpoints/utils.py b/litellm/proxy/video_endpoints/utils.py index a38226cc253..b2ec2eda761 100644 --- a/litellm/proxy/video_endpoints/utils.py +++ b/litellm/proxy/video_endpoints/utils.py @@ -47,7 +47,7 @@ def get_custom_provider_from_data(data: dict[str, Any]) -> str | None: return None -def encode_character_id_in_response(response: Any, custom_llm_provider: str, model_id: str | None) -> Any: +def encode_character_id_in_response(response: Any, custom_llm_provider: str, model_id: str | None) -> object: if isinstance(response, dict) and response.get("id"): response["id"] = encode_character_id_with_provider( character_id=response["id"], diff --git a/litellm/rag/rag_query.py b/litellm/rag/rag_query.py index 9325547c17d..9f608dbc99a 100644 --- a/litellm/rag/rag_query.py +++ b/litellm/rag/rag_query.py @@ -70,7 +70,7 @@ class RAGQuery: return None @staticmethod - def build_context_message(context_chunks: list[Any]) -> ChatCompletionUserMessage: + def build_context_message(context_chunks: Sequence[_SearchResultView | str | None]) -> ChatCompletionUserMessage: """ Process search results and build a context message. """ diff --git a/litellm/router.py b/litellm/router.py index cdede46ef5c..97d4b481369 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3384,7 +3384,7 @@ class Router: async def _aresponses_streaming_iterator( self, response: "BaseResponsesAPIStreamingIterator", - initial_kwargs: dict[str, Any], + initial_kwargs: dict[str, object], ) -> "BaseResponsesAPIStreamingIterator": """ Wrap a Responses-API streaming iterator so MidStreamFallbackError @@ -5596,7 +5596,7 @@ class Router: async def _aanthropic_messages_streaming_iterator( self, response: AsyncIterator[bytes], - initial_kwargs: dict[str, Any], # mutable-ok: mutated in-place before re-entering the fallback chain + initial_kwargs: dict[str, object], # mutable-ok: mutated in-place before re-entering the fallback chain ) -> AsyncIterator[bytes]: """ Wrap an anthropic_messages (/v1/messages) streaming response so a diff --git a/litellm/secret_managers/google_secret_manager.py b/litellm/secret_managers/google_secret_manager.py index 6d6543cee4b..15e8783b96b 100644 --- a/litellm/secret_managers/google_secret_manager.py +++ b/litellm/secret_managers/google_secret_manager.py @@ -1,7 +1,10 @@ import base64 import os +from collections.abc import Mapping from typing import Final +from pydantic import ConfigDict, TypeAdapter + import litellm from litellm._logging import verbose_logger from litellm.caching.caching import InMemoryCache @@ -11,6 +14,9 @@ from litellm.llms.custom_httpx.http_handler import get_httpx_client from litellm.types.proxy.common_proxy_errors import CommonProxyErrors from litellm.types.secret_managers.main import KeyManagementSystem +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_BASE64_TEXT: Final = TypeAdapter(str, config=ConfigDict(strict=True, hide_input_in_errors=True)) + class GoogleSecretManager(GCSBucketBase): def __init__( @@ -86,12 +92,14 @@ class GoogleSecretManager(GCSBucketBase): ) # Parse the JSON response and return the secret value - secret_data: Final = response.json() - _base64_encoded_value: Final = secret_data.get("payload", {}).get("data") + secret_data: Final = _JSON_OBJECT.validate_python(response.json()) + _base64_encoded_value: Final = _JSON_OBJECT.validate_python(secret_data.get("payload", {})).get("data") # decode the base64 encoded value if _base64_encoded_value is not None: - _decoded_value: Final = base64.b64decode(_base64_encoded_value).decode("utf-8") + _decoded_value: Final = base64.b64decode(_BASE64_TEXT.validate_python(_base64_encoded_value)).decode( + "utf-8" + ) self.cache.set_cache(secret_name, _decoded_value) # Cache the retrieved secret return _decoded_value diff --git a/litellm/utils.py b/litellm/utils.py index 3f10d2a3ebd..4733051f4f6 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5756,7 +5756,7 @@ def get_model_cost_key(potential_key: str) -> str | None: _get_model_cost_key = get_model_cost_key -def _get_model_info_from_model_cost(key: str) -> dict[str, Any]: +def _get_model_info_from_model_cost(key: str) -> dict[str, object]: return litellm.model_cost[key] @@ -8525,7 +8525,7 @@ _get_bundled_model_cost_map = get_bundled_model_cost_map def _get_model_cost_entry_for_provider_config( model: str, provider: LlmProviders, -) -> dict[str, Any]: +) -> dict[str, object]: candidate_keys: Final = (model, f"{provider.value}/{model}") for model_key in candidate_keys: model_info = litellm.model_cost.get(model_key) diff --git a/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py b/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py index 1c87fb7564d..09f09170241 100644 --- a/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py +++ b/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py @@ -618,3 +618,53 @@ class TestHandlerIntegration: api_base=None, ) assert result is not None + + +class TestStreamingContentType: + @pytest.mark.asyncio + @pytest.mark.parametrize( + "response_headers, expected_events", + [ + ([("content-type", "application/json")], [{"jsonrpc": "2.0", "id": "req-001", "result": {}}]), + ( + [("Content-Type", "Application/JSON; charset=utf-8")], + [{"jsonrpc": "2.0", "id": "req-001", "result": {}}], + ), + ( + [("content-type", "text/plain"), ("CONTENT-TYPE", "application/json")], + [{"jsonrpc": "2.0", "id": "req-001", "result": {}}], + ), + ( + [("content-type", "application/json"), ("content-type", "text/event-stream")], + [{"jsonrpc": "2.0", "id": "req-001", "result": {}}], + ), + ([], []), + ([("content-type", "")], []), + ([("content-type", "text/event-stream")], []), + ([("x-content-type", "application/json")], []), + ([("content-type", "application"), ("content-type", "/json")], []), + ], + ) + async def test_streaming_reads_a_single_json_body_only_when_the_content_type_says_json( + self, httpx_transport, response_headers, expected_events + ): + from litellm.a2a_protocol.providers.bedrock_agentcore.config import ( + BedrockAgentCoreA2AConfig, + ) + + with respx.mock(assert_all_called=True) as router: + router.post(url__regex=r".*/invocations.*").mock( + return_value=httpx.Response( + 200, headers=response_headers, content=b'{"jsonrpc": "2.0", "id": "req-001", "result": {}}' + ) + ) + events = [ + event + async for event in BedrockAgentCoreA2AConfig().handle_streaming( + request_id="req-001", + params=SAMPLE_PARAMS, + litellm_params=SAMPLE_LITELLM_PARAMS, + ) + ] + + assert events == expected_events diff --git a/tests/unit/integrations/focus/test_focus_logger.py b/tests/unit/integrations/focus/test_focus_logger.py new file mode 100644 index 00000000000..5ff5c40df47 --- /dev/null +++ b/tests/unit/integrations/focus/test_focus_logger.py @@ -0,0 +1,92 @@ +import pytest + +from litellm.integrations.focus.focus_logger import FocusLogger +from litellm.proxy import proxy_server + +FOCUS_LOCK = "cronjob_lock:focus_export_usage_data" +VANTAGE_LOCK = "cronjob_lock:vantage_export_usage_data" + + +class _InMemoryRedis: + def __init__(self, locks: dict[str, str]) -> None: + self.locks = locks + self.granted: list[tuple[str, str]] = [] + + async def async_set_cache(self, key: str, value: str, nx: bool, ttl: int) -> bool: + if nx and key in self.locks: + return False + self.locks[key] = value + self.granted.append((key, value)) + return True + + async def async_get_cache(self, key: str) -> str | None: + return self.locks.get(key) + + async def async_delete_cache(self, key: str) -> int: + return 1 if self.locks.pop(key, None) is not None else 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("proxy_logging_registered", [True, False]) +async def test_scheduled_focus_export_without_a_redis_lock_runs_until_the_database_boundary( + monkeypatch: pytest.MonkeyPatch, proxy_logging_registered: bool +) -> None: + monkeypatch.setattr(proxy_server.proxy_logging_obj.db_spend_update_writer.pod_lock_manager, "redis_cache", None) + if not proxy_logging_registered: + monkeypatch.setattr(proxy_server, "proxy_logging_obj", None) + monkeypatch.setattr(proxy_server, "prisma_client", None) + logger = FocusLogger(provider="s3", destination_config={"bucket_name": "focus-exports"}) + + with pytest.raises(RuntimeError, match="Database not connected"): + await logger.initialize_focus_export_job() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("locks_before", "locks_after", "granted"), + [ + pytest.param({}, {}, [(FOCUS_LOCK, "this-pod")], id="lock-free"), + pytest.param({FOCUS_LOCK: "this-pod"}, {}, [], id="lock-already-held-by-this-pod"), + pytest.param( + {VANTAGE_LOCK: "another-pod"}, + {VANTAGE_LOCK: "another-pod"}, + [(FOCUS_LOCK, "this-pod")], + id="vantage-lock-held-by-another-pod", + ), + ], +) +async def test_scheduled_focus_export_runs_under_the_focus_pod_lock_and_releases_it_when_the_export_fails( + monkeypatch: pytest.MonkeyPatch, + locks_before: dict[str, str], + locks_after: dict[str, str], + granted: list[tuple[str, str]], +) -> None: + lock_manager = proxy_server.proxy_logging_obj.db_spend_update_writer.pod_lock_manager + redis = _InMemoryRedis(dict(locks_before)) + monkeypatch.setattr(lock_manager, "redis_cache", redis) + monkeypatch.setattr(lock_manager, "pod_id", "this-pod") + monkeypatch.setattr(proxy_server, "prisma_client", None) + logger = FocusLogger(provider="s3", destination_config={"bucket_name": "focus-exports"}) + + with pytest.raises(RuntimeError, match="Database not connected"): + await logger.initialize_focus_export_job() + + assert redis.locks == locks_after + assert redis.granted == granted + + +@pytest.mark.asyncio +async def test_scheduled_focus_export_is_skipped_while_another_pod_holds_the_focus_lock( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lock_manager = proxy_server.proxy_logging_obj.db_spend_update_writer.pod_lock_manager + redis = _InMemoryRedis({FOCUS_LOCK: "another-pod"}) + monkeypatch.setattr(lock_manager, "redis_cache", redis) + monkeypatch.setattr(lock_manager, "pod_id", "this-pod") + monkeypatch.setattr(proxy_server, "prisma_client", None) + logger = FocusLogger(provider="s3", destination_config={"bucket_name": "focus-exports"}) + + await logger.initialize_focus_export_job() + + assert redis.locks == {FOCUS_LOCK: "another-pod"} + assert redis.granted == [] diff --git a/tests/unit/integrations/test_argilla.py b/tests/unit/integrations/test_argilla.py new file mode 100644 index 00000000000..b8948acf011 --- /dev/null +++ b/tests/unit/integrations/test_argilla.py @@ -0,0 +1,110 @@ +import asyncio +import json +import re + +import httpx +import pytest +import respx + +import litellm +from litellm.integrations.argilla import ArgillaLogger + +ARGILLA_BASE_URL = "https://argilla.example.test" +USER_MESSAGE = {"role": "user", "content": "where is my order?"} + + +async def _cancel_background_tasks() -> None: + background_tasks = [task for task in asyncio.all_tasks() if task is not asyncio.current_task()] + for task in background_tasks: + task.cancel() + await asyncio.gather(*background_tasks, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("transformation_object", "standard_logging_object", "expected_fields"), + [ + pytest.param( + {"prompt": "messages", "answer": "response"}, + {"messages": [USER_MESSAGE], "response": "it ships tomorrow"}, + {"prompt": [USER_MESSAGE], "answer": "it ships tomorrow"}, + id="message-list-and-text-response", + ), + pytest.param( + {"answer": "response", "prompt": "messages"}, + {"messages": USER_MESSAGE, "response": {"choices": [{"message": {"content": "it ships tomorrow"}}]}}, + {"answer": "it ships tomorrow", "prompt": [USER_MESSAGE]}, + id="single-message-and-chat-completion-response", + ), + pytest.param( + {"prompt": "messages", "context": "messages", "answer": "response", "label": "response"}, + {"messages": [USER_MESSAGE, USER_MESSAGE], "response": {"choices": [{"message": {}}]}}, + { + "prompt": [USER_MESSAGE, USER_MESSAGE], + "context": [USER_MESSAGE, USER_MESSAGE], + "answer": "", + "label": "", + }, + id="payload-field-reused-by-several-argilla-fields", + ), + ], +) +async def test_argilla_logger_sends_the_record_shaped_by_the_configured_field_mapping( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + transformation_object: dict[str, str], + standard_logging_object: dict[str, object], + expected_fields: dict[str, object], +) -> None: + monkeypatch.setattr(litellm, "argilla_transformation_object", transformation_object) + respx_mock.get(f"{ARGILLA_BASE_URL}/api/v1/me/datasets", params={"name": "support-chats"}).mock( + return_value=httpx.Response(200, json={"items": [{"id": "dataset-403"}]}) + ) + bulk_route = respx_mock.post(f"{ARGILLA_BASE_URL}/api/v1/datasets/dataset-403/records/bulk").mock( + return_value=httpx.Response(200, json={}) + ) + logger = ArgillaLogger( + argilla_api_key="argilla-key", + argilla_dataset_name="support-chats", + argilla_base_url=ARGILLA_BASE_URL, + batch_size=1, + ) + try: + logger.log_success_event({"standard_logging_object": standard_logging_object}, None, None, None) + finally: + await _cancel_background_tasks() + + sent_request = bulk_route.calls.last.request + assert json.loads(sent_request.content) == [{"fields": expected_fields}] + assert sent_request.headers["X-Argilla-Api-Key"] == "argilla-key" + assert logger.argilla_transformation_object is transformation_object + + +@pytest.mark.parametrize( + ("transformation_object", "message"), + [ + pytest.param( + None, + "'litellm.argilla_transformation_object' is required, to log your payload to Argilla.", + id="mapping-not-configured", + ), + pytest.param( + ["messages", "response"], + "'argilla_transformation_object' must be a dictionary, to log your payload to Argilla.", + id="mapping-is-a-list", + ), + pytest.param( + {"prompt": "messages", "spend": "response_cost"}, + "All values in argilla_transformation_object must be a key in SUPPORTED_PAYLOAD_FIELDS, " + "response_cost is not a valid key.", + id="mapping-names-an-unsupported-payload-field", + ), + ], +) +def test_argilla_logger_refuses_to_start_without_a_supported_field_mapping( + monkeypatch: pytest.MonkeyPatch, transformation_object: object, message: str +) -> None: + monkeypatch.setattr(litellm, "argilla_transformation_object", transformation_object) + + with pytest.raises(Exception, match=re.escape(message)): + ArgillaLogger(argilla_api_key="argilla-key") diff --git a/tests/unit/integrations/vantage/__init__.py b/tests/unit/integrations/vantage/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/vantage/test_vantage_logger.py b/tests/unit/integrations/vantage/test_vantage_logger.py new file mode 100644 index 00000000000..8ff2c7fa046 --- /dev/null +++ b/tests/unit/integrations/vantage/test_vantage_logger.py @@ -0,0 +1,92 @@ +import pytest + +from litellm.integrations.vantage.vantage_logger import VantageLogger +from litellm.proxy import proxy_server + +FOCUS_LOCK = "cronjob_lock:focus_export_usage_data" +VANTAGE_LOCK = "cronjob_lock:vantage_export_usage_data" + + +class _InMemoryRedis: + def __init__(self, locks: dict[str, str]) -> None: + self.locks = locks + self.granted: list[tuple[str, str]] = [] + + async def async_set_cache(self, key: str, value: str, nx: bool, ttl: int) -> bool: + if nx and key in self.locks: + return False + self.locks[key] = value + self.granted.append((key, value)) + return True + + async def async_get_cache(self, key: str) -> str | None: + return self.locks.get(key) + + async def async_delete_cache(self, key: str) -> int: + return 1 if self.locks.pop(key, None) is not None else 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("proxy_logging_registered", [True, False]) +async def test_scheduled_vantage_export_without_a_redis_lock_runs_until_the_database_boundary( + monkeypatch: pytest.MonkeyPatch, proxy_logging_registered: bool +) -> None: + monkeypatch.setattr(proxy_server.proxy_logging_obj.db_spend_update_writer.pod_lock_manager, "redis_cache", None) + if not proxy_logging_registered: + monkeypatch.setattr(proxy_server, "proxy_logging_obj", None) + monkeypatch.setattr(proxy_server, "prisma_client", None) + logger = VantageLogger(api_key="vantage-key", integration_token="vantage-token") + + with pytest.raises(RuntimeError, match="Database not connected"): + await logger.initialize_focus_export_job() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("locks_before", "locks_after", "granted"), + [ + pytest.param({}, {}, [(VANTAGE_LOCK, "this-pod")], id="lock-free"), + pytest.param({VANTAGE_LOCK: "this-pod"}, {}, [], id="lock-already-held-by-this-pod"), + pytest.param( + {FOCUS_LOCK: "another-pod"}, + {FOCUS_LOCK: "another-pod"}, + [(VANTAGE_LOCK, "this-pod")], + id="focus-lock-held-by-another-pod", + ), + ], +) +async def test_scheduled_vantage_export_runs_under_the_vantage_pod_lock_and_releases_it_when_the_export_fails( + monkeypatch: pytest.MonkeyPatch, + locks_before: dict[str, str], + locks_after: dict[str, str], + granted: list[tuple[str, str]], +) -> None: + lock_manager = proxy_server.proxy_logging_obj.db_spend_update_writer.pod_lock_manager + redis = _InMemoryRedis(dict(locks_before)) + monkeypatch.setattr(lock_manager, "redis_cache", redis) + monkeypatch.setattr(lock_manager, "pod_id", "this-pod") + monkeypatch.setattr(proxy_server, "prisma_client", None) + logger = VantageLogger(api_key="vantage-key", integration_token="vantage-token") + + with pytest.raises(RuntimeError, match="Database not connected"): + await logger.initialize_focus_export_job() + + assert redis.locks == locks_after + assert redis.granted == granted + + +@pytest.mark.asyncio +async def test_scheduled_vantage_export_is_skipped_while_another_pod_holds_the_vantage_lock( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lock_manager = proxy_server.proxy_logging_obj.db_spend_update_writer.pod_lock_manager + redis = _InMemoryRedis({VANTAGE_LOCK: "another-pod"}) + monkeypatch.setattr(lock_manager, "redis_cache", redis) + monkeypatch.setattr(lock_manager, "pod_id", "this-pod") + monkeypatch.setattr(proxy_server, "prisma_client", None) + logger = VantageLogger(api_key="vantage-key", integration_token="vantage-token") + + await logger.initialize_focus_export_job() + + assert redis.locks == {VANTAGE_LOCK: "another-pod"} + assert redis.granted == [] diff --git a/tests/unit/llms/a2a/chat/test_a2a_chat_transformation.py b/tests/unit/llms/a2a/chat/test_a2a_chat_transformation.py index a419e73a2e1..76bce26e2e0 100644 --- a/tests/unit/llms/a2a/chat/test_a2a_chat_transformation.py +++ b/tests/unit/llms/a2a/chat/test_a2a_chat_transformation.py @@ -1,11 +1,15 @@ """Tests for litellm/llms/a2a/chat/transformation.py response transform.""" +import uuid from unittest.mock import MagicMock +import httpx import pytest +from pydantic import ValidationError from litellm.llms.a2a.chat.streaming_iterator import A2AModelResponseIterator from litellm.llms.a2a.chat.transformation import A2AConfig +from litellm.llms.a2a.common_utils import A2AError from litellm.types.utils import ModelResponse @@ -121,3 +125,104 @@ def test_resolve_agent_config_from_registry_returns_the_explicit_headers_object_ assert resolved[2] is headers assert resolved[:2] == ("http://explicit.example", "explicit-key") assert optional_params == {"stream": True} + + +def _transform(payload: object, headers: dict[str, str] | None = None) -> ModelResponse: + return A2AConfig().transform_response( + model="a2a/test-agent", + raw_response=httpx.Response(200, json=payload, headers=headers), + model_response=ModelResponse(), + logging_obj=None, + request_data={}, + messages=[{"role": "user", "content": "hi there agent"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +@pytest.mark.parametrize( + ("result", "expected_text"), + [ + ({"kind": "message", "parts": [{"kind": "text", "text": "direct"}]}, "direct"), + ({"message": {"parts": [{"kind": "text", "text": "nested"}]}}, "nested"), + ({"kind": "task", "status": {"message": {"parts": [{"kind": "text", "text": "status"}]}}}, "status"), + ({"kind": "task", "artifacts": [{"parts": [{"kind": "text", "text": "artifact"}]}]}, "artifact"), + ({"kind": "task", "status": {"state": "submitted"}}, ""), + (None, ""), + ("plain text", ""), + ([{"kind": "text", "text": "not a result object"}], ""), + ], +) +def test_transform_response_extracts_the_text_of_each_result_shape(result: object, expected_text: str): + response = _transform({"jsonrpc": "2.0", "id": "resp-1", "result": result, "extra": None}) + + assert response.choices[0].message.content == expected_text + assert response.choices[0].finish_reason == "stop" + assert response.id == "resp-1" + assert response.model == "a2a/test-agent" + + +@pytest.mark.parametrize("payload", [{}, {"jsonrpc": "2.0", "result": {"parts": [{"kind": "text", "text": "4"}]}}]) +def test_transform_response_generates_an_id_when_the_agent_sends_none(payload: dict[str, object]): + response = _transform(payload) + + assert str(uuid.UUID(response.id)) == response.id + + +@pytest.mark.parametrize( + ("error", "expected_message"), + [ + ({"code": -32600, "message": "Invalid Request"}, "A2A error: Invalid Request"), + ({"code": -32600}, "A2A error: Unknown error"), + ({}, "A2A error: Unknown error"), + ({"message": None}, "A2A error: None"), + ({"message": 429}, "A2A error: 429"), + ({"message": {"detail": "nested"}}, "A2A error: {'detail': 'nested'}"), + ], +) +def test_transform_response_raises_the_agent_error_with_its_message(error: dict[str, object], expected_message: str): + payload = {"jsonrpc": "2.0", "id": "1", "error": error, "result": {"parts": [{"kind": "text", "text": "4"}]}} + + with pytest.raises(A2AError) as exc_info: + _transform(payload, headers={"x-agent": "billing"}) + + assert exc_info.value.message == expected_message + assert exc_info.value.status_code == 200 + assert exc_info.value.headers["x-agent"] == "billing" + + +@pytest.mark.parametrize("error", [None, "429 Too Many Requests", 429, False, ["Request timed out"]]) +def test_transform_response_rejects_an_error_member_that_is_not_an_object(error: object): + with pytest.raises(ValidationError) as exc_info: + _transform({"jsonrpc": "2.0", "id": "1", "error": error, "result": {"parts": []}}) + + assert [(detail["type"], detail["loc"]) for detail in exc_info.value.errors()] == [("dict_type", ())] + assert "input_value" not in str(exc_info.value) + + +@pytest.mark.parametrize("payload", [[], [{"result": {"parts": []}}], ["rate limit"], "upstream unavailable"]) +def test_transform_response_rejects_a_body_that_is_not_an_object(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert [(detail["type"], detail["loc"]) for detail in exc_info.value.errors()] == [("dict_type", ())] + assert "input_value" not in str(exc_info.value) + + +def test_transform_response_reports_an_undecodable_body_as_an_agent_error(): + with pytest.raises(A2AError) as exc_info: + A2AConfig().transform_response( + model="a2a/test-agent", + raw_response=httpx.Response(502, content=b"Bad Gateway"), + model_response=ModelResponse(), + logging_obj=None, + request_data={}, + messages=[{"role": "user", "content": "hi there agent"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert exc_info.value.status_code == 502 + assert exc_info.value.message.startswith("Failed to parse A2A response: ") diff --git a/tests/unit/llms/aiml/image_generation/test_aiml_image_generation_transformation.py b/tests/unit/llms/aiml/image_generation/test_aiml_image_generation_transformation.py index 5ac4c7c4643..ea9ace5e994 100644 --- a/tests/unit/llms/aiml/image_generation/test_aiml_image_generation_transformation.py +++ b/tests/unit/llms/aiml/image_generation/test_aiml_image_generation_transformation.py @@ -1,7 +1,9 @@ import os +from datetime import datetime +import httpx import pytest - +from pydantic import ValidationError os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" @@ -9,13 +11,11 @@ import litellm litellm.model_cost = litellm.get_model_cost_map(url="") -from litellm.llms.aiml.image_generation.cost_calculator import ( - cost_calculator as aiml_cost_calculator, -) +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.aiml.image_generation.transformation import ( AimlImageGenerationConfig, ) -from litellm.types.utils import ImageObject, ImageResponse +from litellm.types.utils import ImageResponse def test_openai_style_model_supports_full_openai_param_surface(): @@ -130,3 +130,126 @@ def test_openai_style_unsupported_param_dropped_with_drop_params(): assert mapped == {} +def _transform_response(payload: object) -> ImageResponse: + return AimlImageGenerationConfig().transform_image_generation_response( + model="flux-pro", + raw_response=httpx.Response(200, json=payload), + model_response=ImageResponse(), + logging_obj=Logging( + model="flux-pro", + messages=[], + stream=False, + call_type="image_generation", + start_time=datetime(2026, 1, 1), + litellm_call_id="aiml-image-generation-test", + function_id="aiml-image-generation-test", + ), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +@pytest.mark.parametrize( + ("payload", "expected"), + [ + ( + { + "data": [ + {"url": "https://cdn.aiml.example/a.png", "revised_prompt": "a red fox", "width": 1024}, + "not-an-image", + {"b64_json": "QUJD", "revised_prompt": "a blue fox"}, + [], + {"b64_json": "", "image_base64": "REVG"}, + {"content_type": "image/png"}, + ] + }, + [ + (None, "https://cdn.aiml.example/a.png", "a red fox"), + ("QUJD", None, "a blue fox"), + ("REVG", None, None), + ], + ), + ( + { + "output": { + "choices": [ + {"image_base64": "QUJD", "url": "https://cdn.aiml.example/ignored.png"}, + "skipped", + {"url": "https://cdn.aiml.example/b.png"}, + {"finish_reason": "stop"}, + ] + } + }, + [("QUJD", None, None), (None, "https://cdn.aiml.example/b.png", None)], + ), + ( + { + "images": [ + {"url": "https://cdn.aiml.example/c.png", "image_base64": "aWdub3JlZA=="}, + [1, 2], + {"image_base64": "REVG"}, + {"seed": 7}, + ] + }, + [(None, "https://cdn.aiml.example/c.png", None), ("REVG", None, None)], + ), + ({"output": {"choices": {"finish_reason": "stop"}}}, []), + ({"output": {"choices": "pending"}}, []), + ({"images": {"count": 0}}, []), + ({"images": ""}, []), + ({"data": "pending"}, []), + ({"id": "gen-1"}, []), + ([], []), + ], +) +def test_transform_image_generation_response_reads_images_from_every_response_shape( + payload: object, expected: list[tuple[str | None, str | None, str | None]] +): + result = _transform_response(payload) + + assert [(image.b64_json, image.url, image.revised_prompt) for image in result.data] == expected + + +@pytest.mark.parametrize( + "payload", + [ + {"data": [7]}, + {"data": [None]}, + {"data": ["aiml-secret-url-token"]}, + {"data": [["b64_json", "aiml-secret-url-token"]]}, + {"output": {"choices": [False]}}, + {"output": {"choices": ["aiml-secret image_base64"]}}, + {"output": {"choices": None}}, + {"images": [1.5]}, + {"images": ["aiml-secret-url-token"]}, + {"images": {"url": "aiml-secret-url-token"}}, + {"images": 3}, + ], +) +def test_transform_image_generation_response_rejects_malformed_image_entries_without_echoing_them(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform_response(payload) + + assert "aiml-secret" not in str(exc_info.value) + + +def _rejection_text(payload: object) -> str: + with pytest.raises(ValidationError) as exc_info: + _transform_response(payload) + + return str(exc_info.value) + + +@pytest.mark.parametrize("position", [403, 429]) +@pytest.mark.parametrize("malformed_entry", [None, 7, "url"]) +def test_transform_image_generation_response_rejection_text_is_the_same_wherever_the_malformed_entry_sits( + position: int, malformed_entry: object +): + first = [malformed_entry] + late = [*[{"seed": 1}] * position, malformed_entry] + + assert _rejection_text({"data": late}) == _rejection_text({"data": first}) + assert _rejection_text({"output": {"choices": late}}) == _rejection_text({"output": {"choices": first}}) + assert _rejection_text({"images": late}) == _rejection_text({"images": first}) diff --git a/tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py b/tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py index e9736d47537..b5c7b337806 100644 --- a/tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py +++ b/tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py @@ -6,7 +6,9 @@ OpenAI-compatible file operations and Anthropic's Files API format. """ import asyncio +import datetime import io +import json import threading import time @@ -16,6 +18,7 @@ from openai.types.file_deleted import FileDeleted from pydantic import ValidationError from unittest.mock import Mock, patch +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.anthropic.files.transformation import ( AnthropicFilesConfig, ANTHROPIC_FILES_API_BASE, @@ -663,3 +666,151 @@ class TestProviderConfigRegistration: ) assert config is not None assert isinstance(config, AnthropicFilesConfig) + + +def _listed_files(body: bytes) -> list[OpenAIFileObject]: + logging_obj = Logging( + model="", + messages=[], + stream=False, + call_type="file_list", + start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-id", + function_id="function-id", + ) + return AnthropicFilesConfig().transform_list_files_response( + raw_response=httpx.Response(200, content=body), logging_obj=logging_obj, litellm_params={} + ) + + +def _anthropic_file(index: int) -> dict[str, object]: + return { + "id": f"file-{index}", + "type": "file", + "filename": f"doc-{index}.pdf", + "mime_type": "application/pdf", + "size_bytes": index, + "created_at": "2025-01-01T00:00:00Z", + "downloadable": False, + } + + +def test_list_files_response_maps_every_anthropic_file_in_provider_order() -> None: + result = _listed_files( + json.dumps( + { + "data": [ + _anthropic_file(100), + {"id": "file-b", "filename": "b.jsonl", "bytes": 200, "created_at": "2025-01-02T00:00:00Z"}, + {"id": "file-c", "filename": "c.jsonl", "size_bytes": "300", "purpose": "batch", "created_at": ""}, + ], + "has_more": False, + "first_id": "file-100", + "last_id": "file-c", + } + ).encode() + ) + + assert [(listed.id, listed.filename, listed.bytes, listed.purpose) for listed in result] == [ + ("file-100", "doc-100.pdf", 100, "messages"), + ("file-b", "b.jsonl", 200, "messages"), + ("file-c", "c.jsonl", 300, "batch"), + ] + assert [type(listed.bytes) for listed in result] == [int, int, int] + assert [listed.created_at for listed in result[:2]] == [1735689600, 1735776000] + assert result[0] == OpenAIFileObject( + id="file-100", + bytes=100, + created_at=1735689600, + filename="doc-100.pdf", + object="file", + purpose="messages", + status="uploaded", + status_details=None, + ) + + +@pytest.mark.parametrize( + "body", + [ + b"{}", + b'{"has_more": false}', + b'{"data": []}', + b'{"data": ""}', + b'{"data": {}}', + b'{"data": [], "unknown": [1, {"nested": null}]}', + ], +) +def test_list_files_response_without_any_file_is_an_empty_list(body: bytes) -> None: + assert _listed_files(body) == [] + + +@pytest.mark.parametrize("body", [b'["file-abc123"]', b'"file-abc123"', b"null", b"7"]) +def test_list_files_response_rejects_a_body_that_is_not_a_json_object(body: bytes) -> None: + with pytest.raises(ValidationError) as exc_info: + _listed_files(body) + + assert [error["type"] for error in exc_info.value.errors()] == ["dict_type"] + + +@pytest.mark.parametrize("body", [b'{"data": null}', b'{"data": 7}', b'{"data": 0}', b'{"data": true}']) +def test_list_files_response_rejects_data_that_cannot_be_iterated(body: bytes) -> None: + with pytest.raises(ValidationError) as exc_info: + _listed_files(body) + + assert [error["type"] for error in exc_info.value.errors()] == ["iterable_type"] + + +@pytest.mark.parametrize( + "data", + [ + ["file-abc123"], + [None], + [7], + [["file-abc123"]], + [_anthropic_file(0), "file-abc123"], + "file-abc123", + {"file-abc123": _anthropic_file(0)}, + ], +) +def test_list_files_response_rejects_a_listed_file_that_is_not_a_json_object(data: object) -> None: + with pytest.raises(ValidationError) as exc_info: + _listed_files(json.dumps({"data": data}).encode()) + + assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [("dict_type", ())] + + +def test_list_files_response_error_for_a_late_malformed_file_does_not_name_its_position() -> None: + files = [_anthropic_file(index) for index in range(403)] + + with pytest.raises(ValidationError) as exc_info: + _listed_files(json.dumps({"data": [*files, "file-abc123"]}).encode()) + + assert "403" not in str(exc_info.value) + assert "429" not in str(exc_info.value) + assert "file-abc123" not in str(exc_info.value) + + +def test_list_files_response_reports_every_invalid_field_of_a_listed_file() -> None: + with pytest.raises(ValidationError) as exc_info: + _listed_files(b'{"data": [{"id": null, "filename": 7, "size_bytes": "many", "purpose": "unknown"}]}') + + assert [error["loc"] for error in exc_info.value.errors()] == [("id",), ("bytes",), ("filename",), ("purpose",)] + assert "many" in str(exc_info.value) + + +@pytest.mark.parametrize( + "data, expected_errors", + [ + ([{"id": None}, "file-abc123"], [("string_type", ("id",))]), + (["file-abc123", {"id": None}], [("dict_type", ())]), + ([{"id": None}, {"filename": 7}], [("string_type", ("id",))]), + ], +) +def test_list_files_response_reports_only_the_first_bad_file_in_provider_order( + data: object, expected_errors: object +) -> None: + with pytest.raises(ValidationError) as exc_info: + _listed_files(json.dumps({"data": data}).encode()) + + assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == expected_errors diff --git a/tests/unit/llms/azure_ai/anthropic/__init__.py b/tests/unit/llms/azure_ai/anthropic/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/__init__.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py new file mode 100644 index 00000000000..29380c6fc53 --- /dev/null +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -0,0 +1,188 @@ +import json +import operator +from typing import Final + +import httpx +import pytest +import respx +from pydantic import ValidationError + +import litellm +from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import AzureAIAnthropicTokenCounter +from litellm.types.utils import TokenCountResponse + +API_BASE: Final = "https://my-resource.services.ai.azure.com" +COUNT_URL: Final = f"{API_BASE}/anthropic/v1/messages/count_tokens" + + +async def _count(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> TokenCountResponse | None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + respx.post(COUNT_URL).mock(return_value=response) + return await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}}, + request_model="azure_ai/claude-sonnet-4-5", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("payload", "expected_total"), + [ + ({"input_tokens": 15}, 15), + ({"input_tokens": "15"}, 15), + ({"input_tokens": 15.0}, 15), + ({"input_tokens": True}, 1), + ({"input_tokens": -3}, -3), + ({}, 0), + ({"type": "error", "error": {"message": "429 rate limit"}}, 0), + ({"input_tokens": 7, "context_management": {"original_input_tokens": 9}, "nested": {"a": [1, None]}}, 7), + ], +) +async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload( + monkeypatch: pytest.MonkeyPatch, payload: dict[str, object], expected_total: int +) -> None: + result: Final = await _count(monkeypatch, httpx.Response(200, json=payload)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": expected_total, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "original_response": payload, + "error": False, + "error_message": None, + "status_code": None, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "rejected_count", + [ + None, + 15.5, + "Request timed out", + "429 Too Many Requests", + [15], + [*range(403), "Request timed out"], + {"value": 429}, + ], +) +async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_validation_message( + monkeypatch: pytest.MonkeyPatch, rejected_count: object +) -> None: + with pytest.raises(ValidationError) as rejection: + TokenCountResponse.model_validate( + { + "total_tokens": rejected_count, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + } + ) + + result: Final = await _count(monkeypatch, httpx.Response(200, json={"input_tokens": rejected_count})) + + assert result is not None + assert [error["loc"] for error in rejection.value.errors()] == [("total_tokens",)] + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "original_response": None, + "error": True, + "error_message": str(rejection.value), + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [[*({"input_tokens": position} for position in range(403)), "429"], "input_tokens", 15, 0, False], +) +async def test_count_payload_that_is_not_an_object_becomes_an_error_response_naming_the_failed_lookup( + monkeypatch: pytest.MonkeyPatch, payload: object +) -> None: + with pytest.raises(AttributeError) as failed_lookup: + operator.attrgetter("get")(payload) + + result: Final = await _count(monkeypatch, httpx.Response(200, json=payload)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "original_response": None, + "error": True, + "error_message": str(failed_lookup.value), + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +async def test_null_count_payload_falls_back_to_local_counting(monkeypatch: pytest.MonkeyPatch) -> None: + assert await _count(monkeypatch, httpx.Response(200, content=b"null")) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("upstream_status", "upstream_text"), + [(429, "rate limit exceeded"), (403, "forbidden"), (408, "Request timed out")], +) +async def test_upstream_error_status_becomes_an_error_response_with_that_status_and_text( + monkeypatch: pytest.MonkeyPatch, upstream_status: int, upstream_text: str +) -> None: + result: Final = await _count(monkeypatch, httpx.Response(upstream_status, text=upstream_text)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "original_response": None, + "error": True, + "error_message": upstream_text, + "status_code": upstream_status, + } + ) + + +@pytest.mark.asyncio +async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with pytest.raises(json.JSONDecodeError) as undecodable: + json.loads("not json") + + result: Final = await _count(monkeypatch, httpx.Response(200, text="not json")) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "original_response": None, + "error": True, + "error_message": f"CountTokens processing error: {undecodable.value}", + "status_code": 500, + } + ) diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py index 4db786668b8..6f326bfec44 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py @@ -1,8 +1,8 @@ -import asyncio -import json from unittest.mock import Mock +import httpx import pytest +from pydantic import ValidationError # Ensure the project root is on the import path so `litellm` can be imported when # tests are executed from any working directory. @@ -13,6 +13,21 @@ from litellm.llms.bedrock.chat.invoke_transformations.amazon_qwen2_transformatio from litellm.types.utils import ModelResponse +def _transform(body: object) -> ModelResponse: + return AmazonQwen2Config().transform_response( + model="qwen2/test-model", + raw_response=httpx.Response(200, json=body), + model_response=ModelResponse(), + logging_obj=Mock(), + request_data={}, + messages=[{"role": "user", "content": "Hello!"}], + optional_params={}, + litellm_params={}, + encoding=None, + api_key="test-key", + ) + + def test_qwen2_get_supported_params(): """Test that Qwen2 config returns correct supported parameters""" config = AmazonQwen2Config() @@ -253,10 +268,89 @@ def test_qwen2_transform_response_without_usage(): assert result.choices[0]["finish_reason"] == "stop" +@pytest.mark.parametrize( + ("usage", "expected_counts"), + [ + pytest.param({}, (0, 0, 0), id="no-counts"), + pytest.param( + {"prompt_tokens": None, "completion_tokens": None, "total_tokens": None}, (0, 0, 0), id="null-counts" + ), + pytest.param({"prompt_tokens": 7}, (7, 0, 0), id="only-prompt-count"), + pytest.param( + {"prompt_tokens": "7", "completion_tokens": 3.0, "total_tokens": ""}, (7, 3, 0), id="loosely-typed-counts" + ), + pytest.param( + {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3, "latency_ms": [12, None]}, + (1, 2, 3), + id="unmodelled-field", + ), + ], +) +def test_qwen2_transform_response_reads_usage_counts_that_are_missing_null_or_loosely_typed( + usage: dict[str, object], expected_counts: tuple[int, int, int] +) -> None: + result = _transform({"text": "Hi", "usage": usage}) + + reported = result.usage + assert (reported.prompt_tokens, reported.completion_tokens, reported.total_tokens) == expected_counts + assert result.choices[0].message.content == "Hi" + + +def test_qwen2_transform_response_yields_empty_content_when_the_body_carries_no_generated_text() -> None: + result = _transform({"stop_reason": "length", "metrics": {"latency_ms": 12}}) + + assert result.choices[0].message.content == "" + assert result.choices[0].finish_reason == "stop" + + +@pytest.mark.parametrize( + "body", + [ + pytest.param({"text": ["secret-output"]}, id="generated-text-is-a-list"), + pytest.param({"text": None, "debug": "secret-output"}, id="generated-text-is-null"), + pytest.param({"text": "Hi", "usage": ["secret-output"]}, id="usage-is-a-list"), + pytest.param(["secret-output"], id="body-is-a-list"), + ], +) +def test_qwen2_transform_response_rejects_a_malformed_body_without_echoing_it(body: object) -> None: + with pytest.raises(ValidationError) as exc_info: + _transform(body) + + assert "secret-output" not in str(exc_info.value) + + +def test_qwen2_transform_response_reports_every_count_that_is_not_a_number_with_its_value() -> None: + with pytest.raises(ValidationError) as exc_info: + _transform({"text": "Hi", "usage": {"prompt_tokens": "many", "total_tokens": [4, 2]}}) + + assert exc_info.value.title == "Usage" + assert {error["loc"]: error["input"] for error in exc_info.value.errors()} == { + ("prompt_tokens",): "many", + ("total_tokens",): [4, 2], + } + assert "input_value='many'" in str(exc_info.value) + assert "input_value=[4, 2]" in str(exc_info.value) + + +@pytest.mark.parametrize( + ("body", "expected_content"), + [ + pytest.param({"generation": "", "text": "from text"}, "from text", id="empty-generation"), + pytest.param({"generation": None, "text": "from text"}, "from text", id="null-generation"), + pytest.param({"generation": "from generation", "text": None}, "from generation", id="null-text"), + pytest.param({"generation": None, "text": ""}, "", id="both-empty"), + ], +) +def test_qwen2_transform_response_falls_back_to_text_when_generation_is_empty_or_null( + body: dict[str, object], expected_content: str +) -> None: + assert _transform(body).choices[0].message.content == expected_content + + def test_qwen2_provider_detection(): """Test that Qwen2 provider is correctly detected from model names""" - from litellm.utils import ProviderConfigManager from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager # Test with qwen2/ prefix config = ProviderConfigManager.get_provider_chat_config( diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py index e011b1fca2b..a147dbb2c9b 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py @@ -1,8 +1,8 @@ -import asyncio -import json from unittest.mock import Mock +import httpx import pytest +from pydantic import ValidationError # Ensure the project root is on the import path so `litellm` can be imported when # tests are executed from any working directory. @@ -13,6 +13,21 @@ from litellm.llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformatio from litellm.types.utils import ModelResponse +def _transform(body: object) -> ModelResponse: + return AmazonQwen3Config().transform_response( + model="qwen3/test-model", + raw_response=httpx.Response(200, json=body), + model_response=ModelResponse(), + logging_obj=Mock(), + request_data={}, + messages=[{"role": "user", "content": "Hello!"}], + optional_params={}, + litellm_params={}, + encoding=None, + api_key="test-key", + ) + + def test_qwen3_get_supported_params(): """Test that Qwen3 config returns correct supported parameters""" config = AmazonQwen3Config() @@ -181,10 +196,74 @@ def test_qwen3_transform_response_without_usage(): assert result.choices[0]["finish_reason"] == "stop" +@pytest.mark.parametrize( + ("usage", "expected_counts"), + [ + pytest.param({}, (0, 0, 0), id="no-counts"), + pytest.param( + {"prompt_tokens": None, "completion_tokens": None, "total_tokens": None}, (0, 0, 0), id="null-counts" + ), + pytest.param({"prompt_tokens": 7}, (7, 0, 0), id="only-prompt-count"), + pytest.param( + {"prompt_tokens": "7", "completion_tokens": 3.0, "total_tokens": ""}, (7, 3, 0), id="loosely-typed-counts" + ), + pytest.param( + {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3, "latency_ms": [12, None]}, + (1, 2, 3), + id="unmodelled-field", + ), + ], +) +def test_qwen3_transform_response_reads_usage_counts_that_are_missing_null_or_loosely_typed( + usage: dict[str, object], expected_counts: tuple[int, int, int] +) -> None: + result = _transform({"generation": "Hi", "usage": usage}) + + reported = result.usage + assert (reported.prompt_tokens, reported.completion_tokens, reported.total_tokens) == expected_counts + assert result.choices[0].message.content == "Hi" + + +def test_qwen3_transform_response_yields_empty_content_when_the_body_carries_no_generated_text() -> None: + result = _transform({"stop_reason": "length", "metrics": {"latency_ms": 12}}) + + assert result.choices[0].message.content == "" + assert result.choices[0].finish_reason == "stop" + + +@pytest.mark.parametrize( + "body", + [ + pytest.param({"generation": ["secret-output"]}, id="generated-text-is-a-list"), + pytest.param({"generation": None, "debug": "secret-output"}, id="generated-text-is-null"), + pytest.param({"generation": "Hi", "usage": ["secret-output"]}, id="usage-is-a-list"), + pytest.param(["secret-output"], id="body-is-a-list"), + ], +) +def test_qwen3_transform_response_rejects_a_malformed_body_without_echoing_it(body: object) -> None: + with pytest.raises(ValidationError) as exc_info: + _transform(body) + + assert "secret-output" not in str(exc_info.value) + + +def test_qwen3_transform_response_reports_every_count_that_is_not_a_number_with_its_value() -> None: + with pytest.raises(ValidationError) as exc_info: + _transform({"generation": "Hi", "usage": {"prompt_tokens": "many", "total_tokens": [4, 2]}}) + + assert exc_info.value.title == "Usage" + assert {error["loc"]: error["input"] for error in exc_info.value.errors()} == { + ("prompt_tokens",): "many", + ("total_tokens",): [4, 2], + } + assert "input_value='many'" in str(exc_info.value) + assert "input_value=[4, 2]" in str(exc_info.value) + + def test_qwen3_provider_detection(): """Test that Qwen3 provider is correctly detected from model names""" - from litellm.utils import ProviderConfigManager from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager # Test with qwen3/ prefix config = ProviderConfigManager.get_provider_chat_config( diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py index edf93e53dc2..d9b01c4e75a 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py @@ -401,6 +401,27 @@ def test_router_deployment_with_a_non_numeric_stream_chunk_size_gets_a_400_befor send.assert_not_called() +@pytest.mark.parametrize( + ("model", "expected_provider"), + [ + pytest.param("us.anthropic.claude-3-5-sonnet-20240620-v1:0", "anthropic", id="region-prefixed-model-id"), + pytest.param( + "arn:aws:bedrock:us-east-1:123:inference-profile/us.meta.llama3-2-90b-instruct-v1:0", + "meta", + id="inference-profile-arn", + ), + pytest.param("my-anthropic-and-cohere-model", "cohere", id="two-providers-named-earliest-listed-wins"), + pytest.param("us.twelvelabs.pegasus-1-2-v1:0", "twelvelabs", id="provider-listed-late"), + pytest.param("unknown-model", None, id="no-provider-named"), + pytest.param("", None, id="empty-model"), + ], +) +def test_get_bedrock_invoke_provider_finds_a_provider_named_anywhere_in_the_model( + model: str, expected_provider: str | None +) -> None: + assert AmazonInvokeConfig.get_bedrock_invoke_provider(model) == expected_provider + + @pytest.mark.respx(assert_all_called=False) def test_transform_request_meta_llama(monkeypatch, respx_mock): monkeypatch.setattr(litellm, "known_tokenizer_config", dict(litellm.known_tokenizer_config)) diff --git a/tests/unit/llms/brave/__init__.py b/tests/unit/llms/brave/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/brave/search/__init__.py b/tests/unit/llms/brave/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/brave/search/test_brave_search_transformation.py b/tests/unit/llms/brave/search/test_brave_search_transformation.py new file mode 100644 index 00000000000..8523c7c462d --- /dev/null +++ b/tests/unit/llms/brave/search/test_brave_search_transformation.py @@ -0,0 +1,123 @@ +import httpx +import pytest +from pydantic import ValidationError + +from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult +from litellm.llms.brave.search.transformation import BraveSearchConfig + + +def _transform(payload: object, query: str = "q=litellm") -> SearchResponse: + request = httpx.Request("GET", f"https://api.search.brave.com/res/v1/web/search?{query}") + return BraveSearchConfig().transform_search_response( + raw_response=httpx.Response(200, json=payload, request=request), logging_obj=None + ) + + +def _entry(index: int) -> dict[str, str]: + return {"title": f"Title {index}", "url": f"https://example.com/{index}", "description": f"Snippet {index}"} + + +def test_brave_results_are_mapped_section_by_section_with_normalized_dates(): + response = _transform( + { + "type": "search", + "query": {"original": "litellm"}, + "web": { + "type": "search", + "results": [ + { + "title": "LiteLLM", + "url": "https://example.com/litellm", + "description": "Call every LLM API", + "page_age": "2024-03-02T10:00:00", + "age": "January 1, 2020", + "fetched_content_timestamp": 1700000000, + "meta_url": {"hostname": "example.com"}, + }, + {"title": "Docs", "url": "https://example.com/docs", "description": "Docs", "page_age": None}, + ], + }, + "news": { + "results": [ + {"title": "News", "url": "https://example.com/news", "description": "Fresh", "age": "March 5, 2024"} + ] + }, + "videos": {"results": [{"title": "Video", "url": "https://example.com/video"}]}, + } + ) + + assert response.object == "search" + assert response.results == [ + SearchResult( + title="LiteLLM", + url="https://example.com/litellm", + snippet="Call every LLM API", + date="2024-03-02", + last_updated="2023-11-14", + ), + SearchResult(title="Docs", url="https://example.com/docs", snippet="Docs"), + SearchResult(title="News", url="https://example.com/news", snippet="Fresh", date="2024-03-05"), + SearchResult(title="Video", url="https://example.com/video", snippet=""), + ] + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"type": "search", "query": {"original": "litellm"}}, + {"web": {}}, + {"web": {"results": []}, "news": {"results": []}}, + {"web": {"results": ""}}, + {"web": {"results": {}}}, + ], +) +def test_brave_response_without_results_is_empty(payload): + assert _transform(payload).results == [] + + +def test_brave_result_missing_optional_fields_defaults_to_empty_strings_and_no_dates(): + assert _transform({"web": {"results": [{"profile": {"name": "Example"}}]}}).results == [ + SearchResult(title="", url="", snippet="", date=None, last_updated=None) + ] + + +def test_brave_count_limits_results_across_sections_without_reading_entries_past_the_limit(): + response = _transform( + { + "web": {"results": [_entry(1), _entry(2), "not a result object"]}, + "news": {"results": ["not a result object"]}, + }, + query="q=litellm&count=2", + ) + + assert [result.title for result in response.results] == ["Title 1", "Title 2"] + + +def test_brave_result_filter_restricts_the_sections_that_are_read(): + response = _transform( + {"web": {"results": [_entry(1)]}, "news": {"results": [_entry(2)]}, "videos": None}, + query="q=litellm&result_filter=news", + ) + + assert [result.title for result in response.results] == ["Title 2"] + + +@pytest.mark.parametrize("field", ["title", "url", "description"]) +@pytest.mark.parametrize("value", [None, 7, ["text"]]) +def test_brave_result_with_non_string_field_is_rejected(field, value): + with pytest.raises(ValidationError): + _transform({"web": {"results": [{**_entry(1), field: value}]}}) + + +def test_brave_result_with_several_non_string_fields_reports_every_field(): + with pytest.raises(ValidationError) as exc_info: + _transform({"web": {"results": [{"title": 7, "url": None, "description": ["text"]}]}}) + + assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",)] + + +@pytest.mark.parametrize("payload", [{"web": {"results": [_entry(1)]}}, {}, [], "plain text body"]) +def test_brave_non_integer_count_is_reported_whatever_the_response_body_is(payload): + with pytest.raises(ValueError, match="invalid literal for int"): + _transform(payload, query="q=litellm&count=abc") diff --git a/tests/unit/llms/cohere/embed/test_transformation.py b/tests/unit/llms/cohere/embed/test_transformation.py index d7a2e604fc6..c0f355cfb57 100644 --- a/tests/unit/llms/cohere/embed/test_transformation.py +++ b/tests/unit/llms/cohere/embed/test_transformation.py @@ -6,6 +6,7 @@ from typing import Final import httpx import pytest import respx +from pydantic import ValidationError import litellm from litellm.caching.llm_caching_handler import LLMClientCache @@ -13,6 +14,144 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.cohere.embed.transformation import CohereEmbeddingConfig from litellm.types.utils import EmbeddingResponse +FLOAT_VECTORS = [[0.1, 0.2], [0.3, 0.4]] +INT8_VECTORS = [[1, -2], [3, 4]] + + +def _transform(payload: object) -> EmbeddingResponse: + logging_obj = Logging( + model="embed-english-v3.0", + messages=[], + stream=False, + call_type="embedding", + start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-id", + function_id="function-id", + ) + logging_obj.model_call_details["input"] = ["hello", "world"] + return CohereEmbeddingConfig().transform_embedding_response( + model="embed-english-v3.0", + raw_response=httpx.Response(200, json=payload), + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + api_key="sk-test", + request_data={"model": "embed-english-v3.0", "texts": ["hello", "world"]}, + optional_params={}, + litellm_params={}, + ) + + +def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]: + assert response.usage is not None + return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) + + +def test_cohere_embedding_response_lists_every_embedding_type_in_provider_order(): + response = _transform( + { + "id": "embed-id", + "embeddings": {"float": FLOAT_VECTORS, "int8": INT8_VECTORS}, + "meta": {"billed_units": {"input_tokens": 5}}, + } + ) + + assert response.object == "list" + assert response.model == "embed-english-v3.0" + assert response.data == [ + {"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}, + {"object": "embedding", "index": 1, "embedding": [0.3, 0.4]}, + {"object": "embedding", "index": 0, "embedding": [1, -2]}, + {"object": "embedding", "index": 1, "embedding": [3, 4]}, + ] + assert _token_counts(response) == (5, 0, 5) + + +@pytest.mark.parametrize( + ("meta", "expected_counts", "expected_details"), + [ + ({"billed_units": {"input_tokens": 5}}, (5, 0, 5), (5, None)), + ({"billed_units": {"images": 3}}, (3, 0, 3), (None, 3)), + ({"billed_units": {"input_tokens": 5, "images": 3}}, (8, 0, 8), (5, 3)), + ({"billed_units": {"input_tokens": 0}}, (0, 0, 0), (0, None)), + ( + {"api_version": {"version": "2"}, "billed_units": {"input_tokens": 2, "output_tokens": 9}}, + (2, 0, 2), + (2, None), + ), + ], +) +def test_cohere_embedding_usage_comes_from_billed_units(meta, expected_counts, expected_details): + response = _transform({"embeddings": {"float": FLOAT_VECTORS}, "meta": meta}) + + assert response.usage is not None + details = response.usage.prompt_tokens_details + assert details is not None + assert _token_counts(response) == expected_counts + assert (details.text_tokens, details.image_tokens) == expected_details + + +@pytest.mark.parametrize( + "payload", + [ + {"embeddings": {"float": FLOAT_VECTORS}}, + {"embeddings": {"float": FLOAT_VECTORS}, "meta": {}}, + {"embeddings": {"float": FLOAT_VECTORS}, "meta": {"billed_units": {}}}, + {"embeddings": {"float": FLOAT_VECTORS}, "meta": {"billed_units": {"input_tokens": None, "images": None}}}, + ], +) +def test_cohere_embedding_usage_counts_input_tokens_when_billed_units_are_absent(payload): + response = _transform(payload) + + assert response.usage is not None + assert response.usage.prompt_tokens_details is None + assert _token_counts(response) == (2, 0, 2) + + +@pytest.mark.parametrize( + ("embeddings", "expected_vectors"), + [ + ({}, []), + ({"float": []}, []), + ({"float": [], "int8": INT8_VECTORS}, INT8_VECTORS), + ({"base64": ["AAAA", "BBBB"]}, ["AAAA", "BBBB"]), + ({"float": [[0.1], "x", None, 3]}, [[0.1], "x", None, 3]), + ], +) +def test_cohere_embedding_rows_pass_through_whatever_each_type_lists(embeddings, expected_vectors): + response = _transform({"embeddings": embeddings, "meta": {"billed_units": {"input_tokens": 1}}}) + + assert [row["embedding"] for row in response.data] == expected_vectors + + +@pytest.mark.parametrize("embeddings", [FLOAT_VECTORS, [], None, "vectors", 3]) +def test_cohere_embeddings_that_are_not_keyed_by_type_are_rejected(embeddings): + with pytest.raises(ValidationError): + _transform({"embeddings": embeddings, "meta": {"billed_units": {"input_tokens": 1}}}) + + +@pytest.mark.parametrize("vectors", [None, 3, False, 1.5]) +def test_cohere_embedding_type_without_a_list_of_vectors_is_rejected(vectors): + with pytest.raises(ValidationError): + _transform({"embeddings": {"float": vectors}, "meta": {"billed_units": {"input_tokens": 1}}}) + + +@pytest.mark.parametrize("meta", [None, [], "meta", 3]) +def test_cohere_embedding_meta_that_is_not_an_object_is_rejected(meta): + with pytest.raises(ValidationError): + _transform({"embeddings": {"float": FLOAT_VECTORS}, "meta": meta}) + + +@pytest.mark.parametrize("payload", [[1], "text", 3, True]) +def test_cohere_embedding_body_that_is_not_an_object_is_rejected(payload): + with pytest.raises(ValidationError): + _transform(payload) + + +def test_cohere_embedding_body_without_embeddings_is_rejected(): + with pytest.raises(KeyError): + _transform({"meta": {"billed_units": {"input_tokens": 1}}}) + + COHERE_V2_EMBED_URL: Final = "https://api.cohere.ai/v2/embed" diff --git a/tests/unit/llms/cometapi/embed/__init__.py b/tests/unit/llms/cometapi/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cometapi/embed/test_cometapi_embed_transformation.py b/tests/unit/llms/cometapi/embed/test_cometapi_embed_transformation.py new file mode 100644 index 00000000000..d59f6b4e153 --- /dev/null +++ b/tests/unit/llms/cometapi/embed/test_cometapi_embed_transformation.py @@ -0,0 +1,89 @@ +import datetime + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.cometapi.embed.transformation import CometAPIEmbeddingConfig +from litellm.types.utils import EmbeddingResponse + +EMBEDDING_ROW = {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]} + + +def _transform(payload: object) -> EmbeddingResponse: + logging_obj = Logging( + model="cometapi/text-embedding-3-small", + messages=[], + stream=False, + call_type="embedding", + start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-id", + function_id="function-id", + ) + return CometAPIEmbeddingConfig().transform_embedding_response( + model="text-embedding-3-small", + raw_response=httpx.Response(200, json=payload), + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + api_key="sk-test", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + +def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]: + assert response.usage is not None + return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) + + +def test_cometapi_embedding_response_carries_provider_fields_and_usage(): + response = _transform( + { + "object": "list", + "model": "text-embedding-3-small", + "data": [EMBEDDING_ROW], + "usage": {"prompt_tokens": 8, "total_tokens": 9}, + } + ) + + assert response.model == "text-embedding-3-small" + assert response.object == "list" + assert response.data == [EMBEDDING_ROW] + assert _token_counts(response) == (8, 0, 9) + + +@pytest.mark.parametrize( + ("payload", "expected"), + [ + ({"data": [EMBEDDING_ROW]}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 4}}, (4, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 6}}, (0, 0, 6)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": None, "total_tokens": None}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": "", "total_tokens": []}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": "12", "total_tokens": 12.0}}, (12, 0, 12)), + ( + {"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 3, "total_tokens": 5, "completion_tokens": 2}}, + (3, 0, 5), + ), + ], +) +def test_cometapi_embedding_usage_tolerates_missing_null_and_loosely_typed_counts(payload, expected): + response = _transform(payload) + + assert response.data == [EMBEDDING_ROW] + assert _token_counts(response) == expected + + +@pytest.mark.parametrize("usage", [None, [], "tokens", 3]) +def test_cometapi_embedding_usage_that_is_not_an_object_is_rejected(usage): + with pytest.raises(ValidationError): + _transform({"data": [EMBEDDING_ROW], "usage": usage}) + + +@pytest.mark.parametrize("usage", [{"prompt_tokens": "many"}, {"total_tokens": 1.5}, {"prompt_tokens": [1]}]) +def test_cometapi_embedding_token_count_that_is_not_an_integer_is_rejected(usage): + with pytest.raises(ValidationError): + _transform({"data": [EMBEDDING_ROW], "usage": usage}) diff --git a/tests/unit/llms/dataforseo/search/test_dataforseo_search_transformation.py b/tests/unit/llms/dataforseo/search/test_dataforseo_search_transformation.py index 0f20c5b1108..d94f6a53c89 100644 --- a/tests/unit/llms/dataforseo/search/test_dataforseo_search_transformation.py +++ b/tests/unit/llms/dataforseo/search/test_dataforseo_search_transformation.py @@ -2,13 +2,18 @@ Unit tests for DataForSEO Search functionality. """ +import datetime import os from unittest.mock import AsyncMock, patch +import httpx import pytest +from pydantic import ValidationError import litellm +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult +from litellm.llms.dataforseo.search.transformation import DataForSEOSearchConfig from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -54,6 +59,119 @@ async def test_dataforseo_search_basic(): assert len(response.results[0].snippet) > 0 +def _transform(payload: object) -> SearchResponse: + logging_obj = Logging( + model="dataforseo/search", + messages=[], + stream=False, + call_type="search", + start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-id", + function_id="function-id", + ) + return DataForSEOSearchConfig().transform_search_response( + raw_response=httpx.Response(200, json=payload), logging_obj=logging_obj + ) + + +def _successful_task(result: object) -> dict[str, object]: + return {"status_code": 20000, "status_message": "Ok.", "tasks": [{"status_code": 20000, "result": result}]} + + +def test_dataforseo_keeps_only_organic_items_of_the_first_result_in_provider_order(): + response = _transform( + _successful_task( + [ + { + "keyword": "litellm", + "items_count": 4, + "items": [ + { + "type": "organic", + "rank_absolute": 1, + "title": "LiteLLM", + "url": "https://example.com/litellm", + "description": "Call every LLM API", + "links": None, + }, + {"type": "people_also_ask", "items": [{"title": "What is LiteLLM?"}]}, + {"type": "paid", "title": "Ad", "url": "https://example.com/ad", "description": "Sponsored"}, + { + "type": "organic", + "rank_absolute": 4, + "title": "Docs", + "url": "https://example.com/docs", + "description": "Docs", + }, + ], + }, + "a second result entry is never read", + ] + ) + ) + + assert response.object == "search" + assert response.results == [ + SearchResult(title="LiteLLM", url="https://example.com/litellm", snippet="Call every LLM API"), + SearchResult(title="Docs", url="https://example.com/docs", snippet="Docs"), + ] + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"status_code": 20000, "tasks": []}, + {"tasks": [{"status_code": 40501, "status_message": "Invalid Field.", "result": None}]}, + {"tasks": [{"status_code": 20000}]}, + _successful_task([]), + _successful_task(""), + _successful_task([{"keyword": "litellm", "items_count": 0}]), + _successful_task([{"items": []}]), + _successful_task([{"items": ""}]), + _successful_task([{"items": {}}]), + _successful_task([{"items": [{"type": "paid", "title": "Ad"}, {"title": "no type"}]}]), + [], + "plain text body", + ], +) +def test_dataforseo_response_without_organic_items_is_empty(payload): + assert _transform(payload).results == [] + + +def test_dataforseo_organic_item_missing_optional_fields_defaults_to_empty_strings(): + assert _transform(_successful_task([{"items": [{"type": "organic"}]}])).results == [ + SearchResult(title="", url="", snippet="", date=None, last_updated=None) + ] + + +@pytest.mark.parametrize("field", ["title", "url", "description"]) +@pytest.mark.parametrize("value", [None, 7, ["text"]]) +def test_dataforseo_organic_item_with_non_string_field_is_rejected(field, value): + item = {"type": "organic", "title": "T", "url": "https://example.com", "description": "D", field: value} + + with pytest.raises(ValidationError): + _transform(_successful_task([{"items": [item]}])) + + +def test_dataforseo_organic_item_with_several_non_string_fields_reports_every_field(): + item = {"type": "organic", "title": 7, "url": None, "description": ["text"]} + + with pytest.raises(ValidationError) as exc_info: + _transform(_successful_task([{"items": [item]}])) + + assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",)] + + +def test_dataforseo_malformed_item_is_rejected_without_naming_its_position(): + items = [{}] * 429 + ["not-an-object"] + + with pytest.raises(ValidationError) as exc_info: + _transform(_successful_task([{"items": items}])) + + assert "429" not in str(exc_info.value) + + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/llms/duckduckgo/search/test_duckduckgo_search_transformation.py b/tests/unit/llms/duckduckgo/search/test_duckduckgo_search_transformation.py index 68419891166..484a724ad7c 100644 --- a/tests/unit/llms/duckduckgo/search/test_duckduckgo_search_transformation.py +++ b/tests/unit/llms/duckduckgo/search/test_duckduckgo_search_transformation.py @@ -5,8 +5,11 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx +from pydantic import ValidationError import litellm +from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult +from litellm.llms.duckduckgo.search.transformation import DuckDuckGoSearchConfig class TestDuckDuckGoSearchMocked: @@ -297,3 +300,133 @@ def test_duckduckgo_sync_search_returns_typed_results_without_a_limit(respx_mock for result in response.results ) assert response.results[-1].url == "https://example.com/9" + + +def _transform(payload: object, query: str = "q=python&format=json") -> SearchResponse: + request = httpx.Request("GET", f"https://api.duckduckgo.com/?{query}") + return DuckDuckGoSearchConfig().transform_search_response( + raw_response=httpx.Response(200, json=payload, request=request), logging_obj=None + ) + + +def _topic(name: str) -> dict[str, object]: + return { + "FirstURL": f"https://duckduckgo.com/{name}", + "Icon": {"Height": "", "URL": "", "Width": ""}, + "Result": f'{name}', + "Text": f"{name} - About {name}.", + } + + +def test_duckduckgo_abstract_comes_first_then_related_and_nested_topics_in_provider_order(): + long_text = "word " * 12 + response = _transform( + { + "Abstract": "", + "AbstractSource": "Wikipedia", + "AbstractText": "Python is a high-level programming language.", + "AbstractURL": "https://en.wikipedia.org/wiki/Python", + "Answer": "", + "Heading": "Python", + "ImageHeight": 0, + "Infobox": "", + "RelatedTopics": [ + _topic("Python_programming"), + {"FirstURL": "https://duckduckgo.com/Long", "Text": long_text}, + {"Name": "Related Topics", "Topics": [_topic("Indus")]}, + ], + "Results": [], + "Type": "A", + "meta": {"attribution": None, "src_options": {"is_fanon": 0}}, + } + ) + + assert response.object == "search" + assert response.results == [ + SearchResult( + title="Python", + url="https://en.wikipedia.org/wiki/Python", + snippet="Python is a high-level programming language.", + ), + SearchResult( + title="Python_programming", + url="https://duckduckgo.com/Python_programming", + snippet="About Python_programming.", + ), + SearchResult(title=long_text[:50] + "...", url="https://duckduckgo.com/Long", snippet=long_text), + SearchResult(title="Indus", url="https://duckduckgo.com/Indus", snippet="About Indus."), + ] + + +def test_duckduckgo_abstract_without_heading_gets_an_empty_title(): + assert _transform({"AbstractURL": "https://example.com", "AbstractText": "Summary"}).results == [ + SearchResult(title="", url="https://example.com", snippet="Summary") + ] + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"Abstract": "", "AbstractText": "", "AbstractURL": "", "Heading": "", "RelatedTopics": [], "Results": []}, + {"AbstractURL": "https://example.com", "AbstractText": "", "Heading": "Disambiguation"}, + {"AbstractURL": "", "AbstractText": "Summary without a link", "Heading": None}, + {"RelatedTopics": ""}, + {"RelatedTopics": {}}, + {"RelatedTopics": ["not a topic", None, 3, ["nested"]]}, + {"RelatedTopics": [{"Name": "Group without topics"}, {"FirstURL": "https://example.com"}, {"Text": "no url"}]}, + {"RelatedTopics": [{"Name": "Empty group", "Topics": []}]}, + ], +) +def test_duckduckgo_response_without_abstract_or_linked_topics_is_empty(payload): + assert _transform(payload).results == [] + + +def test_duckduckgo_non_object_related_topics_are_skipped_and_the_rest_are_kept(): + response = _transform({"RelatedTopics": ["not a topic", _topic("First"), None, _topic("Second")]}) + + assert [result.title for result in response.results] == ["First", "Second"] + + +@pytest.mark.parametrize( + ("query", "expected_titles"), + [ + ("q=python&_max_results=1", ["Python"]), + ("q=python&_max_results=2", ["Python", "First"]), + ("q=python&_max_results=3", ["Python", "First", "Nested"]), + ("q=python&_max_results=not-a-number", ["Python", "First", "Nested", "Second"]), + ("q=python", ["Python", "First", "Nested", "Second"]), + ], +) +def test_duckduckgo_max_results_query_param_limits_the_results_including_the_abstract(query, expected_titles): + response = _transform( + { + "AbstractURL": "https://en.wikipedia.org/wiki/Python", + "AbstractText": "Python is a language.", + "Heading": "Python", + "RelatedTopics": [_topic("First"), {"Topics": [_topic("Nested")]}, _topic("Second")], + }, + query=query, + ) + + assert [result.title for result in response.results] == expected_titles + + +@pytest.mark.parametrize( + ("field", "value"), + [("Heading", None), ("Heading", 7), ("AbstractURL", 7), ("AbstractText", ["Summary"])], +) +def test_duckduckgo_abstract_with_non_string_field_is_rejected(field, value): + payload = {"AbstractURL": "https://example.com", "AbstractText": "Summary", "Heading": "Title", field: value} + + with pytest.raises(ValidationError): + _transform(payload) + + +def test_duckduckgo_abstract_with_several_non_string_fields_reports_every_field(): + payload = {"Heading": 7, "AbstractURL": ["https://example.com"], "AbstractText": {"text": "Summary"}} + + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",)] diff --git a/tests/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index 6b556ff241f..abcd33869b7 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -5,8 +5,10 @@ from unittest.mock import Mock import httpx import pytest import respx +from pydantic import ValidationError import litellm +from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult from litellm.llms.exa_ai.search.transformation import ExaAISearchConfig EXA_SEARCH_URL: Final = "https://api.exa.ai/search" @@ -88,3 +90,110 @@ def test_search_sends_max_results_as_num_results(respx_mock: respx.MockRouter) - } assert route.calls.last.request.headers["x-api-key"] == "test-exa-key" assert [result.url for result in response.results] == ["https://example.com/ml"] + + +def _transform(payload: object) -> SearchResponse: + return ExaAISearchConfig().transform_search_response(httpx.Response(200, json=payload), logging_obj=Mock()) + + +def test_transform_search_response_maps_every_exa_result_in_provider_order(): + response: Final = _transform( + { + "requestId": "b5947044c4b78efa9552a7c89b306d95", + "resolvedSearchType": "neural", + "results": [ + { + "id": "https://example.com/first", + "title": "First", + "url": "https://example.com/first", + "publishedDate": "2024-01-05T00:00:00.000Z", + "author": "Ada", + "score": 0.42, + "text": "First text", + "highlights": ["ignored when text is present"], + "highlightScores": [0.5], + "image": None, + }, + {"title": "Second", "url": "https://example.com/second", "highlights": ["one", "two"]}, + {"title": "Third", "url": "https://example.com/third", "summary": "Summary", "publishedDate": None}, + ], + "costDollars": {"total": 0.005}, + } + ) + + assert response.object == "search" + assert response.results == [ + SearchResult( + title="First", url="https://example.com/first", snippet="First text", date="2024-01-05T00:00:00.000Z" + ), + SearchResult(title="Second", url="https://example.com/second", snippet="one\n\ntwo"), + SearchResult(title="Third", url="https://example.com/third", snippet="Summary"), + ] + + +def test_transform_search_response_result_without_optional_fields_gets_empty_defaults(): + assert _transform({"results": [{}]}).results == [SearchResult(title="", url="", snippet="")] + + +@pytest.mark.parametrize( + ("highlights", "expected_snippet"), + [ + (None, "a summary"), + ("", "a summary"), + ({}, "a summary"), + ([""], "a summary"), + (["", ""], "\n\n"), + ("abc", "a\n\nb\n\nc"), + ({"first": 1, "second": 2}, "first\n\nsecond"), + ], +) +def test_transform_search_response_joins_whatever_the_highlights_value_iterates_over( + highlights: object, expected_snippet: str +): + response: Final = _transform( + { + "results": [ + {"title": "Title", "url": "https://example.com", "highlights": highlights, "summary": "a summary"} + ] + } + ) + + assert response.results[0].snippet == expected_snippet + + +@pytest.mark.parametrize("payload", [{}, {"results": []}, {"results": ""}, {"results": {}}, {"requestId": "abc"}]) +def test_transform_search_response_without_results_is_empty(payload: dict[str, object]): + assert _transform(payload).results == [] + + +@pytest.mark.parametrize( + "result", + [ + {"title": None, "url": "https://example.com"}, + {"title": "Title", "url": 7}, + {"title": "Title", "url": "https://example.com", "text": 7}, + {"title": "Title", "url": "https://example.com", "summary": ["a summary"]}, + {"title": "Title", "url": "https://example.com", "publishedDate": 1704412800}, + ], +) +def test_transform_search_response_rejects_result_with_wrongly_typed_field(result: dict[str, object]): + with pytest.raises(ValidationError): + _transform({"results": [result]}) + + +def test_transform_search_response_reports_every_wrongly_typed_field_of_a_result(): + result: Final = {"title": 7, "url": None, "text": ["text"], "publishedDate": 1704412800} + + with pytest.raises(ValidationError) as exc_info: + _transform({"results": [result]}) + + assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",), ("date",)] + + +def test_transform_search_response_rejects_malformed_result_without_naming_its_position(): + results = [{}] * 429 + ["not-an-object"] + + with pytest.raises(ValidationError) as exc_info: + _transform({"results": results}) + + assert "429" not in str(exc_info.value) diff --git a/tests/unit/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py b/tests/unit/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py index 6c55760b625..11078f64a9f 100644 --- a/tests/unit/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py +++ b/tests/unit/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py @@ -6,7 +6,9 @@ from pathlib import Path import httpx import pytest +from pydantic import ValidationError +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.fal_ai.image_edit import FalAIImageEditConfig from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.router import GenericLiteLLMParams @@ -156,3 +158,69 @@ def test_transform_request_requires_an_image(image): litellm_params=GenericLiteLLMParams(), headers={}, ) + + +def _transform_response(payload: object) -> ImageResponse: + return FalAIImageEditConfig().transform_image_edit_response( + model="openai/gpt-image-2.5/flare/edit", raw_response=httpx.Response(200, json=payload), logging_obj=None + ) + + +def test_transform_response_maps_image_objects_and_bare_urls_in_provider_order(): + response = _transform_response( + { + "images": [ + "https://fal.media/a.png", + {"b64_json": "QUJD", "width": 0, "height": "768", "content_type": 5}, + 7, + None, + {"url": "https://fal.media/b.png", "width": 512}, + ], + "seed": 42, + } + ) + + assert [(image.url, image.b64_json, image.provider_specific_fields) for image in response.data] == [ + ("https://fal.media/a.png", None, None), + (None, "QUJD", None), + ("https://fal.media/b.png", None, {"width": 512}), + ] + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"detail": "rate limit exceeded"}, + {"images": None}, + {"images": "https://fal.media/a.png"}, + {"images": {"url": "https://fal.media/a.png"}}, + {"images": 429}, + ], +) +def test_transform_response_without_an_image_list_is_empty(payload): + assert _transform_response(payload).data == [] + + +@pytest.mark.parametrize( + "payload", [7, True, "https://fal.media/a.png", [{"url": "https://fal.media/a.png"}], [{"images": []}]] +) +def test_transform_response_rejects_non_object_bodies_without_echoing_them(payload): + with pytest.raises(ValidationError) as exc_info: + _transform_response(payload) + + assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [("dict_type", ())] + assert "input_value" not in str(exc_info.value) + assert "fal.media" not in str(exc_info.value) + + +def test_transform_response_reports_a_body_that_is_not_json_as_a_provider_error(): + with pytest.raises(BaseLLMException) as exc_info: + FalAIImageEditConfig().transform_image_edit_response( + model="openai/gpt-image-2.5/flare/edit", + raw_response=httpx.Response(502, content=b"502 Bad Gateway"), + logging_obj=None, + ) + + assert exc_info.value.status_code == 502 + assert str(exc_info.value).startswith("Error parsing Fal AI image edit response: ") diff --git a/tests/unit/llms/fal_ai/image_generation/test_bria_transformation.py b/tests/unit/llms/fal_ai/image_generation/test_bria_transformation.py new file mode 100644 index 00000000000..76e9b514805 --- /dev/null +++ b/tests/unit/llms/fal_ai/image_generation/test_bria_transformation.py @@ -0,0 +1,93 @@ +import httpx +import pytest +from pydantic import ValidationError + +from litellm.llms.fal_ai.image_generation.bria_transformation import FalAIBriaConfig +from litellm.types.utils import ImageObject, ImageResponse + + +def _transform(payload: object, model_response: ImageResponse | None = None) -> ImageResponse: + return FalAIBriaConfig().transform_image_generation_response( + model="bria/text-to-image/3.2", + raw_response=httpx.Response(200, json=payload), + model_response=model_response or ImageResponse(), + logging_obj=None, + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_transform_image_generation_response_maps_the_single_image(): + response = _transform( + { + "image": { + "url": "https://fal.media/a.png", + "content_type": "image/png", + "file_name": "a.png", + "file_size": 123456, + "width": 1024, + "height": 1024, + }, + "images": [{"url": "https://fal.media/ignored.png"}], + } + ) + + assert [(image.url, image.b64_json, image.provider_specific_fields) for image in response.data] == [ + ("https://fal.media/a.png", None, None) + ] + + +def test_transform_image_generation_response_keeps_an_image_without_a_url(): + response = _transform({"image": {"content_type": "image/png", "url": None}}) + + assert [(image.url, image.b64_json) for image in response.data] == [(None, None)] + + +def test_transform_image_generation_response_appends_to_images_already_on_the_response(): + model_response = ImageResponse(data=[ImageObject(url="https://fal.media/existing.png")]) + + response = _transform({"image": {"url": "https://fal.media/new.png"}}, model_response) + + assert response is model_response + assert [image.url for image in response.data] == ["https://fal.media/existing.png", "https://fal.media/new.png"] + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"detail": "rate limit exceeded"}, + {"image": None}, + {"image": {}}, + {"image": "https://fal.media/a.png"}, + {"image": ["https://fal.media/a.png"]}, + {"image": 429}, + {"images": [{"url": "https://fal.media/a.png"}]}, + ], +) +def test_transform_image_generation_response_without_an_image_object_is_empty(payload: dict[str, object]): + assert _transform(payload).data == [] + + +@pytest.mark.parametrize( + "payload", [7, True, "https://fal.media/a.png", [{"image": {"url": "https://fal.media/a.png"}}]] +) +def test_transform_image_generation_response_rejects_non_object_bodies_without_echoing_them(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [("dict_type", ())] + assert "input_value" not in str(exc_info.value) + assert "fal.media" not in str(exc_info.value) + + +@pytest.mark.parametrize("url", [7, ["Request timed out"], {"code": 429}]) +def test_transform_image_generation_response_reports_a_url_of_the_wrong_type_with_its_value(url: object): + with pytest.raises(ValidationError) as exc_info: + _transform({"image": {"url": url}}) + + assert [(error["type"], error["loc"], error["input"]) for error in exc_info.value.errors()] == [ + ("string_type", ("url",), url) + ] diff --git a/tests/unit/llms/fal_ai/image_generation/test_fal_ai_image_generation_transformation.py b/tests/unit/llms/fal_ai/image_generation/test_fal_ai_image_generation_transformation.py new file mode 100644 index 00000000000..a7313c32c57 --- /dev/null +++ b/tests/unit/llms/fal_ai/image_generation/test_fal_ai_image_generation_transformation.py @@ -0,0 +1,97 @@ +import httpx +import pytest +from pydantic import ValidationError + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.fal_ai.image_generation.transformation import FalAIImageGenerationConfig +from litellm.types.utils import ImageObject, ImageResponse + + +def _transform(payload: object, model_response: ImageResponse | None = None) -> ImageResponse: + return FalAIImageGenerationConfig().transform_image_generation_response( + model="fal-ai/some-new-model", + raw_response=httpx.Response(200, json=payload), + model_response=model_response or ImageResponse(), + logging_obj=None, + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_transform_image_generation_response_maps_image_objects_and_bare_urls(): + response = _transform( + { + "images": [ + {"url": "https://fal.media/a.png", "width": 1024, "height": 768, "content_type": "image/png"}, + "https://fal.media/b.png", + {"b64_json": "QUJD", "width": 0, "height": "768", "content_type": 5}, + 7, + None, + ["https://fal.media/ignored.png"], + ], + "seed": 42, + "timings": {"inference": 2.5}, + } + ) + + assert [(image.url, image.b64_json, image.provider_specific_fields) for image in response.data] == [ + ("https://fal.media/a.png", None, {"width": 1024, "height": 768, "content_type": "image/png"}), + ("https://fal.media/b.png", None, None), + (None, "QUJD", None), + ] + + +def test_transform_image_generation_response_appends_to_images_already_on_the_response(): + model_response = ImageResponse(data=[ImageObject(url="https://fal.media/existing.png")]) + + response = _transform({"images": ["https://fal.media/new.png"]}, model_response) + + assert response is model_response + assert [image.url for image in response.data] == ["https://fal.media/existing.png", "https://fal.media/new.png"] + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"detail": "rate limit exceeded"}, + {"images": None}, + {"images": "https://fal.media/a.png"}, + {"images": {"url": "https://fal.media/a.png"}}, + {"images": 429}, + {"image": {"url": "https://fal.media/a.png"}}, + ], +) +def test_transform_image_generation_response_without_an_image_list_is_empty(payload: dict[str, object]): + assert _transform(payload).data == [] + + +@pytest.mark.parametrize( + "payload", [7, True, "https://fal.media/a.png", [{"url": "https://fal.media/a.png"}], [{"images": []}]] +) +def test_transform_image_generation_response_rejects_non_object_bodies_without_echoing_them(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [("dict_type", ())] + assert "input_value" not in str(exc_info.value) + assert "fal.media" not in str(exc_info.value) + + +def test_transform_image_generation_response_reports_a_body_that_is_not_json_as_a_provider_error(): + with pytest.raises(BaseLLMException) as exc_info: + FalAIImageGenerationConfig().transform_image_generation_response( + model="fal-ai/some-new-model", + raw_response=httpx.Response(502, content=b"502 Bad Gateway"), + model_response=ImageResponse(), + logging_obj=None, + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert exc_info.value.status_code == 502 + assert str(exc_info.value).startswith("Error transforming image generation response: ") diff --git a/tests/unit/llms/fal_ai/image_generation/test_imagen4_transformation.py b/tests/unit/llms/fal_ai/image_generation/test_imagen4_transformation.py new file mode 100644 index 00000000000..4866016a62d --- /dev/null +++ b/tests/unit/llms/fal_ai/image_generation/test_imagen4_transformation.py @@ -0,0 +1,112 @@ +import httpx +import pytest +from pydantic import ValidationError + +from litellm.llms.fal_ai.image_generation.imagen4_transformation import FalAIImagen4Config +from litellm.types.utils import ImageObject, ImageResponse + + +def _transform(payload: object, model_response: ImageResponse | None = None) -> ImageResponse: + return FalAIImagen4Config().transform_image_generation_response( + model="fal-ai/imagen4/preview", + raw_response=httpx.Response(200, json=payload), + model_response=model_response or ImageResponse(), + logging_obj=None, + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_transform_image_generation_response_maps_image_objects_bare_urls_and_seed(): + response = _transform( + { + "images": [ + { + "url": "https://fal.media/a.png", + "content_type": "image/png", + "file_name": "z9RV14K95DvU.png", + "file_size": 4404019, + }, + "https://fal.media/b.png", + {"content_type": "image/png"}, + 7, + None, + ["https://fal.media/ignored.png"], + ], + "seed": 42, + } + ) + + assert [(image.url, image.b64_json, image.provider_specific_fields) for image in response.data] == [ + ("https://fal.media/a.png", None, None), + ("https://fal.media/b.png", None, None), + (None, None, None), + ] + assert repr(response._hidden_params["seed"]) == "42" + + +@pytest.mark.parametrize("seed", [0, None, "42", 1.0, [1, 2], {"value": 42}]) +def test_transform_image_generation_response_stores_the_seed_exactly_as_received(seed: object): + response = _transform({"images": [], "seed": seed}) + + assert repr(response._hidden_params["seed"]) == repr(seed) + + +def test_transform_image_generation_response_appends_to_images_already_on_the_response(): + model_response = ImageResponse(data=[ImageObject(url="https://fal.media/existing.png")]) + + response = _transform({"images": ["https://fal.media/new.png"]}, model_response) + + assert response is model_response + assert [image.url for image in response.data] == ["https://fal.media/existing.png", "https://fal.media/new.png"] + assert "seed" not in response._hidden_params + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"detail": "rate limit exceeded"}, + {"images": None}, + {"images": "https://fal.media/a.png"}, + {"images": {"url": "https://fal.media/a.png"}}, + {"images": 429}, + {"image": {"url": "https://fal.media/a.png"}}, + ], +) +def test_transform_image_generation_response_without_an_image_list_is_empty(payload: dict[str, object]): + response = _transform(payload) + + assert response.data == [] + assert "seed" not in response._hidden_params + + +@pytest.mark.parametrize( + "payload", [7, True, "https://fal.media/a.png", [{"url": "https://fal.media/a.png"}], [{"images": [], "seed": 1}]] +) +def test_transform_image_generation_response_rejects_non_object_bodies_without_echoing_them(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [("dict_type", ())] + assert "input_value" not in str(exc_info.value) + assert "fal.media" not in str(exc_info.value) + + +@pytest.mark.parametrize("position", [0, 403]) +def test_transform_image_generation_response_keeps_earlier_images_when_a_url_has_the_wrong_type(position: int): + model_response = ImageResponse() + images = [{"url": f"https://fal.media/{index}.png"} for index in range(position)] + [{"url": 7}] + + with pytest.raises(ValidationError) as exc_info: + _transform({"images": images, "seed": 42}, model_response) + + assert [(error["type"], error["loc"], error["input"]) for error in exc_info.value.errors()] == [ + ("string_type", ("url",), 7) + ] + assert [image.url for image in model_response.data] == [ + f"https://fal.media/{index}.png" for index in range(position) + ] + assert "seed" not in model_response._hidden_params diff --git a/tests/unit/llms/fastcrw/search/test_transformation.py b/tests/unit/llms/fastcrw/search/test_transformation.py index adf8fec087c..a8becec7840 100644 --- a/tests/unit/llms/fastcrw/search/test_transformation.py +++ b/tests/unit/llms/fastcrw/search/test_transformation.py @@ -1,9 +1,12 @@ import os from unittest.mock import Mock, patch +import httpx import pytest +from pydantic import ValidationError import litellm +from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult from litellm.llms.fastcrw.search.transformation import FastCRWSearchConfig @@ -180,3 +183,102 @@ def test_transform_search_response_non_list_data(): _resp({"success": True, "data": {"unexpected": "shape"}}), logging_obj=Mock() ) assert resp.results == [] + + +def _transform(payload: object) -> SearchResponse: + return _config().transform_search_response( + raw_response=httpx.Response(200, json=payload), logging_obj=None + ) + + +def _entry(index: int) -> dict[str, str]: + return { + "title": f"Title {index}", + "url": f"https://example.com/{index}", + "description": f"Description {index}", + } + + +def test_transform_search_response_keeps_provider_order_and_snippet_fallbacks(): + response = _transform( + { + "success": True, + "warning": None, + "data": [ + {**_entry(1), "markdown": "# Page 1", "metadata": {"statusCode": 200}}, + {**_entry(2), "markdown": ""}, + {**_entry(3), "markdown": None}, + {"title": "Title 4", "url": "https://example.com/4"}, + {"metadata": {"statusCode": 200}}, + ], + } + ) + + assert response.object == "search" + assert response.results == [ + SearchResult(title="Title 1", url="https://example.com/1", snippet="# Page 1"), + SearchResult( + title="Title 2", url="https://example.com/2", snippet="Description 2" + ), + SearchResult( + title="Title 3", url="https://example.com/3", snippet="Description 3" + ), + SearchResult(title="Title 4", url="https://example.com/4", snippet=""), + SearchResult(title="", url="", snippet="", date=None, last_updated=None), + ] + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"success": False, "error": "rate limit"}, + {"success": True, "data": None}, + {"success": True, "data": "rate limit"}, + {"success": True, "data": 429}, + ], +) +def test_transform_search_response_without_a_data_list_is_empty(payload): + assert _transform(payload).results == [] + + +@pytest.mark.parametrize("payload", [[_entry(1)], "upstream unavailable", True, 429]) +def test_transform_search_response_rejects_a_body_that_is_not_an_object(payload): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert [error["type"] for error in exc_info.value.errors()] == ["dict_type"] + assert "input_value" not in str(exc_info.value) + + +@pytest.mark.parametrize("position", [0, 1, 403]) +@pytest.mark.parametrize("entry", ["rate limit", None, 429, [_entry(1)]]) +def test_transform_search_response_rejects_an_entry_that_is_not_an_object( + position, entry +): + entries = [_entry(index) for index in range(position)] + [entry] + + with pytest.raises(ValidationError) as exc_info: + _transform({"success": True, "data": entries}) + + assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [ + ("dict_type", ()) + ] + assert "input_value" not in str(exc_info.value) + + +@pytest.mark.parametrize("position", [0, 403]) +def test_transform_search_response_reports_every_field_of_wrong_type(position): + entries = [_entry(index) for index in range(position)] + [ + {"title": 7, "url": None, "markdown": ["Request timed out"]} + ] + + with pytest.raises(ValidationError) as exc_info: + _transform({"success": True, "data": entries}) + + assert [(error["loc"], error["input"]) for error in exc_info.value.errors()] == [ + (("title",), 7), + (("url",), None), + (("snippet",), ["Request timed out"]), + ] + assert "Request timed out" in str(exc_info.value) diff --git a/tests/unit/llms/firecrawl/search/test_firecrawl_search_transformation.py b/tests/unit/llms/firecrawl/search/test_firecrawl_search_transformation.py new file mode 100644 index 00000000000..571466ac86f --- /dev/null +++ b/tests/unit/llms/firecrawl/search/test_firecrawl_search_transformation.py @@ -0,0 +1,160 @@ +import httpx +import pytest +from pydantic import ValidationError + +from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult +from litellm.llms.firecrawl.search.transformation import FirecrawlSearchConfig + + +def _transform(payload: object) -> SearchResponse: + return FirecrawlSearchConfig().transform_search_response( + raw_response=httpx.Response(200, json=payload), logging_obj=None + ) + + +def _entry(index: int) -> dict[str, str]: + return {"title": f"Title {index}", "url": f"https://example.com/{index}", "description": f"Description {index}"} + + +def _news(index: int) -> dict[str, str]: + return {"title": f"News {index}", "url": f"https://example.com/news/{index}", "snippet": f"Snippet {index}"} + + +def test_firecrawl_cloud_response_lists_web_results_before_news_results(): + response = _transform( + { + "success": True, + "data": { + "web": [ + {**_entry(1), "markdown": "# Page 1", "metadata": {"statusCode": 200}}, + {**_entry(2), "markdown": ""}, + {**_entry(3), "markdown": None, "snippet": "ignored for web results"}, + {"position": 4}, + ], + "images": [{"title": "Logo", "imageUrl": "https://example.com/logo.png"}], + "news": [ + {**_news(1), "date": "2024-03-05", "markdown": "# News 1"}, + {**_news(2), "date": None, "description": "ignored for news results"}, + {"title": "News 3", "url": "https://example.com/news/3", "description": "ignored for news results"}, + ], + }, + "creditsUsed": 2, + } + ) + + assert response.object == "search" + assert response.results == [ + SearchResult(title="Title 1", url="https://example.com/1", snippet="# Page 1"), + SearchResult(title="Title 2", url="https://example.com/2", snippet="Description 2"), + SearchResult(title="Title 3", url="https://example.com/3", snippet="Description 3"), + SearchResult(title="", url="", snippet=""), + SearchResult(title="News 1", url="https://example.com/news/1", snippet="# News 1", date="2024-03-05"), + SearchResult(title="News 2", url="https://example.com/news/2", snippet="Snippet 2"), + SearchResult(title="News 3", url="https://example.com/news/3", snippet="", date=None, last_updated=None), + ] + + +def test_firecrawl_self_hosted_response_maps_the_flat_data_list(): + response = _transform( + { + "success": True, + "data": [ + {**_entry(1), "markdown": "# Page 1"}, + {**_entry(2), "date": "2024-03-05"}, + {"metadata": {"statusCode": 200}}, + ], + } + ) + + assert response.results == [ + SearchResult(title="Title 1", url="https://example.com/1", snippet="# Page 1"), + SearchResult(title="Title 2", url="https://example.com/2", snippet="Description 2", date=None), + SearchResult(title="", url="", snippet=""), + ] + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"success": False, "error": "rate limit"}, + {"data": None}, + {"data": "rate limit"}, + {"data": 429}, + {"data": []}, + {"data": {}}, + {"data": {"images": [{"imageUrl": "https://example.com/logo.png"}]}}, + {"data": {"web": [], "news": []}}, + {"data": {"web": "", "news": {}}}, + ], +) +def test_firecrawl_response_without_web_or_news_results_is_empty(payload): + assert _transform(payload).results == [] + + +@pytest.mark.parametrize("payload", [[_entry(1)], "upstream unavailable", True, 429]) +def test_firecrawl_body_that_is_not_an_object_is_rejected_without_echoing_it(payload): + with pytest.raises(ValidationError) as exc_info: + _transform(payload) + + assert [error["type"] for error in exc_info.value.errors()] == ["dict_type"] + assert "input_value" not in str(exc_info.value) + + +@pytest.mark.parametrize("section", ["web", "news"]) +@pytest.mark.parametrize("value", [None, 429, True]) +def test_firecrawl_section_that_cannot_be_iterated_is_rejected_without_echoing_it(section, value): + with pytest.raises(ValidationError) as exc_info: + _transform({"data": {"web": [_entry(1)], "news": [_news(1)], section: value}}) + + assert [error["type"] for error in exc_info.value.errors()] == ["iterable_type"] + assert "input_value" not in str(exc_info.value) + + +def _payload_with_entries(shape: str, entries: list[object]) -> dict[str, object]: + if shape == "self-hosted": + return {"data": entries} + if shape == "web": + return {"data": {"web": entries}} + return {"data": {"web": [_entry(0)], "news": entries}} + + +@pytest.mark.parametrize("shape", ["self-hosted", "web", "news"]) +@pytest.mark.parametrize("position", [0, 403]) +@pytest.mark.parametrize("entry", ["rate limit", None, 429, [_entry(1)]]) +def test_firecrawl_entry_that_is_not_an_object_is_rejected_without_its_position(shape, position, entry): + entries = [_entry(index) for index in range(position)] + [entry] + + with pytest.raises(ValidationError) as exc_info: + _transform(_payload_with_entries(shape, entries)) + + assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [("dict_type", ())] + assert "input_value" not in str(exc_info.value) + + +@pytest.mark.parametrize("shape", ["self-hosted", "web", "news"]) +@pytest.mark.parametrize("position", [0, 403]) +def test_firecrawl_entry_with_several_fields_of_wrong_type_reports_every_field(shape, position): + entries = [_entry(index) for index in range(position)] + [ + {"title": 7, "url": None, "markdown": ["Request timed out"]} + ] + + with pytest.raises(ValidationError) as exc_info: + _transform(_payload_with_entries(shape, entries)) + + assert [(error["loc"], error["input"]) for error in exc_info.value.errors()] == [ + (("title",), 7), + (("url",), None), + (("snippet",), ["Request timed out"]), + ] + assert "Request timed out" in str(exc_info.value) + + +def test_firecrawl_news_entry_with_non_string_date_reports_the_date_next_to_the_other_fields(): + with pytest.raises(ValidationError) as exc_info: + _transform({"data": {"news": [{**_news(1), "title": 7, "date": 1709596800}]}}) + + assert [(error["loc"], error["input"]) for error in exc_info.value.errors()] == [ + (("title",), 7), + (("date",), 1709596800), + ] diff --git a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 7b25484ea90..16e2221dc84 100644 --- a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -5,10 +5,11 @@ from unittest.mock import MagicMock, patch import httpx import pytest import respx +from pydantic import ValidationError import litellm -from litellm.litellm_core_utils.get_supported_openai_params import get_supported_openai_params from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY +from litellm.litellm_core_utils.get_supported_openai_params import get_supported_openai_params from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id @@ -2011,6 +2012,121 @@ VISION_MODEL = next( ) +_CHAT_REPLY: Final = { + "id": "chat-1", + "object": "chat.completion", + "created": 1, + "model": "accounts/fireworks/models/llama-v3p1-8b-instruct", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "pong"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, +} +_FIREWORKS_BASE: Final = "https://api.fireworks.ai/inference/v1" +_MODELS_URL: Final = f"{_FIREWORKS_BASE}/accounts/acct/models" + + +def _complete_against_reply(reply: object) -> ModelResponse: + return litellm.completion( + model="fireworks_ai/accounts/fireworks/models/llama-v3p1-8b-instruct", + messages=[{"role": "user", "content": "ping"}], + api_key="fw-test-key", + client=HTTPHandler( + client=httpx.Client( + transport=httpx.MockTransport(lambda request: httpx.Response(200, json=reply, request=request)) + ) + ), + ) + + +def test_completion_keeps_unmodelled_reply_fields_and_nested_fireworks_metrics() -> None: + perf_metrics: Final = {"ttft": 0.1, "per_stage": [{"name": "prefill", "ms": None}]} + + response: Final = _complete_against_reply( + { + **_CHAT_REPLY, + "system_fingerprint": None, + "service_tier": "default", + "x-trace": {"spans": [1, 2]}, + "perf_metrics": perf_metrics, + "prompt_token_ids": [7, 8, 9], + } + ) + + assert response.model == "fireworks_ai/accounts/fireworks/models/llama-v3p1-8b-instruct" + assert response.choices[0].message.content == "pong" + assert response.model_dump()["service_tier"] == "default" + assert response.model_dump()["x-trace"] == {"spans": [1, 2]} + assert response._hidden_params["fireworks_perf_metrics"] == perf_metrics + assert response._hidden_params["fireworks_prompt_token_ids"] == [7, 8, 9] + + +@pytest.mark.parametrize( + "reply", + [pytest.param([_CHAT_REPLY], id="list"), pytest.param("pong", id="string"), pytest.param(7, id="number")], +) +def test_completion_with_non_object_reply_raises_api_connection_error(reply: object) -> None: + with pytest.raises(litellm.APIConnectionError): + _complete_against_reply(reply) + + +@pytest.mark.parametrize( + ("body", "expected"), + [ + pytest.param( + {"models": [{"name": "accounts/acct/models/a", "state": {"ready": [True]}}, {"name": "b"}], "totalSize": 2}, + ["fireworks_ai/accounts/acct/models/a", "fireworks_ai/b"], + id="listed-models", + ), + pytest.param({"models": []}, [], id="empty-list"), + pytest.param({"models": ""}, [], id="models-empty-string"), + pytest.param({"models": {}}, [], id="models-empty-object"), + pytest.param({"models": [{"name": ""}]}, ["fireworks_ai/"], id="empty-name"), + ], +) +def test_get_models_prefixes_each_listed_name( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, body: dict[str, object], expected: list[str] +) -> None: + monkeypatch.setenv("FIREWORKS_ACCOUNT_ID", "acct") + respx_mock.get(_MODELS_URL).respond(200, json=body) + + assert FireworksAIConfig().get_models(api_key="fw-test-key", api_base=_FIREWORKS_BASE) == expected + + +@pytest.mark.parametrize( + ("body", "missing_key"), + [ + pytest.param({"data": []}, "models", id="no-models-key"), + pytest.param({"models": [{"name": "a"}, {"id": "b"}]}, "name", id="entry-without-name"), + ], +) +def test_get_models_raises_key_error_for_missing_field( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, body: dict[str, object], missing_key: str +) -> None: + monkeypatch.setenv("FIREWORKS_ACCOUNT_ID", "acct") + respx_mock.get(_MODELS_URL).respond(200, json=body) + + with pytest.raises(KeyError, match=missing_key): + FireworksAIConfig().get_models(api_key="fw-test-key", api_base=_FIREWORKS_BASE) + + +@pytest.mark.parametrize( + "body", + [ + pytest.param([{"name": "a"}], id="top-level-list"), + pytest.param({"models": None}, id="models-null"), + pytest.param({"models": ["a"]}, id="entry-not-an-object"), + pytest.param({"models": [{"name": 7}]}, id="name-not-a-string"), + ], +) +def test_get_models_rejects_malformed_listing( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, body: object +) -> None: + monkeypatch.setenv("FIREWORKS_ACCOUNT_ID", "acct") + respx_mock.get(_MODELS_URL).respond(200, json=body) + + with pytest.raises(ValidationError): + FireworksAIConfig().get_models(api_key="fw-test-key", api_base=_FIREWORKS_BASE) + + def _fireworks_chat_client() -> MagicMock: body: Final = { "id": "chat-user-attribution", diff --git a/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py b/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py index 2a69926414a..138aafcbfcd 100644 --- a/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py +++ b/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py @@ -460,6 +460,8 @@ def test_imagen_generation_response_keeps_predictions_without_image_bytes(): ("gemini/imagen-4.0-generate-001", {"predictions": [{"bytesBase64Encoded": "a"}, "not an object"]}), ("gemini-3.1-flash-image-preview", {"usageMetadata": "not an object"}), ("gemini-3.1-flash-image-preview", {"usageMetadata": ["not an object"]}), + ("gemini/imagen-4.0-generate-001", ["not an object"]), + ("gemini/imagen-4.0-generate-001", "not an object"), ], ) def test_image_generation_response_rejects_malformed_payloads_without_echoing_them(model: str, payload: object): @@ -467,3 +469,117 @@ def test_image_generation_response_rejects_malformed_payloads_without_echoing_th _transform_response(model, payload) assert "input_value" not in str(exc_info.value) + + +_GEMINI_IMAGE_MODEL = "gemini-3.1-flash-image-preview" + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"candidates": []}, + {"candidates": ""}, + {"candidates": {}}, + {"candidates": [{}]}, + {"candidates": [{"finishReason": "SAFETY"}]}, + {"candidates": [{"content": {}}]}, + {"candidates": [{"content": {"parts": ""}}]}, + {"candidates": [{"content": {"parts": "plain text"}}]}, + {"candidates": [{"content": {"parts": {"text": "only text"}}}]}, + {"candidates": [{"content": {"parts": [{"text": "only text"}, "plain text", ["text"], [], ""]}}]}, + {"candidates": [{"content": {"parts": [{"inlineData": {"mimeType": "image/png"}}]}}]}, + {"candidates": [{"content": {"parts": [{"inlineData": "no image"}, {"inlineData": ["mimeType"]}]}}]}, + {"candidates": [{"content": {"parts": [{"inlineData": ""}, {"inlineData": []}, {"inlineData": {}}]}}]}, + {"candidates": [{"content": {"parts": [{"inline_data": {"data": "snake-case-is-not-read"}}]}}]}, + ], +) +def test_gemini_image_generation_response_without_inline_image_data_has_no_images(payload: object): + assert _transform_response(_GEMINI_IMAGE_MODEL, payload).data == [] + + +def test_gemini_image_generation_response_keeps_image_order_and_thought_signatures(): + result = _transform_response( + _GEMINI_IMAGE_MODEL, + { + "candidates": [ + { + "content": { + "parts": [ + {"text": "caption"}, + {"inlineData": {"mimeType": "image/png", "data": "first"}, "thoughtSignature": "sig-1"}, + "plain text", + ] + } + }, + { + "content": { + "parts": [ + {"inlineData": {"data": "second"}}, + {"inlineData": {"data": None}, "thoughtSignature": ""}, + {"inlineData": {"data": "fourth"}, "thoughtSignature": {"nested": ["sig"]}}, + ] + } + }, + ] + }, + ) + + assert [(image.b64_json, image.provider_specific_fields) for image in result.data or []] == [ + ("first", {"thought_signature": "sig-1"}), + ("second", None), + (None, None), + ("fourth", {"thought_signature": {"nested": ["sig"]}}), + ] + + +@pytest.mark.parametrize( + "payload", + [ + ["gemini-secret"], + "gemini-secret", + {"candidates": None}, + {"candidates": 7}, + {"candidates": "gemini-secret"}, + {"candidates": ["gemini-secret"]}, + {"candidates": [{"content": None}]}, + {"candidates": [{"content": "gemini-secret"}]}, + {"candidates": [{"content": ["gemini-secret"]}]}, + {"candidates": [{"content": {"parts": None}}]}, + {"candidates": [{"content": {"parts": [None]}}]}, + {"candidates": [{"content": {"parts": [7]}}]}, + {"candidates": [{"content": {"parts": ["inlineData gemini-secret"]}}]}, + {"candidates": [{"content": {"parts": [["inlineData", "gemini-secret"]]}}]}, + {"candidates": [{"content": {"parts": {"inlineData": "gemini-secret"}}}]}, + {"candidates": [{"content": {"parts": [{"inlineData": None}]}}]}, + {"candidates": [{"content": {"parts": [{"inlineData": 7}]}}]}, + {"candidates": [{"content": {"parts": [{"inlineData": "data gemini-secret"}]}}]}, + {"candidates": [{"content": {"parts": [{"inlineData": ["data", "gemini-secret"]}]}}]}, + ], +) +def test_gemini_image_generation_response_rejects_malformed_candidates_without_echoing_them(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform_response(_GEMINI_IMAGE_MODEL, payload) + + assert "gemini-secret" not in str(exc_info.value) + + +def _rejection_text(payload: object) -> str: + with pytest.raises(ValidationError) as exc_info: + _transform_response(_GEMINI_IMAGE_MODEL, payload) + + return str(exc_info.value) + + +@pytest.mark.parametrize("position", [403, 429]) +@pytest.mark.parametrize("malformed_entry", [None, 7, "inlineData"]) +def test_gemini_image_generation_response_rejection_text_is_the_same_wherever_the_malformed_entry_sits( + position: int, malformed_entry: object +): + late_parts = [*[{"text": "caption"}] * position, malformed_entry] + late_candidates = [*[{"finishReason": "STOP"}] * position, malformed_entry] + + assert _rejection_text({"candidates": [{"content": {"parts": late_parts}}]}) == _rejection_text( + {"candidates": [{"content": {"parts": [malformed_entry]}}]} + ) + assert _rejection_text({"candidates": late_candidates}) == _rejection_text({"candidates": [malformed_entry]}) diff --git a/tests/unit/llms/google_pse/__init__.py b/tests/unit/llms/google_pse/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/google_pse/search/__init__.py b/tests/unit/llms/google_pse/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/google_pse/search/test_google_pse_search_transformation.py b/tests/unit/llms/google_pse/search/test_google_pse_search_transformation.py new file mode 100644 index 00000000000..9234e46d5b7 --- /dev/null +++ b/tests/unit/llms/google_pse/search/test_google_pse_search_transformation.py @@ -0,0 +1,110 @@ +import datetime + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult +from litellm.llms.google_pse.search.transformation import GooglePSESearchConfig + + +def _transform(payload: object) -> SearchResponse: + logging_obj = Logging( + model="google_pse/search", + messages=[], + stream=False, + call_type="search", + start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-id", + function_id="function-id", + ) + return GooglePSESearchConfig().transform_search_response( + raw_response=httpx.Response(200, json=payload), logging_obj=logging_obj + ) + + +def test_google_pse_items_keep_provider_order_and_ignore_search_metadata(): + response = _transform( + { + "kind": "customsearch#search", + "url": {"type": "application/json", "template": "https://www.googleapis.com/customsearch/v1?q={q}"}, + "queries": {"request": [{"title": "Google Custom Search - litellm", "totalResults": "2", "count": 2}]}, + "searchInformation": {"searchTime": 0.21, "formattedSearchTime": "0.21", "totalResults": "2"}, + "items": [ + { + "kind": "customsearch#result", + "title": "LiteLLM", + "htmlTitle": "LiteLLM", + "link": "https://example.com/litellm", + "displayLink": "example.com", + "snippet": "Call every LLM API", + "htmlSnippet": "Call every LLM API", + "pagemap": {"metatags": [{"og:title": "LiteLLM"}]}, + }, + {"title": "Docs", "link": "https://example.com/docs", "snippet": "Docs"}, + ], + } + ) + + assert response.object == "search" + assert response.results == [ + SearchResult(title="LiteLLM", url="https://example.com/litellm", snippet="Call every LLM API"), + SearchResult(title="Docs", url="https://example.com/docs", snippet="Docs"), + ] + + +def test_google_pse_item_without_title_link_or_snippet_gets_empty_strings(): + assert _transform({"items": [{"kind": "customsearch#result"}]}).results == [ + SearchResult(title="", url="", snippet="", date=None, last_updated=None) + ] + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"kind": "customsearch#search", "searchInformation": {"totalResults": "0"}}, + {"items": []}, + {"items": ""}, + {"items": {}}, + ], +) +def test_google_pse_response_without_items_is_empty(payload): + assert _transform(payload).results == [] + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("title", None), + ("title", 7), + ("link", None), + ("link", ["https://example.com"]), + ("snippet", None), + ("snippet", {"text": "snippet"}), + ], +) +def test_google_pse_item_with_field_of_wrong_type_is_rejected(field, value): + item = {"title": "T", "link": "https://example.com", "snippet": "S", field: value} + + with pytest.raises(ValidationError): + _transform({"items": [item]}) + + +def test_google_pse_item_with_several_fields_of_wrong_type_reports_every_field(): + item = {"title": 7, "link": None, "snippet": ["snippet"]} + + with pytest.raises(ValidationError) as exc_info: + _transform({"items": [item]}) + + assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",)] + + +def test_google_pse_malformed_item_is_rejected_without_naming_its_position(): + items = [{}] * 429 + ["not-an-object"] + + with pytest.raises(ValidationError) as exc_info: + _transform({"items": items}) + + assert "429" not in str(exc_info.value) diff --git a/tests/unit/llms/infinity/embedding/__init__.py b/tests/unit/llms/infinity/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/infinity/embedding/test_infinity_embedding_transformation.py b/tests/unit/llms/infinity/embedding/test_infinity_embedding_transformation.py new file mode 100644 index 00000000000..e890360d725 --- /dev/null +++ b/tests/unit/llms/infinity/embedding/test_infinity_embedding_transformation.py @@ -0,0 +1,89 @@ +import datetime + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.infinity.embedding.transformation import InfinityEmbeddingConfig +from litellm.types.utils import EmbeddingResponse + +EMBEDDING_ROW = {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]} + + +def _transform(payload: object) -> EmbeddingResponse: + logging_obj = Logging( + model="infinity/bge-small", + messages=[], + stream=False, + call_type="embedding", + start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-id", + function_id="function-id", + ) + return InfinityEmbeddingConfig().transform_embedding_response( + model="bge-small", + raw_response=httpx.Response(200, json=payload), + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + api_key="sk-test", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + +def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]: + assert response.usage is not None + return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) + + +def test_infinity_embedding_response_carries_provider_fields_and_usage(): + response = _transform( + { + "object": "list", + "model": "bge-small", + "data": [EMBEDDING_ROW], + "usage": {"prompt_tokens": 8, "total_tokens": 9}, + } + ) + + assert response.model == "bge-small" + assert response.object == "list" + assert response.data == [EMBEDDING_ROW] + assert _token_counts(response) == (8, 0, 9) + + +@pytest.mark.parametrize( + ("payload", "expected"), + [ + ({"data": [EMBEDDING_ROW]}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 4}}, (4, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 6}}, (0, 0, 6)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": None, "total_tokens": None}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": "", "total_tokens": []}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": "12", "total_tokens": 12.0}}, (12, 0, 12)), + ( + {"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 3, "total_tokens": 5, "completion_tokens": 2}}, + (3, 0, 5), + ), + ], +) +def test_infinity_embedding_usage_tolerates_missing_null_and_loosely_typed_counts(payload, expected): + response = _transform(payload) + + assert response.data == [EMBEDDING_ROW] + assert _token_counts(response) == expected + + +@pytest.mark.parametrize("usage", [None, [], "tokens", 3]) +def test_infinity_embedding_usage_that_is_not_an_object_is_rejected(usage): + with pytest.raises(ValidationError): + _transform({"data": [EMBEDDING_ROW], "usage": usage}) + + +@pytest.mark.parametrize("usage", [{"prompt_tokens": "many"}, {"total_tokens": 1.5}, {"prompt_tokens": [1]}]) +def test_infinity_embedding_token_count_that_is_not_an_integer_is_rejected(usage): + with pytest.raises(ValidationError): + _transform({"data": [EMBEDDING_ROW], "usage": usage}) diff --git a/tests/unit/llms/langgraph/chat/test_langgraph_transformation.py b/tests/unit/llms/langgraph/chat/test_langgraph_transformation.py index 881da2ce15b..238f7dd2f26 100644 --- a/tests/unit/llms/langgraph/chat/test_langgraph_transformation.py +++ b/tests/unit/llms/langgraph/chat/test_langgraph_transformation.py @@ -20,11 +20,14 @@ Non-streaming: import asyncio import importlib +import json import os +import httpx import pytest import litellm +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -232,3 +235,36 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), "cohere_key": getattr(litellm, "cohere_key", None), } + + +@pytest.mark.asyncio +async def test_langgraph_streaming_acompletion_without_a_client_posts_through_the_cached_langgraph_client(monkeypatch): + sent_requests = [] + + async def events(): + yield b'data: ["messages", [[{"type": "ai", "content": "100"}, {}]]]\n\n' + yield b'data: ["metadata", {"run_id": "run-1"}]\n\n' + + def respond(request): + sent_requests.append(request) + return httpx.Response(200, content=events()) + + cached_client = get_async_httpx_client(llm_provider="langgraph", params={}) + monkeypatch.setattr(cached_client, "client", httpx.AsyncClient(transport=httpx.MockTransport(respond))) + + response = await litellm.acompletion( + model="langgraph/agent", + messages=[{"role": "user", "content": "What is 25 * 4?"}], + api_base="http://langgraph.test", + stream=True, + ) + streamed = [(chunk.choices[0].delta.content, chunk.choices[0].finish_reason) async for chunk in response] + + assert streamed == [("100", None), (None, "stop")] + assert [(request.method, str(request.url)) for request in sent_requests] == [ + ("POST", "http://langgraph.test/runs/stream") + ] + assert json.loads(sent_requests[0].content) == { + "assistant_id": "agent", + "input": {"messages": [{"role": "human", "content": "What is 25 * 4?"}]}, + } diff --git a/tests/unit/llms/lemonade/chat/__init__.py b/tests/unit/llms/lemonade/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/lemonade/chat/test_lemonade_chat_transformation.py b/tests/unit/llms/lemonade/chat/test_lemonade_chat_transformation.py new file mode 100644 index 00000000000..6f177d9bf8b --- /dev/null +++ b/tests/unit/llms/lemonade/chat/test_lemonade_chat_transformation.py @@ -0,0 +1,172 @@ +import pytest +import respx +from pydantic import ValidationError + +from litellm.llms.lemonade.chat.transformation import LemonadeChatConfig + +LEMONADE_BASE = "http://lemonade.test/api/v1" +MODEL = "Qwen3-4B-GGUF" + + +def _expected_model_info(**overrides: object) -> dict[str, object]: + return { + "key": f"lemonade/{MODEL}", + "litellm_provider": "lemonade", + "mode": "chat", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "max_tokens": None, + "max_input_tokens": None, + "max_output_tokens": None, + **overrides, + } + + +@pytest.mark.parametrize( + ("body", "expected"), + [ + pytest.param( + {"object": "list", "data": [{"id": MODEL, "recipe_options": {"ctx_size": [4096]}}, {"id": "nomic-embed"}]}, + [f"lemonade/{MODEL}", "lemonade/nomic-embed"], + id="listed-models", + ), + pytest.param({"data": []}, [], id="empty-list"), + pytest.param({"object": "list"}, [], id="data-key-missing"), + pytest.param({"data": ""}, [], id="data-empty-string"), + pytest.param({"data": {}}, [], id="data-empty-object"), + pytest.param({"data": [{"id": ""}]}, ["lemonade/"], id="empty-id"), + ], +) +def test_get_models_prefixes_each_listed_id( + respx_mock: respx.MockRouter, body: dict[str, object], expected: list[str] +) -> None: + respx_mock.get(f"{LEMONADE_BASE}/models").respond(200, json=body) + + assert LemonadeChatConfig().get_models(api_base=LEMONADE_BASE) == expected + + +def test_get_models_raises_key_error_for_entry_without_id(respx_mock: respx.MockRouter) -> None: + respx_mock.get(f"{LEMONADE_BASE}/models").respond(200, json={"data": [{"id": MODEL}, {"name": "no-id"}]}) + + with pytest.raises(KeyError, match="id"): + LemonadeChatConfig().get_models(api_base=LEMONADE_BASE) + + +@pytest.mark.parametrize( + "body", + [ + pytest.param([{"id": MODEL}], id="top-level-list"), + pytest.param({"data": None}, id="data-null"), + pytest.param({"data": [MODEL]}, id="entry-not-an-object"), + pytest.param({"data": [{"id": 7}]}, id="id-not-a-string"), + ], +) +def test_get_models_rejects_malformed_listing(respx_mock: respx.MockRouter, body: object) -> None: + respx_mock.get(f"{LEMONADE_BASE}/models").respond(200, json=body) + + with pytest.raises(ValidationError): + LemonadeChatConfig().get_models(api_base=LEMONADE_BASE) + + +@pytest.mark.parametrize( + ("body", "expected"), + [ + pytest.param({}, _expected_model_info(), id="empty-object"), + pytest.param( + {"max_tokens": "2048", "max_output_tokens": 4096, "max_input_tokens": 8192}, + _expected_model_info(max_tokens=2048, max_output_tokens=4096, max_input_tokens=8192), + id="numeric-string-and-ints", + ), + pytest.param( + {"max_output_tokens": 4096}, + _expected_model_info(max_tokens=4096, max_output_tokens=4096), + id="max-tokens-falls-back-to-output", + ), + pytest.param( + {"max_tokens": True, "max_output_tokens": 0, "max_input_tokens": 12.5}, + _expected_model_info(), + id="bool-zero-and-float-ignored", + ), + pytest.param( + {"max_tokens": None, "max_output_tokens": "many", "recipe_options": [1], "provider_specific_entry": "x"}, + _expected_model_info(provider_specific_entry={"recipe_options": [1]}), + id="null-text-and-non-object-options-ignored", + ), + pytest.param( + {"provider_specific_entry": {"recipe_options": {"ctx_size": "32768"}, "labels": ["hot"]}}, + _expected_model_info( + max_input_tokens=32768, + provider_specific_entry={"recipe_options": {"ctx_size": "32768"}, "labels": ["hot"]}, + ), + id="nested-entry-kept", + ), + ], +) +def test_get_model_info_reads_token_limits_from_server_payload( + respx_mock: respx.MockRouter, body: dict[str, object], expected: dict[str, object] +) -> None: + respx_mock.get(f"{LEMONADE_BASE}/models/{MODEL}").respond(200, json=body) + + assert LemonadeChatConfig().get_model_info(model=f"lemonade/{MODEL}", api_base=LEMONADE_BASE) == expected + + +@pytest.mark.parametrize("body", [pytest.param([{"max_tokens": 5}], id="list"), pytest.param("loaded", id="string")]) +def test_get_model_info_raises_for_non_object_payload_instead_of_defaulting( + respx_mock: respx.MockRouter, body: object +) -> None: + respx_mock.get(f"{LEMONADE_BASE}/models/{MODEL}").respond(200, json=body) + + with pytest.raises(ValidationError): + LemonadeChatConfig().get_model_info(model=f"lemonade/{MODEL}", api_base=LEMONADE_BASE) + + +def test_get_model_info_defaults_when_server_errors(respx_mock: respx.MockRouter) -> None: + respx_mock.get(f"{LEMONADE_BASE}/models/{MODEL}").respond(503, json={"max_tokens": 5}) + + assert ( + LemonadeChatConfig().get_model_info(model=f"lemonade/{MODEL}", api_base=LEMONADE_BASE) == _expected_model_info() + ) + + +@pytest.mark.parametrize( + ("supplied", "stored"), + [ + pytest.param({}, {}, id="nothing-supplied"), + pytest.param( + {"repeat_penalty": 0.0, "n": 0, "stop": "", "tools": []}, + {"repeat_penalty": 0.0, "n": 0, "stop": "", "tools": []}, + id="falsy-values-are-kept", + ), + pytest.param({"top_k": 40, "top_p": None, "max_tokens": None}, {"top_k": 40}, id="none-is-skipped"), + ], +) +def test_init_stores_each_supplied_value_except_none_as_class_level_config( + supplied: dict[str, object], stored: dict[str, object] +) -> None: + class ScopedLemonadeConfig(LemonadeChatConfig): + pass + + ScopedLemonadeConfig(**supplied) + + assert ScopedLemonadeConfig.get_config() == stored + + +def test_init_stores_the_supplied_object_itself_rather_than_a_copy() -> None: + class ScopedLemonadeConfig(LemonadeChatConfig): + pass + + logit_bias = {"50256": -100} + + ScopedLemonadeConfig(logit_bias=logit_bias) + + assert ScopedLemonadeConfig.get_config()["logit_bias"] is logit_bias + + +def test_init_called_again_with_none_keeps_the_value_stored_earlier() -> None: + class ScopedLemonadeConfig(LemonadeChatConfig): + pass + + ScopedLemonadeConfig(top_k=40) + ScopedLemonadeConfig(top_k=None, n=2) + + assert ScopedLemonadeConfig.get_config() == {"top_k": 40, "n": 2} diff --git a/tests/unit/llms/linkup/search/test_linkup_search_transformation.py b/tests/unit/llms/linkup/search/test_linkup_search_transformation.py index c98990b5a55..80967c6e539 100644 --- a/tests/unit/llms/linkup/search/test_linkup_search_transformation.py +++ b/tests/unit/llms/linkup/search/test_linkup_search_transformation.py @@ -2,12 +2,18 @@ Tests for Linkup Search API integration. """ +import datetime import os from unittest.mock import Mock, patch +import httpx import pytest +from pydantic import ValidationError import litellm +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult +from litellm.llms.linkup.search.transformation import LinkupSearchConfig from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -99,6 +105,100 @@ class TestLinkupSearchTransformation: assert "Microsoft Cloud revenue" in first_result.snippet +def _transform(payload: object) -> SearchResponse: + logging_obj = Logging( + model="linkup/search", + messages=[], + stream=False, + call_type="search", + start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-id", + function_id="function-id", + ) + return LinkupSearchConfig().transform_search_response( + raw_response=httpx.Response(200, json=payload), logging_obj=logging_obj + ) + + +def test_linkup_text_and_image_results_keep_provider_order_and_fields(): + response = _transform( + { + "results": [ + {"type": "text", "name": "Annual Report", "url": "https://example.com/report", "content": "Revenue up"}, + {"type": "image", "name": "Chart", "url": "https://example.com/chart.png", "content": "A chart"}, + {"type": "image", "url": "https://example.com/logo.png"}, + {"name": "Untyped", "url": "https://example.com/untyped", "content": "Defaults to text"}, + ] + } + ) + + assert response.results == [ + SearchResult(title="Annual Report", url="https://example.com/report", snippet="Revenue up"), + SearchResult(title="Chart", url="https://example.com/chart.png", snippet="A chart"), + SearchResult(title="https://example.com/logo.png", url="https://example.com/logo.png", snippet=""), + SearchResult(title="Untyped", url="https://example.com/untyped", snippet="Defaults to text"), + ] + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"answer": "A sourced answer without a results array"}, + {"results": []}, + {"results": ""}, + {"results": {}}, + {"results": [{"type": "video", "name": "Clip", "url": "https://example.com/clip"}]}, + ], +) +def test_linkup_response_without_text_or_image_results_is_empty(payload): + response = _transform(payload) + + assert response.object == "search" + assert response.results == [] + + +def test_linkup_result_missing_optional_fields_defaults_to_empty_strings(): + response = _transform( + {"results": [{"type": "text", "favicon": "https://example.com/favicon.ico"}, {"type": "image"}]} + ) + + assert response.results == [ + SearchResult(title="", url="", snippet=""), + SearchResult(title="", url="", snippet=""), + ] + + +@pytest.mark.parametrize( + "result", + [ + {"type": "text", "name": None, "url": "https://example.com", "content": "c"}, + {"type": "text", "name": "Title", "url": 7, "content": "c"}, + {"type": "image", "name": "Title", "url": "https://example.com", "content": None}, + ], +) +def test_linkup_result_with_non_string_field_is_rejected(result): + with pytest.raises(ValidationError): + _transform({"results": [result]}) + + +@pytest.mark.parametrize("result_type", ["text", "image"]) +def test_linkup_result_with_several_non_string_fields_reports_every_field(result_type): + with pytest.raises(ValidationError) as exc_info: + _transform({"results": [{"type": result_type, "name": 7, "url": None, "content": ["content"]}]}) + + assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",)] + + +def test_linkup_malformed_result_is_rejected_without_naming_its_position(): + results = [{}] * 429 + ["not-an-object"] + + with pytest.raises(ValidationError) as exc_info: + _transform({"results": results}) + + assert "429" not in str(exc_info.value) + + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py b/tests/unit/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py index d819a79cef1..87e2c165515 100644 --- a/tests/unit/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py +++ b/tests/unit/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py @@ -1,9 +1,11 @@ import base64 +import json from typing import Final from unittest.mock import MagicMock import httpx import pytest +from pydantic import ValidationError import litellm from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig @@ -195,3 +197,113 @@ def test_transform_response_invalid_base64_raises(): raw_response=raw_response, logging_obj=MagicMock(), ) + + +_SPEECH_AUDIO: Final = b"ID3-fake-mp3-bytes" +_SPEECH_AUDIO_B64: Final = base64.b64encode(_SPEECH_AUDIO).decode() + + +@pytest.mark.parametrize( + ("request_body", "expected_content_type"), + [ + pytest.param(b'{"input": "hi", "response_format": "opus"}', "audio/ogg", id="known-format"), + pytest.param(b'{"input": "hi", "response_format": "aac"}', "audio/mpeg", id="unknown-format"), + pytest.param(b'{"input": "hi", "response_format": null}', "audio/mpeg", id="null-format"), + pytest.param(b'{"input": "hi", "response_format": ["wav"]}', "audio/mpeg", id="non-string-format"), + pytest.param(b'{"input": "hi", "ref_audio": "bXktdm9pY2Utc2FtcGxl"}', "audio/mpeg", id="format-omitted"), + pytest.param(b"", "audio/mpeg", id="request-without-body"), + ], +) +def test_transform_response_labels_audio_with_the_requested_format(request_body: bytes, expected_content_type: str): + raw_response: Final = httpx.Response( + 200, + json={"id": None, "audio_data": _SPEECH_AUDIO_B64, "usage": {"characters": 2}, "voices": ["en_paul_neutral"]}, + request=httpx.Request("POST", SPEECH_URL, content=request_body), + ) + result: Final = MistralTextToSpeechConfig().transform_text_to_speech_response( + model="voxtral-mini-tts-2603", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + assert result.content == _SPEECH_AUDIO + assert result.response.headers["content-type"] == expected_content_type + + +@pytest.mark.parametrize( + ("payload", "expected_keys"), + [ + pytest.param({}, "()", id="empty-object"), + pytest.param({"audio_data": None}, "('audio_data',)", id="null-audio"), + pytest.param({"audio_data": ""}, "('audio_data',)", id="empty-audio"), + pytest.param({"id": "tts-1", "audio_data": 7}, "('id', 'audio_data')", id="numeric-audio"), + pytest.param({"audio_data": [_SPEECH_AUDIO_B64], "id": "tts-1"}, "('audio_data', 'id')", id="list-audio"), + ], +) +def test_transform_response_without_usable_audio_reports_the_response_keys( + payload: dict[str, object], expected_keys: str +): + raw_response: Final = httpx.Response(200, json=payload, request=httpx.Request("POST", SPEECH_URL)) + with pytest.raises(MistralTextToSpeechException) as exc_info: + MistralTextToSpeechConfig().transform_text_to_speech_response( + model="voxtral-mini-tts-2603", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + assert exc_info.value.status_code == 500 + assert exc_info.value.message == f"No audio_data in Mistral speech response. Response keys: {expected_keys}" + + +def test_transform_response_non_json_body_raises_provider_error_with_upstream_status(): + raw_response: Final = httpx.Response( + 502, + text="bad gateway", + request=httpx.Request("POST", SPEECH_URL), + ) + with pytest.raises(MistralTextToSpeechException) as exc_info: + MistralTextToSpeechConfig().transform_text_to_speech_response( + model="voxtral-mini-tts-2603", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + assert exc_info.value.status_code == 502 + assert exc_info.value.message == "Non-JSON response from Mistral speech API: bad gateway" + + +@pytest.mark.parametrize( + "response_body", + [b'["secret-audio-payload"]', b'"secret-audio-payload"', b"7", b"null", b"true"], +) +def test_transform_response_json_body_that_is_not_an_object_is_rejected_without_echoing_it(response_body: bytes): + raw_response: Final = httpx.Response(200, content=response_body, request=httpx.Request("POST", SPEECH_URL)) + with pytest.raises(ValidationError) as exc_info: + MistralTextToSpeechConfig().transform_text_to_speech_response( + model="voxtral-mini-tts-2603", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + assert "secret-audio-payload" not in str(exc_info.value) + assert "input_value" not in str(exc_info.value) + + +@pytest.mark.parametrize( + ("request_body", "expected_error"), + [ + pytest.param(b'["wav"]', ValidationError, id="json-array-body"), + pytest.param(b'"wav"', ValidationError, id="json-string-body"), + pytest.param(b"response_format=wav", json.JSONDecodeError, id="form-encoded-body"), + ], +) +def test_transform_response_rejects_a_request_body_that_is_not_a_json_object( + request_body: bytes, expected_error: type[Exception] +): + raw_response: Final = httpx.Response( + 200, + json={"audio_data": _SPEECH_AUDIO_B64}, + request=httpx.Request("POST", SPEECH_URL, content=request_body), + ) + with pytest.raises(expected_error): + MistralTextToSpeechConfig().transform_text_to_speech_response( + model="voxtral-mini-tts-2603", + raw_response=raw_response, + logging_obj=MagicMock(), + ) diff --git a/tests/unit/llms/nlp_cloud/__init__.py b/tests/unit/llms/nlp_cloud/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nlp_cloud/chat/__init__.py b/tests/unit/llms/nlp_cloud/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nlp_cloud/chat/test_nlp_cloud_chat_transformation.py b/tests/unit/llms/nlp_cloud/chat/test_nlp_cloud_chat_transformation.py new file mode 100644 index 00000000000..35f8d933cee --- /dev/null +++ b/tests/unit/llms/nlp_cloud/chat/test_nlp_cloud_chat_transformation.py @@ -0,0 +1,47 @@ +import pytest + +from litellm.llms.nlp_cloud.chat.transformation import NLPCloudConfig + + +@pytest.mark.parametrize( + ("supplied", "stored"), + [ + pytest.param({}, {}, id="nothing-supplied"), + pytest.param( + {"max_length": 0, "length_no_input": False, "end_sequence": "", "temperature": 0.0}, + {"max_length": 0, "length_no_input": False, "end_sequence": "", "temperature": 0.0}, + id="falsy-values-are-kept", + ), + pytest.param({"top_k": 40, "top_p": None, "num_beams": None}, {"top_k": 40}, id="none-is-skipped"), + ], +) +def test_init_stores_each_supplied_value_except_none_as_class_level_config( + supplied: dict[str, object], stored: dict[str, object] +) -> None: + class ScopedNLPCloudConfig(NLPCloudConfig): + pass + + ScopedNLPCloudConfig(**supplied) + + assert ScopedNLPCloudConfig.get_config() == stored + + +def test_init_stores_the_supplied_object_itself_rather_than_a_copy() -> None: + class ScopedNLPCloudConfig(NLPCloudConfig): + pass + + bad_words = ["darn", "heck"] + + ScopedNLPCloudConfig(bad_words=bad_words) + + assert ScopedNLPCloudConfig.get_config()["bad_words"] is bad_words + + +def test_init_called_again_with_none_keeps_the_value_stored_earlier() -> None: + class ScopedNLPCloudConfig(NLPCloudConfig): + pass + + ScopedNLPCloudConfig(top_k=40) + ScopedNLPCloudConfig(top_k=None, top_p=0.5) + + assert ScopedNLPCloudConfig.get_config() == {"top_k": 40, "top_p": 0.5} diff --git a/tests/unit/llms/ollama/test_ollama_completion_transformation.py b/tests/unit/llms/ollama/test_ollama_completion_transformation.py index 7be377247af..7e6dafe030c 100644 --- a/tests/unit/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_completion_transformation.py @@ -889,6 +889,59 @@ def test_transform_request_leaves_unreadable_images_untouched(payload: str) -> N assert data["images"] == [payload] +async def _complete_in_json_mode(response_text: str) -> ModelResponse: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"model": "llama3", "response": response_text, "done": True}) + + return await litellm.acompletion( + model="ollama/llama3", + messages=[{"role": "user", "content": "hi"}], + api_base="http://ollama.example:11434", + format="json", + client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)), + ) + + +@pytest.mark.parametrize( + ("response_text", "expected_content"), + [ + ('[{"name": "f", "arguments": {}}]', '[{"name": "f", "arguments": {}}]'), + ('["name", "arguments"]', '["name", "arguments"]'), + ('"name arguments"', '"name arguments"'), + ('{"name": "f"}', '{"name": "f"}'), + ('{"b": 1,"a": 2}', '{"b": 1, "a": 2}'), + ("1.50", "1.5"), + ("true", "true"), + ("null", "null"), + ], + ids=["list_of_calls", "list_of_key_names", "text_with_key_names", "name_only", "object", "number", "bool", "null"], +) +async def test_ollama_json_mode_completion_reserializes_json_that_is_not_a_function_call( + response_text: str, expected_content: str +) -> None: + response = await _complete_in_json_mode(response_text) + + assert response.choices[0].message.content == expected_content + assert response.choices[0].message.tool_calls is None + assert response.choices[0].finish_reason == "stop" + + +@pytest.mark.parametrize( + ("arguments_json", "expected_arguments"), + [('{"city":"Paris"}', '{"city": "Paris"}'), ("[1,2]", "[1, 2]"), ('"{}"', '"{}"'), ("null", "null")], + ids=["object", "list", "text", "null"], +) +async def test_ollama_json_mode_completion_turns_a_name_and_arguments_object_into_a_tool_call( + arguments_json: str, expected_arguments: str +) -> None: + response = await _complete_in_json_mode(f'{{"name": "get_weather", "arguments": {arguments_json}}}') + + (tool_call,) = response.choices[0].message.tool_calls + assert (tool_call.function.name, tool_call.function.arguments) == ("get_weather", expected_arguments) + assert response.choices[0].message.content is None + assert response.choices[0].finish_reason == "tool_calls" + + def test_ollama_vision_model_sends_the_image_beside_the_user_prompt(respx_mock: respx.MockRouter) -> None: image_base64: Final = _image_base64("PNG") route: Final = respx_mock.post("http://ollama.test:11434/api/generate").respond( diff --git a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py new file mode 100644 index 00000000000..2e0fc55cda4 --- /dev/null +++ b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py @@ -0,0 +1,187 @@ +import json +import operator +from typing import Final + +import httpx +import pytest +import respx +from pydantic import ValidationError + +import litellm +from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter +from litellm.types.utils import TokenCountResponse + +COUNT_URL: Final = "https://api.openai.com/v1/responses/input_tokens" + + +async def _count(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> TokenCountResponse | None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + respx.post(COUNT_URL).mock(return_value=response) + return await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"api_key": "sk-test"}}, + request_model="openai/gpt-4o", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("payload", "expected_total"), + [ + ({"input_tokens": 15}, 15), + ({"input_tokens": "15"}, 15), + ({"input_tokens": 15.0}, 15), + ({"input_tokens": True}, 1), + ({"input_tokens": -3}, -3), + ({}, 0), + ({"object": "response.input_tokens", "error": {"message": "429 rate limit"}}, 0), + ({"input_tokens": 7, "object": "response.input_tokens", "nested": {"a": [1, None]}}, 7), + ], +) +async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload( + monkeypatch: pytest.MonkeyPatch, payload: dict[str, object], expected_total: int +) -> None: + result: Final = await _count(monkeypatch, httpx.Response(200, json=payload)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": expected_total, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "original_response": payload, + "error": False, + "error_message": None, + "status_code": None, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "rejected_count", + [ + None, + 15.5, + "Request timed out", + "429 Too Many Requests", + [15], + [*range(403), "Request timed out"], + {"value": 429}, + ], +) +async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_validation_message( + monkeypatch: pytest.MonkeyPatch, rejected_count: object +) -> None: + with pytest.raises(ValidationError) as rejection: + TokenCountResponse.model_validate( + { + "total_tokens": rejected_count, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + } + ) + + result: Final = await _count(monkeypatch, httpx.Response(200, json={"input_tokens": rejected_count})) + + assert result is not None + assert [error["loc"] for error in rejection.value.errors()] == [("total_tokens",)] + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "original_response": None, + "error": True, + "error_message": str(rejection.value), + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [[*({"input_tokens": position} for position in range(403)), "429"], "input_tokens", 15, 0, False], +) +async def test_count_payload_that_is_not_an_object_becomes_an_error_response_naming_the_failed_lookup( + monkeypatch: pytest.MonkeyPatch, payload: object +) -> None: + with pytest.raises(AttributeError) as failed_lookup: + operator.attrgetter("get")(payload) + + result: Final = await _count(monkeypatch, httpx.Response(200, json=payload)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "original_response": None, + "error": True, + "error_message": str(failed_lookup.value), + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +async def test_null_count_payload_falls_back_to_local_counting(monkeypatch: pytest.MonkeyPatch) -> None: + assert await _count(monkeypatch, httpx.Response(200, content=b"null")) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("upstream_status", "upstream_text"), + [(429, "rate limit exceeded"), (403, "forbidden"), (408, "Request timed out")], +) +async def test_upstream_error_status_becomes_an_error_response_with_that_status_and_text( + monkeypatch: pytest.MonkeyPatch, upstream_status: int, upstream_text: str +) -> None: + result: Final = await _count(monkeypatch, httpx.Response(upstream_status, text=upstream_text)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "original_response": None, + "error": True, + "error_message": upstream_text, + "status_code": upstream_status, + } + ) + + +@pytest.mark.asyncio +async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with pytest.raises(json.JSONDecodeError) as undecodable: + json.loads("not json") + + result: Final = await _count(monkeypatch, httpx.Response(200, text="not json")) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "original_response": None, + "error": True, + "error_message": f"CountTokens processing error: {undecodable.value}", + "status_code": 500, + } + ) diff --git a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py index 32b47c1c680..59355dfabcb 100644 --- a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py +++ b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py @@ -7,8 +7,7 @@ import pytest import respx import litellm - - +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.llms.openrouter.chat.transformation import ( OpenRouterChatCompletionStreamingHandler, @@ -628,6 +627,88 @@ def test_openrouter_reasoning_effort_high_passes_through(): assert result["reasoning_effort"] == "high" +OPENROUTER_COST_HEADER = "llm_provider-x-litellm-response-cost" + + +def _complete_with_openrouter_usage(usage): + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "gen-1", + "object": "chat.completion", + "created": 1700000000, + "model": "openai/gpt-4o", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}], + "usage": usage, + }, + ) + + return litellm.completion( + model="openrouter/openai/gpt-4o", + messages=[{"role": "user", "content": "hi"}], + api_key="sk-test", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), + ) + + +@pytest.mark.parametrize( + ("cost", "expected"), + [ + pytest.param(0.00015, 0.00015, id="float"), + pytest.param(3, 3.0, id="int"), + pytest.param(0, 0.0, id="zero"), + pytest.param("0.5", 0.5, id="numeric-string"), + pytest.param(" 1.5 ", 1.5, id="padded-numeric-string"), + pytest.param(True, 1.0, id="bool"), + ], +) +def test_openrouter_completion_stores_usage_cost_as_float_cost_header(cost, expected): + response = _complete_with_openrouter_usage( + {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30, "cost": cost, "is_byok": False} + ) + + stored = response._hidden_params["additional_headers"][OPENROUTER_COST_HEADER] + assert type(stored) is float + assert stored == expected + assert response._hidden_params["response_cost"] == expected + + +@pytest.mark.parametrize( + "usage", + [ + pytest.param({"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, id="cost-missing"), + pytest.param({"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30, "cost": None}, id="cost-null"), + pytest.param({"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30, "cost": "abc"}, id="cost-text"), + pytest.param({"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30, "cost": [1]}, id="cost-list"), + pytest.param( + {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30, "cost": {"upstream": 1.0}}, + id="cost-object", + ), + pytest.param({}, id="usage-empty"), + pytest.param(None, id="usage-null"), + ], +) +def test_openrouter_completion_without_usable_cost_returns_response_without_cost_header(usage): + response = _complete_with_openrouter_usage(usage) + + assert response.choices[0].message.content == "hi" + assert OPENROUTER_COST_HEADER not in response._hidden_params["additional_headers"] + + +@pytest.mark.parametrize( + "usage", + [ + pytest.param([{"cost": 1.0}], id="usage-list"), + pytest.param("cost", id="usage-text"), + pytest.param(5, id="usage-number"), + ], +) +def test_openrouter_completion_with_non_object_usage_raises_api_connection_error(usage): + with pytest.raises(litellm.APIConnectionError): + _complete_with_openrouter_usage(usage) + + @respx.mock def test_openrouter_streaming_completion_assembles_text_and_finish_reason() -> None: chunks: Final = ( diff --git a/tests/unit/llms/opensandbox/__init__.py b/tests/unit/llms/opensandbox/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/opensandbox/sandbox/__init__.py b/tests/unit/llms/opensandbox/sandbox/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/opensandbox/sandbox/test_transformation.py b/tests/unit/llms/opensandbox/sandbox/test_transformation.py new file mode 100644 index 00000000000..ec2a47cb92f --- /dev/null +++ b/tests/unit/llms/opensandbox/sandbox/test_transformation.py @@ -0,0 +1,153 @@ +import json +from collections.abc import Sequence + +import httpx +import pytest +from pydantic import ValidationError + +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.opensandbox.sandbox.transformation import OpenSandboxSandboxConfig + +API_BASE = "https://sandbox.test/v1" +PENDING = {"id": "osb_1", "status": {"state": "Pending"}} +RUNNING = {"id": "osb_1", "status": {"state": "Running"}} +EXECD_ENDPOINT = {"endpoint": "execd.local:44772", "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}} + + +def _sandbox_api(created: object, polled: Sequence[object] = ()) -> tuple[AsyncHTTPHandler, list[str]]: + remaining = list(polled) + requested: list[str] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requested.append(f"{request.method} {request.url.path}") + if request.method == "POST": + return httpx.Response(200, content=json.dumps(created).encode()) + if "/endpoints/" in request.url.path: + return httpx.Response(200, content=json.dumps(EXECD_ENDPOINT).encode()) + return httpx.Response(200, content=json.dumps(remaining.pop(0)).encode()) + + return AsyncHTTPHandler(transport=httpx.MockTransport(respond)), requested + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("created", "expected_id"), + [ + pytest.param( + {**RUNNING, "createdAt": "2026-01-01T00:00:00Z", "metadata": {"team": ["a"]}}, "osb_1", id="text-id" + ), + pytest.param({"id": 123, "status": {"state": "Running"}}, "123", id="numeric-id"), + ], +) +async def test_acreate_sandbox_builds_the_handle_from_a_running_create_response(created: object, expected_id: str): + client, requested = _sandbox_api(created) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox(api_key="osb_key", api_base=API_BASE, client=client) + + assert (handle.id, handle.provider, handle.domain) == (expected_id, "opensandbox", API_BASE) + assert type(handle.id) is str + assert handle._hidden_params == { + "api_base": API_BASE, + "api_key": "osb_key", + "execd_endpoint": "execd.local:44772", + "execd_headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, + "use_server_proxy": False, + } + assert requested == ["POST /v1/sandboxes", f"GET /v1/sandboxes/{expected_id}/endpoints/44772"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "created", + [ + pytest.param(PENDING, id="pending"), + pytest.param({"id": "osb_1"}, id="status-missing"), + pytest.param({"id": "osb_1", "status": None}, id="status-null"), + pytest.param({"id": "osb_1", "status": "Running"}, id="status-text"), + pytest.param({"id": "osb_1", "status": {"state": None}}, id="state-null"), + ], +) +async def test_acreate_sandbox_polls_when_the_create_response_does_not_report_running(created: object): + client, requested = _sandbox_api(created, polled=[RUNNING]) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox(api_key="osb_key", api_base=API_BASE, client=client) + + assert handle.id == "osb_1" + assert requested == ["POST /v1/sandboxes", "GET /v1/sandboxes/osb_1", "GET /v1/sandboxes/osb_1/endpoints/44772"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "not_ready", + [ + pytest.param(PENDING, id="pending"), + pytest.param([], id="list"), + pytest.param("Running", id="text"), + pytest.param(None, id="null"), + pytest.param(7, id="number"), + pytest.param({}, id="empty-object"), + pytest.param({"status": ["Running"]}, id="status-list"), + pytest.param({"status": {"state": 5}}, id="state-number"), + ], +) +async def test_acreate_sandbox_keeps_polling_past_a_status_response_without_a_running_state(not_ready: object): + client, requested = _sandbox_api(PENDING, polled=[not_ready, RUNNING]) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="osb_key", api_base=API_BASE, poll_interval=0, client=client + ) + + assert handle.id == "osb_1" + assert requested == [ + "POST /v1/sandboxes", + "GET /v1/sandboxes/osb_1", + "GET /v1/sandboxes/osb_1", + "GET /v1/sandboxes/osb_1/endpoints/44772", + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("state", ["Failed", "Stopping", "Terminated"]) +async def test_acreate_sandbox_raises_when_the_polled_sandbox_reaches_a_terminal_state(state: str): + client, requested = _sandbox_api(PENDING, polled=[[], {"status": {"state": state}}]) + + with pytest.raises(ValueError, match=f"OpenSandbox sandbox osb_1 entered {state}"): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="osb_key", api_base=API_BASE, poll_interval=0, client=client + ) + + assert requested == ["POST /v1/sandboxes", "GET /v1/sandboxes/osb_1", "GET /v1/sandboxes/osb_1"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("created", [["secret-token"], "secret-token", 7, None]) +async def test_acreate_sandbox_rejects_a_create_response_that_is_not_an_object_without_echoing_it(created: object): + client, requested = _sandbox_api(created) + + with pytest.raises(ValidationError) as exc_info: + await OpenSandboxSandboxConfig().acreate_sandbox(api_key="osb_key", api_base=API_BASE, client=client) + + assert [error["type"] for error in exc_info.value.errors()] == ["dict_type"] + assert "secret-token" not in str(exc_info.value) + assert requested == ["POST /v1/sandboxes"] + + +@pytest.mark.asyncio +async def test_public_acreate_sandbox_raises_the_validation_error_for_a_create_response_that_is_not_an_object(): + client, requested = _sandbox_api([RUNNING]) + + with pytest.raises(ValidationError): + await litellm.acreate_sandbox(provider="opensandbox", api_key="osb_key", api_base=API_BASE, client=client) + + assert requested == ["POST /v1/sandboxes"] + + +@pytest.mark.asyncio +async def test_acreate_sandbox_requires_a_sandbox_id(): + client, requested = _sandbox_api({"status": {"state": "Running"}}) + + with pytest.raises(KeyError, match="id"): + await OpenSandboxSandboxConfig().acreate_sandbox(api_key="osb_key", api_base=API_BASE, client=client) + + assert requested == ["POST /v1/sandboxes"] diff --git a/tests/unit/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py b/tests/unit/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py index ad2a585b536..7dcd910113c 100644 --- a/tests/unit/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py +++ b/tests/unit/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py @@ -1,6 +1,14 @@ +import datetime from unittest.mock import patch +import httpx +import pytest +from pydantic import ValidationError + import litellm +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.ovhcloud.embedding.transformation import OVHCloudEmbeddingConfig +from litellm.types.utils import EmbeddingResponse model = "ovhcloud/BGE-M3" @@ -36,3 +44,84 @@ def test_ovhcloud_embeddings(): assert isinstance(response.data[0]["embedding"], list) assert response.model == model assert response.object == "embedding" + + +EMBEDDING_ROW = {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]} + + +def _transform(payload: object) -> EmbeddingResponse: + logging_obj = Logging( + model="ovhcloud/BGE-M3", + messages=[], + stream=False, + call_type="embedding", + start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-id", + function_id="function-id", + ) + return OVHCloudEmbeddingConfig().transform_embedding_response( + model="BGE-M3", + raw_response=httpx.Response(200, json=payload), + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + api_key="sk-test", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + +def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]: + assert response.usage is not None + return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) + + +def test_ovhcloud_embedding_response_carries_provider_fields_and_usage(): + response = _transform( + { + "object": "list", + "model": "BGE-M3", + "data": [EMBEDDING_ROW], + "usage": {"prompt_tokens": 8, "total_tokens": 9}, + } + ) + + assert response.model == "BGE-M3" + assert response.object == "list" + assert response.data == [EMBEDDING_ROW] + assert _token_counts(response) == (8, 0, 9) + + +@pytest.mark.parametrize( + ("payload", "expected"), + [ + ({"data": [EMBEDDING_ROW]}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 4}}, (4, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 6}}, (0, 0, 6)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": None, "total_tokens": None}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": "", "total_tokens": []}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": "12", "total_tokens": 12.0}}, (12, 0, 12)), + ( + {"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 3, "total_tokens": 5, "completion_tokens": 2}}, + (3, 0, 5), + ), + ], +) +def test_ovhcloud_embedding_usage_tolerates_missing_null_and_loosely_typed_counts(payload, expected): + response = _transform(payload) + + assert response.data == [EMBEDDING_ROW] + assert _token_counts(response) == expected + + +@pytest.mark.parametrize("usage", [None, [], "tokens", 3]) +def test_ovhcloud_embedding_usage_that_is_not_an_object_is_rejected(usage): + with pytest.raises(ValidationError): + _transform({"data": [EMBEDDING_ROW], "usage": usage}) + + +@pytest.mark.parametrize("usage", [{"prompt_tokens": "many"}, {"total_tokens": 1.5}, {"prompt_tokens": [1]}]) +def test_ovhcloud_embedding_token_count_that_is_not_an_integer_is_rejected(usage): + with pytest.raises(ValidationError): + _transform({"data": [EMBEDDING_ROW], "usage": usage}) diff --git a/tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py b/tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py index 29d185b686d..2c95265890d 100644 --- a/tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py +++ b/tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py @@ -5,13 +5,16 @@ Tests the response transformation to extract citation tokens and search queries from Perplexity API responses. """ +import datetime from unittest.mock import Mock +import httpx import pytest # Add the project root to Python path from litellm import ModelResponse +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.perplexity.chat.transformation import PerplexityChatConfig from litellm.types.utils import Usage @@ -750,3 +753,167 @@ class TestPerplexityChatTransformation: not hasattr(choice.message, "annotations") or choice.message.annotations is None ) + + +_PERPLEXITY_COMPLETION = { + "id": "cmpl-1", + "object": "chat.completion", + "created": 1700000000, + "model": "sonar", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "Paris [1] is in France [2][7]."}, + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, +} +_FIRST_CITATION = { + "type": "url_citation", + "url_citation": {"url": "https://a.example", "title": "", "start_index": 6, "end_index": 9}, +} + + +def _transform_perplexity_response(body: object) -> ModelResponse: + logging_obj = Logging( + model="sonar", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="perplexity-transform-test", + function_id="perplexity-transform-test", + ) + return PerplexityChatConfig().transform_response( + model="sonar", + raw_response=httpx.Response(200, json=body), + model_response=ModelResponse(), + logging_obj=logging_obj, + request_data={"model": "sonar"}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_transform_response_adds_citation_annotations_and_search_usage_from_the_raw_body(): + response = _transform_perplexity_response( + { + **_PERPLEXITY_COMPLETION, + "citations": ["https://a.example", "https://b.example"], + "search_results": [ + {"url": "https://a.example", "title": "A"}, + {"url": "https://c.example", "title": "C"}, + ], + "usage": {**_PERPLEXITY_COMPLETION["usage"], "num_search_queries": 3}, + } + ) + + assert response.choices[0].message.annotations == [ + { + "type": "url_citation", + "url_citation": {"url": "https://a.example", "title": "A", "start_index": 6, "end_index": 9}, + }, + { + "type": "url_citation", + "url_citation": {"url": "https://b.example", "title": "", "start_index": 23, "end_index": 26}, + }, + ] + assert repr(response.usage.prompt_tokens_details.web_search_requests) == "3" + assert repr(response.usage.citation_tokens) == "8" + assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (10, 5, 15) + assert response.id == "cmpl-1" + + +@pytest.mark.parametrize( + "extra", + [{}, {"citations": None}, {"citations": []}, {"search_results": []}, {"num_search_queries": 0}], +) +def test_transform_response_leaves_a_completion_without_perplexity_fields_unchanged(extra: dict): + response = _transform_perplexity_response({**_PERPLEXITY_COMPLETION, **extra}) + + assert response.choices[0].message.content == "Paris [1] is in France [2][7]." + assert not hasattr(response.choices[0].message, "annotations") + assert not hasattr(response.usage, "citation_tokens") + assert response.usage.prompt_tokens_details is None + assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (10, 5, 15) + + +def test_transform_response_reads_search_queries_from_the_root_of_the_body(): + response = _transform_perplexity_response({**_PERPLEXITY_COMPLETION, "num_search_queries": 2}) + + assert repr(response.usage.prompt_tokens_details.web_search_requests) == "2" + assert not hasattr(response.usage, "citation_tokens") + + +def test_transform_response_annotates_and_counts_citations_when_the_provider_sends_no_usage(): + body = {key: value for key, value in _PERPLEXITY_COMPLETION.items() if key != "usage"} + + response = _transform_perplexity_response({**body, "citations": ["https://a.example"]}) + + assert response.choices[0].message.annotations == [_FIRST_CITATION] + assert repr(response.usage.citation_tokens) == "4" + assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (0, 0, 0) + + +@pytest.mark.parametrize( + ("extra", "citation_tokens_recorded"), + [ + ({"usage": None}, False), + ({"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "num_search_queries": "3"}}, True), + ({"search_results": 3}, True), + ({"search_results": None}, True), + ], +) +def test_transform_response_still_returns_the_completion_when_a_perplexity_field_is_malformed( + extra: dict, citation_tokens_recorded: bool +): + response = _transform_perplexity_response({**_PERPLEXITY_COMPLETION, "citations": ["https://a.example"], **extra}) + + assert response.choices[0].message.content == "Paris [1] is in France [2][7]." + assert not hasattr(response.choices[0].message, "annotations") + assert hasattr(response.usage, "citation_tokens") is citation_tokens_recorded + + +@pytest.mark.parametrize("body", [[], [_PERPLEXITY_COMPLETION], "choices", ""]) +def test_transform_response_rejects_a_body_that_is_not_an_object_as_an_invalid_response_object(body: object): + with pytest.raises(Exception, match="Invalid response object") as exc_info: + _transform_perplexity_response(body) + + assert type(exc_info.value) is Exception + + +@pytest.mark.parametrize("body", [3, 1.5, True, ["error"], "error"]) +def test_transform_response_raises_type_error_for_a_scalar_or_error_bearing_non_object_body(body: object): + with pytest.raises(TypeError): + _transform_perplexity_response(body) + + +def test_transform_response_keeps_body_values_of_every_json_type_as_the_provider_sent_them(): + citations = ["https://a.example", 5, None, 1.5, True, {"url": "https://d.example"}, ["https://e.example"]] + search_results = [{"url": "https://a.example", "title": "A", "score": 0.5, "tags": None}, "loose", 7] + + response = _transform_perplexity_response( + { + **_PERPLEXITY_COMPLETION, + "citations": citations, + "search_results": search_results, + "related_score": 1.5, + "is_final": True, + "trace": {"hops": [1, {"next": None}]}, + } + ) + + assert response.choices[0].message.annotations == [ + { + "type": "url_citation", + "url_citation": {"url": "https://a.example", "title": "A", "start_index": 6, "end_index": 9}, + } + ] + assert response.citations == citations + assert [type(citation) for citation in response.citations] == [str, int, type(None), float, bool, dict, list] + assert response.search_results == search_results + assert [type(result) for result in response.search_results] == [dict, str, int] + assert repr(response.usage.citation_tokens) == "18" diff --git a/tests/unit/llms/perplexity/embedding/test_perplexity_embedding_transformation.py b/tests/unit/llms/perplexity/embedding/test_perplexity_embedding_transformation.py index 6ca7072e7ab..56afc7a1d67 100644 --- a/tests/unit/llms/perplexity/embedding/test_perplexity_embedding_transformation.py +++ b/tests/unit/llms/perplexity/embedding/test_perplexity_embedding_transformation.py @@ -3,12 +3,15 @@ Unit tests for Perplexity embedding transformation logic. """ import base64 +import datetime import struct from unittest.mock import MagicMock import httpx import pytest +from pydantic import ValidationError +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.perplexity.embedding.transformation import ( PerplexityEmbeddingConfig, PerplexityEmbeddingError, @@ -297,3 +300,100 @@ class TestPerplexityEmbeddingProviderConfig: ) assert config is not None assert isinstance(config, PerplexityEmbeddingConfig) + + +FLOAT_ROW = {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]} + + +def _transform(payload: object) -> EmbeddingResponse: + logging_obj = Logging( + model="perplexity/pplx-embed-v1-0.6b", + messages=[], + stream=False, + call_type="embedding", + start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-id", + function_id="function-id", + ) + return PerplexityEmbeddingConfig().transform_embedding_response( + model="pplx-embed-v1-0.6b", + raw_response=httpx.Response(200, json=payload), + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + ) + + +def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]: + assert response.usage is not None + return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) + + +@pytest.mark.parametrize( + ("usage", "expected"), + [ + ({"prompt_tokens": 4, "total_tokens": 6}, (4, 0, 6)), + ({"total_tokens": 6}, (6, 0, 6)), + ({"prompt_tokens": 0, "total_tokens": 6}, (6, 0, 6)), + ({"prompt_tokens": None, "total_tokens": 6}, (6, 0, 6)), + ({"prompt_tokens": 4}, (4, 0, 0)), + ({}, (0, 0, 0)), + ({"prompt_tokens": None, "total_tokens": None}, (0, 0, 0)), + ({"prompt_tokens": "", "total_tokens": []}, (0, 0, 0)), + ({"prompt_tokens": "12", "total_tokens": 12.0}, (12, 0, 12)), + ], +) +def test_perplexity_embedding_usage_falls_back_to_total_tokens_and_tolerates_loose_counts(usage, expected): + response = _transform({"data": [FLOAT_ROW], "usage": usage}) + + assert response.data == [FLOAT_ROW] + assert _token_counts(response) == expected + + +def test_perplexity_embedding_response_without_usage_reports_zero_tokens(): + assert _token_counts(_transform({"data": [FLOAT_ROW]})) == (0, 0, 0) + + +def test_perplexity_embedding_rows_keep_their_fields_in_order_around_the_decoded_vector(): + encoded = base64.b64encode(struct.pack("3b", 127, 0, -127)).decode() + + response = _transform( + {"data": [{"z": 1, "embedding": encoded, "index": 0}, {"index": 1, "embedding": [0.5]}, {"index": 2}]} + ) + + assert [list(row.items()) for row in response.data] == [ + [("z", 1), ("embedding", [1.0, 0.0, -1.0]), ("index", 0)], + [("index", 1), ("embedding", [0.5])], + [("index", 2), ("embedding", None)], + ] + + +@pytest.mark.parametrize("payload", [{}, {"data": []}, {"data": ""}, {"data": {}}]) +def test_perplexity_embedding_response_without_rows_has_no_data(payload): + assert _transform(payload).data == [] + + +@pytest.mark.parametrize("row", [None, 3, "ab", [], [["embedding", [1.0]]]]) +def test_perplexity_embedding_row_that_is_not_an_object_is_rejected(row): + with pytest.raises(ValidationError): + _transform({"data": [FLOAT_ROW, row]}) + + +@pytest.mark.parametrize("usage", [None, [], "tokens", 3]) +def test_perplexity_embedding_usage_that_is_not_an_object_is_rejected(usage): + with pytest.raises(ValidationError): + _transform({"data": [FLOAT_ROW], "usage": usage}) + + +@pytest.mark.parametrize( + "usage", + [ + {"prompt_tokens": "many"}, + {"total_tokens": 1.5}, + {"prompt_tokens": [1]}, + {"prompt_tokens": 4, "total_tokens": 1.5}, + {"prompt_tokens": 4, "total_tokens": "many"}, + ], +) +def test_perplexity_embedding_token_count_that_is_not_an_integer_is_rejected(usage): + with pytest.raises(ValidationError): + _transform({"data": [FLOAT_ROW], "usage": usage}) diff --git a/tests/unit/llms/perplexity/search/test_perplexity_search_transformation.py b/tests/unit/llms/perplexity/search/test_perplexity_search_transformation.py index 267f8eacb08..af4d8c0c706 100644 --- a/tests/unit/llms/perplexity/search/test_perplexity_search_transformation.py +++ b/tests/unit/llms/perplexity/search/test_perplexity_search_transformation.py @@ -1,11 +1,17 @@ +import datetime import json from collections.abc import Iterator from typing import Final +import httpx import pytest import respx +from pydantic import ValidationError import litellm +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult +from litellm.llms.perplexity.search.transformation import PerplexitySearchConfig PERPLEXITY_SEARCH_URL: Final = "https://api.perplexity.ai/search" @@ -66,6 +72,116 @@ def test_search_sends_max_results_in_request_body(respx_mock: respx.MockRouter) assert [result.url for result in response.results] == ["https://example.com/ml"] +def _transform(payload: object) -> SearchResponse: + logging_obj = Logging( + model="perplexity/search", + messages=[], + stream=False, + call_type="search", + start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-id", + function_id="function-id", + ) + return PerplexitySearchConfig().transform_search_response( + raw_response=httpx.Response(200, json=payload), logging_obj=logging_obj + ) + + +def test_perplexity_results_keep_provider_order_fields_and_dates(): + response = _transform( + { + "id": "search-1", + "server_time": None, + "results": [ + { + "title": "LiteLLM", + "url": "https://example.com/litellm", + "snippet": "Call every LLM API", + "date": "2024-01-05", + "last_updated": "2024-02-01", + }, + { + "title": "Docs", + "url": "https://example.com/docs", + "snippet": "Docs", + "date": None, + "last_updated": None, + }, + ], + } + ) + + assert response.object == "search" + assert response.results == [ + SearchResult( + title="LiteLLM", + url="https://example.com/litellm", + snippet="Call every LLM API", + date="2024-01-05", + last_updated="2024-02-01", + ), + SearchResult(title="Docs", url="https://example.com/docs", snippet="Docs", date=None, last_updated=None), + ] + + +@pytest.mark.parametrize( + "payload", + [{}, {"id": "search-1"}, {"id": "search-1", "results": []}, {"results": ""}, {"results": {}}], +) +def test_perplexity_response_without_results_is_empty(payload): + assert _transform(payload).results == [] + + +def test_perplexity_result_missing_optional_fields_defaults_to_empty_strings_and_no_dates(): + assert _transform({"results": [{"score": 0.5}]}).results == [ + SearchResult(title="", url="", snippet="", date=None, last_updated=None) + ] + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("title", None), + ("title", 7), + ("url", None), + ("url", ["https://example.com"]), + ("snippet", None), + ("snippet", {"text": "snippet"}), + ("date", 1704412800), + ("last_updated", ["2024-02-01"]), + ], +) +def test_perplexity_result_with_field_of_wrong_type_is_rejected(field, value): + result = {"title": "T", "url": "https://example.com", "snippet": "S", field: value} + + with pytest.raises(ValidationError): + _transform({"results": [result]}) + + +def test_perplexity_result_with_several_fields_of_wrong_type_reports_every_field(): + result = {"title": 7, "url": None, "snippet": ["snippet"], "date": 1704412800, "last_updated": ["2024-02-01"]} + + with pytest.raises(ValidationError) as exc_info: + _transform({"results": [result]}) + + assert [error["loc"] for error in exc_info.value.errors()] == [ + ("title",), + ("url",), + ("snippet",), + ("date",), + ("last_updated",), + ] + + +def test_perplexity_malformed_result_is_rejected_without_naming_its_position(): + results = [{}] * 429 + ["not-an-object"] + + with pytest.raises(ValidationError) as exc_info: + _transform({"results": results}) + + assert "429" not in str(exc_info.value) + + @pytest.mark.asyncio async def test_router_search_with_search_tools( respx_mock: respx.MockRouter, httpx_transport: None diff --git a/tests/unit/llms/replicate/chat/test_handler.py b/tests/unit/llms/replicate/chat/test_handler.py new file mode 100644 index 00000000000..756a2a21db3 --- /dev/null +++ b/tests/unit/llms/replicate/chat/test_handler.py @@ -0,0 +1,206 @@ +import httpx +import pytest +import respx +from pydantic import ValidationError + +import litellm +from litellm.exceptions import MidStreamFallbackError +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.replicate.chat.handler import handle_prediction_response_streaming +from litellm.llms.replicate.common_utils import ReplicateError + +POLL_URL = "https://api.replicate.com/v1/predictions/test-prediction" +CREATE_URL = "https://api.replicate.com/v1/models/meta/meta-llama-3-8b-instruct/predictions" +MANY_TOKENS = [f"token-{index} " for index in range(500)] + + +@pytest.fixture(autouse=True) +def _poll_without_waiting(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr("litellm.llms.replicate.chat.handler.REPLICATE_POLLING_DELAY_SECONDS", 0) + + +def _poll(*bodies: object) -> list[dict[str, object]]: + with respx.mock: + respx.get(POLL_URL).mock(side_effect=[httpx.Response(200, json=body) for body in bodies]) + return list( + handle_prediction_response_streaming( + POLL_URL, + "replicate-test-key", + lambda _message: None, + headers={"Authorization": "Bearer replicate-test-key"}, + http_client=HTTPHandler(), + ) + ) + + +def test_streaming_poll_yields_only_the_text_added_since_the_previous_poll(): + chunks = _poll( + {"status": "starting"}, + {"status": "processing", "output": ["Hel"]}, + {"status": "processing", "output": ["Hel", "lo"]}, + {"status": "succeeded", "output": ["Hel", "lo", ", world"], "metrics": {"predict_time": 1.5}}, + ) + + assert chunks == [ + {"output": "Hel", "status": "processing"}, + {"output": "lo", "status": "processing"}, + {"output": ", world", "status": "succeeded"}, + ] + + +@pytest.mark.parametrize( + ("output", "text"), + [ + (["Hello", ", ", "world"], "Hello, world"), + ("Hello", "Hello"), + ({"Hel": 1, "lo": [2]}, "Hello"), + (["12", "12.0"], "1212.0"), + ([], ""), + ("", ""), + ({}, ""), + (MANY_TOKENS, "".join(MANY_TOKENS)), + ], +) +def test_streaming_poll_joins_the_prediction_output_the_way_str_join_does(output: object, text: str): + assert _poll({"status": "succeeded", "output": output}) == [{"output": text, "status": "succeeded"}] + + +@pytest.mark.parametrize("status", [None, 1, ["succeeded"], {"state": "succeeded"}]) +def test_streaming_poll_passes_an_unrecognised_status_through_and_keeps_polling(status: object): + chunks = _poll({"status": status, "output": ["a"]}, {"status": "succeeded", "output": ["a", "b"]}) + + assert chunks == [{"output": "a", "status": status}, {"output": "b", "status": "succeeded"}] + + +@pytest.mark.parametrize( + "output", + [ + None, + 5, + 0, + 1.5, + True, + ["a", 1], + ["a", None], + [["a"]], + [{"a": 1}], + [*MANY_TOKENS[:403], 403, *MANY_TOKENS[404:]], + ], +) +def test_streaming_poll_reports_an_output_that_is_not_text_as_unprocessable(output: object): + with pytest.raises(ReplicateError) as exc_info: + _poll({"status": "succeeded", "output": output}) + + assert exc_info.value.status_code == 422 + assert exc_info.value.message == f"Unable to parse response. Got={output}" + + +@pytest.mark.parametrize( + ("body", "message"), + [ + ({"status": "failed", "error": "boom"}, "Error: boom"), + ({"status": "failed", "error": None}, "Error: None"), + ({"status": "failed"}, "Error: "), + ({"status": "failed", "error": {"detail": "bad input"}}, "Error: {'detail': 'bad input'}"), + ], +) +def test_streaming_poll_raises_the_provider_error_of_a_failed_prediction(body: object, message: str): + with pytest.raises(ReplicateError) as exc_info: + _poll(body) + + assert exc_info.value.status_code == 400 + assert exc_info.value.message == message + + +@pytest.mark.parametrize("body", [{}, {"output": ["a"]}]) +def test_streaming_poll_raises_key_error_when_the_prediction_has_no_status(body: object): + with pytest.raises(KeyError) as exc_info: + _poll(body) + + assert exc_info.value.args == ("status",) + + +@pytest.mark.parametrize( + "body", [[{"status": "succeeded"}], [], "Request timed out", ["secret-output", 429], 3, True, 1.5] +) +def test_streaming_poll_rejects_a_body_that_is_not_an_object_without_echoing_it(body: object): + with pytest.raises(ValidationError) as exc_info: + _poll(body) + + assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [("dict_type", ())] + assert "secret-output" not in str(exc_info.value) + assert "429" not in str(exc_info.value) + assert "Request timed out" not in str(exc_info.value) + + +def test_streaming_poll_keeps_the_chunks_yielded_before_a_malformed_body(): + chunks = [] + with respx.mock: + respx.get(POLL_URL).mock( + side_effect=[ + httpx.Response(200, json={"status": "processing", "output": ["a"]}), + httpx.Response(200, json=["not", "an", "object"]), + ] + ) + stream = handle_prediction_response_streaming( + POLL_URL, "replicate-test-key", lambda _message: None, headers={}, http_client=HTTPHandler() + ) + chunks.append(next(stream)) + with pytest.raises(ValidationError): + next(stream) + + assert chunks == [{"output": "a", "status": "processing"}] + + +def test_streaming_poll_retries_after_a_non_200_status(): + with respx.mock: + respx.get(POLL_URL).mock( + side_effect=[ + httpx.Response(503, text="busy"), + httpx.Response(200, json={"status": "succeeded", "output": ["ok"]}), + ] + ) + chunks = list( + handle_prediction_response_streaming( + POLL_URL, "replicate-test-key", lambda _message: None, headers={}, http_client=HTTPHandler() + ) + ) + + assert chunks == [{"output": "ok", "status": "succeeded"}] + + +def _stream_completion(*poll_bodies: object) -> list[str | None]: + with respx.mock: + respx.post(CREATE_URL).mock(return_value=httpx.Response(201, json={"urls": {"get": POLL_URL}})) + respx.get(POLL_URL).mock(side_effect=[httpx.Response(200, json=body) for body in poll_bodies]) + stream = litellm.completion( + model="replicate/meta/meta-llama-3-8b-instruct", + messages=[{"role": "user", "content": "hi"}], + stream=True, + api_key="replicate-test-key", + ) + return [chunk.choices[0].delta.content for chunk in stream] + + +def test_streaming_completion_emits_each_polled_delta_as_content(): + contents = _stream_completion( + {"status": "processing", "output": ["Hel"]}, {"status": "succeeded", "output": ["Hel", "lo"]} + ) + + assert contents == ["Hel", "lo", None] + + +def test_streaming_completion_maps_a_non_text_output_to_unprocessable_entity(): + with pytest.raises(litellm.UnprocessableEntityError) as exc_info: + _stream_completion({"status": "succeeded", "output": ["a", 1]}) + + assert "Unable to parse response. Got=['a', 1]" in str(exc_info.value) + + +@pytest.mark.parametrize("body", [["x"], "Request timed out", ["Request was throttled", 429]]) +def test_streaming_completion_maps_a_non_object_poll_body_to_an_api_error_whatever_text_it_carries(body: object): + with pytest.raises(MidStreamFallbackError) as exc_info: + _stream_completion(body) + + assert type(exc_info.value.original_exception) is litellm.APIError + assert exc_info.value.original_exception.status_code == 500 diff --git a/tests/unit/llms/replicate/chat/test_transformation.py b/tests/unit/llms/replicate/chat/test_transformation.py index 549e915a6d5..fe6dfbc8653 100644 --- a/tests/unit/llms/replicate/chat/test_transformation.py +++ b/tests/unit/llms/replicate/chat/test_transformation.py @@ -3,13 +3,14 @@ from typing import Final from unittest.mock import AsyncMock, Mock, patch import httpx -import litellm import pytest from pydantic import ValidationError +import litellm from litellm.llms.replicate.chat.handler import async_completion from litellm.llms.replicate.chat.handler import completion as replicate_completion from litellm.llms.replicate.chat.transformation import ReplicateConfig +from litellm.llms.replicate.common_utils import ReplicateError from litellm.types.utils import ModelResponse @@ -59,6 +60,106 @@ def test_transform_response_rejects_an_output_that_is_not_made_of_strings_withou assert "secret-output" not in str(exc_info.value) +def test_transform_response_reads_a_succeeded_prediction_that_carries_unmodelled_fields() -> None: + response = _transform( + httpx.Response( + 200, + json={ + "id": "p1", + "status": "succeeded", + "output": ["Hello", " there"], + "metrics": {"predict_time": 1.5, "stages": [1, {"name": None}]}, + "urls": {"get": "https://api.replicate.com/v1/predictions/p1"}, + "error": None, + }, + ) + ) + + assert response.choices[0].message.content == "Hello there" + assert response.model == "replicate/acme/echo-model" + + +@pytest.mark.parametrize( + ("body", "reported"), + [ + pytest.param( + {"status": "failed", "error": "boom", "output": None}, + "{'status': 'failed', 'error': 'boom', 'output': None}", + id="failed", + ), + pytest.param( + {"status": "processing", "output": ["partial"]}, + "{'status': 'processing', 'output': ['partial']}", + id="still-processing", + ), + pytest.param( + {"detail": {"b": [1, 2.5, True], "a": "z"}}, + "{'detail': {'b': [1, 2.5, True], 'a': 'z'}}", + id="no-status", + ), + pytest.param({}, "{}", id="empty-object"), + ], +) +def test_transform_response_reports_an_unsuccessful_prediction_body_as_received( + body: dict[str, object], reported: str +) -> None: + with pytest.raises(ReplicateError) as exc_info: + _transform(httpx.Response(500, json=body, headers={"retry-after": "3"})) + + assert exc_info.value.status_code == 422 + assert exc_info.value.message == f"LiteLLM Error - prediction not succeeded - {reported}" + assert exc_info.value.headers["retry-after"] == "3" + + +@pytest.mark.parametrize( + ("body", "echo"), + [ + pytest.param(["secret-output"], "secret-output", id="list"), + pytest.param("secret-output", "secret-output", id="string"), + pytest.param(4815162342, "4815162342", id="number"), + pytest.param(None, "NoneType", id="null"), + ], +) +def test_transform_response_rejects_a_body_that_is_not_an_object_without_echoing_it(body: object, echo: str) -> None: + with pytest.raises(ValidationError) as exc_info: + _transform(httpx.Response(500, content=json.dumps(body).encode())) + + assert echo not in str(exc_info.value) + + +@pytest.mark.parametrize( + ("supplied", "stored"), + [ + pytest.param({}, {}, id="nothing-supplied"), + pytest.param( + {"system_prompt": "", "max_new_tokens": 0, "temperature": 1, "debug": False}, + {"system_prompt": "", "max_new_tokens": 0, "temperature": 1, "debug": False}, + id="falsy-values-are-kept", + ), + pytest.param({"top_k": 40, "seed": None, "stop_sequences": None}, {"top_k": 40}, id="none-is-skipped"), + ], +) +def test_init_stores_each_supplied_value_except_none_as_class_level_config( + supplied: dict[str, object], stored: dict[str, object] +) -> None: + class ScopedReplicateConfig(ReplicateConfig): + pass + + ScopedReplicateConfig(**supplied) + + assert ScopedReplicateConfig.get_config() == stored + + +def test_init_called_again_with_none_keeps_the_value_stored_earlier() -> None: + class ScopedReplicateConfig(ReplicateConfig): + pass + + ScopedReplicateConfig(top_k=40) + ScopedReplicateConfig(top_k=None, seed=5) + + assert ScopedReplicateConfig.get_config() == {"top_k": 40, "seed": 5} + + @pytest.mark.asyncio @patch("litellm.llms.replicate.chat.handler.asyncio.sleep", new_callable=AsyncMock) @patch("litellm.llms.replicate.chat.handler.get_async_httpx_client") diff --git a/tests/unit/llms/sambanova/embedding/__init__.py b/tests/unit/llms/sambanova/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/sambanova/embedding/test_sambanova_embedding_transformation.py b/tests/unit/llms/sambanova/embedding/test_sambanova_embedding_transformation.py new file mode 100644 index 00000000000..99085214b30 --- /dev/null +++ b/tests/unit/llms/sambanova/embedding/test_sambanova_embedding_transformation.py @@ -0,0 +1,89 @@ +import datetime + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.sambanova.embedding.transformation import SambaNovaEmbeddingConfig +from litellm.types.utils import EmbeddingResponse + +EMBEDDING_ROW = {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]} + + +def _transform(payload: object) -> EmbeddingResponse: + logging_obj = Logging( + model="sambanova/E5-Mistral-7B-Instruct", + messages=[], + stream=False, + call_type="embedding", + start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-id", + function_id="function-id", + ) + return SambaNovaEmbeddingConfig().transform_embedding_response( + model="E5-Mistral-7B-Instruct", + raw_response=httpx.Response(200, json=payload), + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + api_key="sk-test", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + +def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]: + assert response.usage is not None + return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) + + +def test_sambanova_embedding_response_carries_provider_fields_and_usage(): + response = _transform( + { + "object": "list", + "model": "E5-Mistral-7B-Instruct", + "data": [EMBEDDING_ROW], + "usage": {"prompt_tokens": 8, "total_tokens": 9}, + } + ) + + assert response.model == "E5-Mistral-7B-Instruct" + assert response.object == "list" + assert response.data == [EMBEDDING_ROW] + assert _token_counts(response) == (8, 0, 9) + + +@pytest.mark.parametrize( + ("payload", "expected"), + [ + ({"data": [EMBEDDING_ROW]}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 4}}, (4, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 6}}, (0, 0, 6)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": None, "total_tokens": None}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": "", "total_tokens": []}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": "12", "total_tokens": 12.0}}, (12, 0, 12)), + ( + {"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 3, "total_tokens": 5, "completion_tokens": 2}}, + (3, 0, 5), + ), + ], +) +def test_sambanova_embedding_usage_tolerates_missing_null_and_loosely_typed_counts(payload, expected): + response = _transform(payload) + + assert response.data == [EMBEDDING_ROW] + assert _token_counts(response) == expected + + +@pytest.mark.parametrize("usage", [None, [], "tokens", 3]) +def test_sambanova_embedding_usage_that_is_not_an_object_is_rejected(usage): + with pytest.raises(ValidationError): + _transform({"data": [EMBEDDING_ROW], "usage": usage}) + + +@pytest.mark.parametrize("usage", [{"prompt_tokens": "many"}, {"total_tokens": 1.5}, {"prompt_tokens": [1]}]) +def test_sambanova_embedding_token_count_that_is_not_an_integer_is_rejected(usage): + with pytest.raises(ValidationError): + _transform({"data": [EMBEDDING_ROW], "usage": usage}) diff --git a/tests/unit/llms/searchapi/search/test_searchapi_search_transformation.py b/tests/unit/llms/searchapi/search/test_searchapi_search_transformation.py index 117b194b732..8f205e6356e 100644 --- a/tests/unit/llms/searchapi/search/test_searchapi_search_transformation.py +++ b/tests/unit/llms/searchapi/search/test_searchapi_search_transformation.py @@ -5,9 +5,10 @@ from unittest.mock import Mock, patch import httpx import pytest import respx +from pydantic import ValidationError import litellm -from litellm.llms.base_llm.search.transformation import SearchResponse +from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult from litellm.llms.searchapi.search.transformation import SearchAPIConfig @@ -222,6 +223,97 @@ class TestSearchAPIConfig: assert "AND" in result +def _transform(payload: object) -> SearchResponse: + return SearchAPIConfig().transform_search_response(raw_response=httpx.Response(200, json=payload), logging_obj=None) + + +def test_searchapi_organic_results_keep_provider_order_fields_and_dates(): + response = _transform( + { + "search_metadata": {"id": "search_1", "status": "Success"}, + "search_information": {"total_results": 2}, + "organic_results": [ + { + "position": 1, + "title": "LiteLLM", + "link": "https://example.com/litellm", + "snippet": "Call every LLM API", + "date": "Jan 5, 2024", + "sitelinks": {"inline": [{"title": "Docs", "link": "https://example.com/docs"}]}, + }, + {"position": 2, "title": "Docs", "link": "https://example.com/docs", "snippet": "Docs", "date": None}, + ], + "related_searches": [{"query": "litellm proxy"}], + } + ) + + assert response.object == "search" + assert response.results == [ + SearchResult( + title="LiteLLM", url="https://example.com/litellm", snippet="Call every LLM API", date="Jan 5, 2024" + ), + SearchResult(title="Docs", url="https://example.com/docs", snippet="Docs", date=None), + ] + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"error": "Invalid API key"}, + {"search_metadata": {"status": "Success"}, "organic_results": []}, + {"organic_results": ""}, + {"organic_results": {}}, + ], +) +def test_searchapi_response_without_organic_results_is_empty(payload): + assert _transform(payload).results == [] + + +def test_searchapi_result_missing_optional_fields_defaults_to_empty_strings_and_no_date(): + assert _transform({"organic_results": [{"position": 1}]}).results == [ + SearchResult(title="", url="", snippet="", date=None, last_updated=None) + ] + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("title", None), + ("title", 7), + ("link", None), + ("link", ["https://example.com"]), + ("snippet", None), + ("snippet", {"text": "snippet"}), + ("date", 1704412800), + ("date", ["Jan 5, 2024"]), + ], +) +def test_searchapi_result_with_field_of_wrong_type_is_rejected(field, value): + result = {"title": "T", "link": "https://example.com", "snippet": "S", "date": "Jan 5, 2024", field: value} + + with pytest.raises(ValidationError): + _transform({"organic_results": [result]}) + + +def test_searchapi_result_with_several_fields_of_wrong_type_reports_every_field(): + result = {"title": 7, "link": None, "snippet": ["snippet"], "date": 1704412800} + + with pytest.raises(ValidationError) as exc_info: + _transform({"organic_results": [result]}) + + assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",), ("date",)] + + +def test_searchapi_malformed_result_is_rejected_without_naming_its_position(): + results = [{}] * 429 + ["not-an-object"] + + with pytest.raises(ValidationError) as exc_info: + _transform({"organic_results": results}) + + assert "429" not in str(exc_info.value) + + def test_search_request_and_result_mapping( respx_mock: respx.MockRouter, httpx_transport: None ) -> None: diff --git a/tests/unit/llms/searxng/search/test_searxng_search_transformation.py b/tests/unit/llms/searxng/search/test_searxng_search_transformation.py index 5961cec548d..5ca7976491f 100644 --- a/tests/unit/llms/searxng/search/test_searxng_search_transformation.py +++ b/tests/unit/llms/searxng/search/test_searxng_search_transformation.py @@ -10,7 +10,9 @@ from urllib.parse import parse_qs, urlparse import httpx import pytest +from pydantic import ValidationError +from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult from litellm.llms.searxng.search.transformation import SearXNGSearchConfig from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -322,6 +324,109 @@ class TestSearXNGSearchHeaders: assert self.config.get_http_method() == "GET" +def _transform(payload: object) -> SearchResponse: + raw_response = httpx.Response(200, json=payload, request=httpx.Request("GET", "https://searxng.example.com/search")) + return SearXNGSearchConfig().transform_search_response(raw_response=raw_response, logging_obj=None) + + +def test_searxng_results_keep_provider_order_and_ignore_engine_metadata(): + response = _transform( + { + "query": "litellm", + "number_of_results": 0, + "results": [ + { + "url": "https://example.com/litellm", + "title": "LiteLLM", + "content": "Call every LLM API", + "publishedDate": "2025-01-15T00:00:00", + "pubdate": "2025-01-10", + "thumbnail": None, + "engines": ["google", "bing"], + "positions": [1, 2], + "score": 4.0, + }, + {"url": "https://example.com/docs", "title": "Docs", "content": "Docs", "publishedDate": None}, + ], + "answers": [], + "infoboxes": [], + "unresponsive_engines": [["brave", "timeout"]], + } + ) + + assert response.object == "search" + assert response.results == [ + SearchResult( + title="LiteLLM", + url="https://example.com/litellm", + snippet="Call every LLM API", + date="2025-01-15T00:00:00", + ), + SearchResult(title="Docs", url="https://example.com/docs", snippet="Docs", date=None), + ] + + +@pytest.mark.parametrize( + ("dates", "expected"), + [ + ({"publishedDate": "2025-01-15", "pubdate": "2025-01-10"}, "2025-01-15"), + ({"publishedDate": None, "pubdate": "2025-01-10"}, "2025-01-10"), + ({"publishedDate": "", "pubdate": "2025-01-10"}, "2025-01-10"), + ({"publishedDate": 0, "pubdate": "2025-01-10"}, "2025-01-10"), + ({"publishedDate": None, "pubdate": None}, None), + ({"publishedDate": "", "pubdate": ""}, ""), + ({"publishedDate": 0}, None), + ({}, None), + ], +) +def test_searxng_date_prefers_published_date_and_falls_back_to_pubdate_when_it_is_empty(dates, expected): + response = _transform({"results": [{"title": "T", "url": "https://example.com", "content": "C", **dates}]}) + + assert response.results == [SearchResult(title="T", url="https://example.com", snippet="C", date=expected)] + + +@pytest.mark.parametrize("payload", [{}, {"results": ""}, {"results": {}}]) +def test_searxng_response_with_empty_non_list_results_is_empty(payload): + assert _transform(payload).results == [] + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("title", None), + ("title", 7), + ("url", None), + ("content", None), + ("content", ["snippet"]), + ("publishedDate", 1736899200), + ("pubdate", ["2025-01-10"]), + ], +) +def test_searxng_result_with_field_of_wrong_type_is_rejected(field, value): + result = {"title": "T", "url": "https://example.com", "content": "C", field: value} + + with pytest.raises(ValidationError): + _transform({"results": [result]}) + + +def test_searxng_result_with_several_fields_of_wrong_type_reports_every_field(): + result = {"title": 7, "url": None, "content": ["snippet"], "publishedDate": 1736899200} + + with pytest.raises(ValidationError) as exc_info: + _transform({"results": [result]}) + + assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",), ("date",)] + + +def test_searxng_malformed_result_is_rejected_without_naming_its_position(): + results = [{}] * 429 + ["not-an-object"] + + with pytest.raises(ValidationError) as exc_info: + _transform({"results": results}) + + assert "429" not in str(exc_info.value) + + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/llms/serper/search/test_serper_search_transformation.py b/tests/unit/llms/serper/search/test_serper_search_transformation.py index acc9c634c47..92b6984f7ca 100644 --- a/tests/unit/llms/serper/search/test_serper_search_transformation.py +++ b/tests/unit/llms/serper/search/test_serper_search_transformation.py @@ -5,9 +5,13 @@ Tests for Serper Search API integration. import os from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +from pydantic import ValidationError import litellm +from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult +from litellm.llms.serper.search.transformation import SerperSearchConfig from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -194,6 +198,85 @@ class TestSerperSearch: assert len(response.results) == 0 +def _transform(payload: object) -> SearchResponse: + return SerperSearchConfig().transform_search_response( + raw_response=httpx.Response(200, json=payload), logging_obj=None + ) + + +def test_serper_organic_results_keep_provider_order_and_ignore_other_sections(): + response = _transform( + { + "searchParameters": {"q": "litellm", "type": "search", "engine": "google"}, + "knowledgeGraph": {"title": "LiteLLM", "attributes": {"Founded": "2023"}}, + "organic": [ + { + "title": "LiteLLM", + "link": "https://example.com/litellm", + "snippet": "Call every LLM API", + "date": "Jan 15, 2025", + "position": 1, + "sitelinks": [{"title": "Docs", "link": "https://example.com/docs"}], + }, + {"title": "Docs", "link": "https://example.com/docs", "snippet": "Docs", "position": 2}, + {"title": "Blog", "link": "https://example.com/blog", "snippet": "Blog", "date": None}, + ], + "peopleAlsoAsk": [{"question": "What is LiteLLM?", "title": "FAQ", "link": "https://example.com/faq"}], + "credits": 1, + } + ) + + assert response.object == "search" + assert response.results == [ + SearchResult( + title="LiteLLM", url="https://example.com/litellm", snippet="Call every LLM API", date="Jan 15, 2025" + ), + SearchResult(title="Docs", url="https://example.com/docs", snippet="Docs"), + SearchResult(title="Blog", url="https://example.com/blog", snippet="Blog"), + ] + + +def test_serper_result_without_optional_fields_gets_empty_defaults(): + assert _transform({"organic": [{"position": 1}]}).results == [SearchResult(title="", url="", snippet="")] + + +@pytest.mark.parametrize( + "payload", + [{}, {"organic": []}, {"organic": ""}, {"organic": {}}, {"message": "Unauthorized.", "statusCode": 403}], +) +def test_serper_response_without_organic_results_is_empty(payload): + assert _transform(payload).results == [] + + +@pytest.mark.parametrize( + ("field", "value"), + [("title", None), ("title", 7), ("link", None), ("snippet", ["snippet"]), ("date", 1736899200)], +) +def test_serper_result_with_field_of_wrong_type_is_rejected(field, value): + result = {"title": "T", "link": "https://example.com", "snippet": "S", field: value} + + with pytest.raises(ValidationError): + _transform({"organic": [result]}) + + +def test_serper_result_with_several_fields_of_wrong_type_reports_every_field(): + result = {"title": 7, "link": None, "snippet": ["snippet"], "date": 1736899200} + + with pytest.raises(ValidationError) as exc_info: + _transform({"organic": [result]}) + + assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",), ("date",)] + + +def test_serper_malformed_result_is_rejected_without_naming_its_position(): + results = [{}] * 429 + ["not-an-object"] + + with pytest.raises(ValidationError) as exc_info: + _transform({"organic": results}) + + assert "429" not in str(exc_info.value) + + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/llms/tavily/search/test_tavily_search_transformation.py b/tests/unit/llms/tavily/search/test_tavily_search_transformation.py index bd43588fb96..f66aaffcf48 100644 --- a/tests/unit/llms/tavily/search/test_tavily_search_transformation.py +++ b/tests/unit/llms/tavily/search/test_tavily_search_transformation.py @@ -5,9 +5,13 @@ Tests for Tavily Search API integration. import os from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +from pydantic import ValidationError import litellm +from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult +from litellm.llms.tavily.search.transformation import TavilySearchConfig from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -90,6 +94,93 @@ class TestTavilySearch: assert first_result.snippet == "This is a test snippet for result 1" +def _transform(payload: object) -> SearchResponse: + return TavilySearchConfig().transform_search_response( + raw_response=httpx.Response(200, json=payload), logging_obj=None + ) + + +def test_tavily_results_keep_provider_order_and_ignore_answer_and_usage_metadata(): + response = _transform( + { + "query": "litellm", + "answer": "LiteLLM is a gateway.", + "follow_up_questions": None, + "images": [{"url": "https://example.com/image.png", "description": "logo"}], + "results": [ + { + "title": "LiteLLM", + "url": "https://example.com/litellm", + "content": "Call every LLM API", + "score": 0.81, + "raw_content": None, + "favicon": "https://example.com/favicon.ico", + }, + {"title": "Docs", "url": "https://example.com/docs", "content": "Docs", "raw_content": "Raw docs"}, + ], + "response_time": 1.67, + "usage": {"credits": 1}, + "request_id": "123e4567-e89b-12d3-a456-426614174111", + } + ) + + assert response.object == "search" + assert response.results == [ + SearchResult(title="LiteLLM", url="https://example.com/litellm", snippet="Call every LLM API"), + SearchResult(title="Docs", url="https://example.com/docs", snippet="Docs"), + ] + + +def test_tavily_result_without_title_url_or_content_gets_empty_strings(): + assert _transform({"results": [{"score": 0.5, "raw_content": None}]}).results == [ + SearchResult(title="", url="", snippet="", date=None, last_updated=None) + ] + + +@pytest.mark.parametrize( + "payload", + [{}, {"query": "litellm", "answer": None}, {"results": []}, {"results": ""}, {"results": {}}], +) +def test_tavily_response_without_results_is_empty(payload): + assert _transform(payload).results == [] + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("title", None), + ("title", 7), + ("url", None), + ("url", ["https://example.com"]), + ("content", None), + ("content", {"text": "snippet"}), + ], +) +def test_tavily_result_with_field_of_wrong_type_is_rejected(field, value): + result = {"title": "T", "url": "https://example.com", "content": "S", field: value} + + with pytest.raises(ValidationError): + _transform({"results": [result]}) + + +def test_tavily_result_with_several_fields_of_wrong_type_reports_every_field(): + result = {"title": 7, "url": None, "content": ["snippet"]} + + with pytest.raises(ValidationError) as exc_info: + _transform({"results": [result]}) + + assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",)] + + +def test_tavily_malformed_result_is_rejected_without_naming_its_position(): + results = [{}] * 429 + ["not-an-object"] + + with pytest.raises(ValidationError) as exc_info: + _transform({"results": results}) + + assert "429" not in str(exc_info.value) + + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/llms/vertex_ai/fine_tuning/test_handler.py b/tests/unit/llms/vertex_ai/fine_tuning/test_handler.py index ce9def6ab1d..5447d553f30 100644 --- a/tests/unit/llms/vertex_ai/fine_tuning/test_handler.py +++ b/tests/unit/llms/vertex_ai/fine_tuning/test_handler.py @@ -3,7 +3,11 @@ import json from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch +import google.auth +import google.auth.credentials +import httpx import pytest +import respx import litellm from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -572,3 +576,105 @@ def _restore_litellm_state(state) -> None: for attr, value in state.items(): if hasattr(litellm, attr): setattr(litellm, attr, value) + + +_VERTEX_TUNING_JOB = { + "name": "projects/test-project/locations/us-central1/tuningJobs/42", + "tunedModelDisplayName": "gemini-1.0-pro-002-tuned", + "baseModel": "gemini-1.0-pro-002", + "supervisedTuningSpec": {"trainingDatasetUri": "gs://bucket/train.jsonl"}, + "state": "JOB_STATE_RUNNING", + "createTime": "2024-12-31T22:40:20.211140Z", + "updateTime": "2024-12-31T22:40:20.211140Z", +} +_BOTH_CALL_STYLES = pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) + + +class _StaticGoogleCredentials(google.auth.credentials.Credentials): + def refresh(self, request): + self.token = "test-token" + + +@pytest.fixture +def vertex_tuning_endpoint(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> respx.Route: + monkeypatch.setattr(google.auth, "default", lambda scopes: (_StaticGoogleCredentials(), "test-project")) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + return respx_mock.post(host="us-central1-aiplatform.googleapis.com") + + +def _vertex_json_response(payload: object) -> httpx.Response: + return httpx.Response(200, content=json.dumps(payload).encode()) + + +async def _created_vertex_tuning_job(run_async: bool, endpoint: respx.Route, payload: object): + endpoint.mock(return_value=_vertex_json_response(payload)) + pending = VertexFineTuningAPI().create_fine_tuning_job( + _is_async=run_async, + create_fine_tuning_job_data=FineTuningJobCreate( + training_file="gs://bucket/train.jsonl", model="gemini-1.0-pro-002" + ), + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials=None, + api_base=None, + timeout=600.0, + ) + return await pending if run_async else pending + + +@_BOTH_CALL_STYLES +@pytest.mark.asyncio() +async def test_vertex_fine_tuning_job_creation_returns_the_tuning_job_as_an_openai_job( + run_async: bool, vertex_tuning_endpoint: respx.Route +): + job = await _created_vertex_tuning_job(run_async, vertex_tuning_endpoint, _VERTEX_TUNING_JOB) + + assert (job.id, job.created_at, job.fine_tuned_model, job.model, job.status, job.training_file) == ( + "projects/test-project/locations/us-central1/tuningJobs/42", + 1735684820, + "gemini-1.0-pro-002-tuned", + "gemini-1.0-pro-002", + "running", + "gs://bucket/train.jsonl", + ) + + +@_BOTH_CALL_STYLES +@pytest.mark.parametrize("body", [[_VERTEX_TUNING_JOB], [["state", "JOB_STATE_RUNNING"]], "state", 403, 2.5, None]) +@pytest.mark.asyncio() +async def test_vertex_fine_tuning_job_creation_reports_a_tuning_job_body_that_is_not_an_object_as_a_type_error( + run_async: bool, vertex_tuning_endpoint: respx.Route, body: object +): + with pytest.raises(TypeError, match=r"argument after \*\* must be a mapping"): + await _created_vertex_tuning_job(run_async, vertex_tuning_endpoint, body) + + +@pytest.mark.parametrize( + "body", + [ + _VERTEX_TUNING_JOB, + {"candidates": [{"content": {"parts": [{"text": "hello"}]}}], "usageMetadata": {"totalTokenCount": 3}}, + {}, + [{"a": 1}, 2, "x", None], + "text", + 403, + 2.5, + True, + None, + ], +) +@pytest.mark.asyncio() +async def test_vertex_passthrough_post_returns_the_upstream_json_body_unchanged( + vertex_tuning_endpoint: respx.Route, body: object +): + vertex_tuning_endpoint.mock(return_value=_vertex_json_response(body)) + + returned = await VertexFineTuningAPI().pass_through_vertex_ai_POST_request( + request_data={"baseModel": "gemini-1.0-pro-002"}, + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials=None, + request_route="/tuningJobs", + ) + + assert (type(returned), returned) == (type(body), body) diff --git a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py index d3e94e5aa29..756e6441f1f 100644 --- a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py +++ b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py @@ -1,13 +1,18 @@ import base64 import json import os +from datetime import datetime from io import BytesIO from typing import Dict from unittest.mock import MagicMock, patch import httpx import pytest +from pydantic import ValidationError +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.vertex_ai.image_edit.vertex_gemini_transformation import ( VertexAIGeminiImageEditConfig, ) @@ -410,3 +415,184 @@ class TestVertexAIImagenImageEditTransformation: self.config._read_all_bytes(bytearray(b"test_bytearray")) == b"test_bytearray" ) + + +_IMAGEN_EDIT_IMAGE_KEY = "bytesBase64Encoded" + + +def _imagen_edit_predict_response(payload: object) -> httpx.Response: + return httpx.Response(200, content=json.dumps(payload).encode(), headers={"content-type": "application/json"}) + + +def _imagen_edit_response_images(payload: object) -> list[tuple[object, object]]: + response = VertexAIImagenImageEditConfig().transform_image_edit_response( + model="imagen-3.0-capability-001", + raw_response=_imagen_edit_predict_response(payload), + logging_obj=Logging( + model="imagen-3.0-capability-001", + messages=[], + stream=False, + call_type="image_edit", + start_time=datetime(2026, 1, 1), + litellm_call_id="imagen-edit-call", + function_id="imagen-edit-function", + ), + ) + return [(image.b64_json, image.url) for image in response.data] + + +def _imagen_edit_rejection_text(payload: object) -> str: + with pytest.raises(ValidationError) as rejection: + _imagen_edit_response_images(payload) + return str(rejection.value) + + +def _edited_imagen_images(payload: object) -> list[tuple[object, object]]: + transport = httpx.MockTransport(lambda request: _imagen_edit_predict_response(payload)) + response = litellm.image_edit( + model="vertex_ai/imagen-3.0-capability-001", + image=BytesIO(b"image-bytes"), + prompt="add a hat", + api_base="https://vertex.invalid", + vertex_project="test-project", + vertex_location="us-central1", + client=HTTPHandler(client=httpx.Client(transport=transport)), + ) + return [(image.b64_json, image.url) for image in response.data] + + +def test_imagen_image_edit_response_lists_prediction_images_in_order_and_skips_entries_without_image_bytes(): + images = _imagen_edit_response_images( + { + "predictions": [ + {_IMAGEN_EDIT_IMAGE_KEY: "first", "mimeType": "image/png"}, + "plain text", + [], + {}, + ["other"], + {"raiFilteredReason": "blocked"}, + {_IMAGEN_EDIT_IMAGE_KEY: "second"}, + ], + "deployedModelId": "1", + } + ) + + assert images == [("first", None), ("second", None)] + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"predictions": []}, + {"predictions": ""}, + {"predictions": "plain text"}, + {"predictions": _IMAGEN_EDIT_IMAGE_KEY}, + {"predictions": {}}, + {"predictions": {"mimeType": "image/png"}}, + {"predictions": [{}]}, + {"predictions": [{"mimeType": "image/png"}]}, + {"predictions": ["plain", ""]}, + {"predictions": [["a", "b"], []]}, + {"predictions": [[[_IMAGEN_EDIT_IMAGE_KEY]]]}, + {"Predictions": [{_IMAGEN_EDIT_IMAGE_KEY: "a"}]}, + {"predictions": [{"bytes_base64_encoded": "a"}]}, + ], +) +def test_imagen_image_edit_response_without_prediction_image_bytes_has_no_images(payload: object): + assert _imagen_edit_response_images(payload) == [] + + +@pytest.mark.parametrize("image_bytes", [None, "", "aGVsbG8=", "1.5", "true", "é中"]) +def test_imagen_image_edit_response_keeps_the_image_bytes_value_exactly_as_sent(image_bytes: str | None): + images = _imagen_edit_response_images({"predictions": [{_IMAGEN_EDIT_IMAGE_KEY: image_bytes}]}) + + assert repr(images) == repr([(image_bytes, None)]) + + +def test_imagen_image_edit_response_finds_the_image_after_hundreds_of_imageless_predictions(): + images = _imagen_edit_response_images({"predictions": [*[{}] * 403, {_IMAGEN_EDIT_IMAGE_KEY: "late"}]}) + + assert images == [("late", None)] + + +@pytest.mark.parametrize("image_bytes", [5, 1.5, True, ["a"], {"a": 1}]) +def test_imagen_image_edit_response_rejects_non_string_image_bytes_naming_b64_json(image_bytes: object): + with pytest.raises(ValidationError, match="b64_json"): + _imagen_edit_response_images({"predictions": [{_IMAGEN_EDIT_IMAGE_KEY: image_bytes}]}) + + +@pytest.mark.parametrize( + "payload", + [ + {"predictions": None}, + {"predictions": 7}, + {"predictions": 1.5}, + {"predictions": True}, + {"predictions": {_IMAGEN_EDIT_IMAGE_KEY: "vertex-secret"}}, + {"predictions": [None]}, + {"predictions": [7]}, + {"predictions": [1.5]}, + {"predictions": [False]}, + {"predictions": ["vertex-secret " + _IMAGEN_EDIT_IMAGE_KEY]}, + {"predictions": [[_IMAGEN_EDIT_IMAGE_KEY, "vertex-secret"]]}, + {"predictions": [{_IMAGEN_EDIT_IMAGE_KEY: "vertex-secret"}, None]}, + ], +) +def test_imagen_image_edit_response_rejects_malformed_predictions_without_echoing_them(payload: object): + rejection_text = _imagen_edit_rejection_text(payload) + + assert "vertex-secret" not in rejection_text and "input_value" not in rejection_text, rejection_text + + +@pytest.mark.parametrize("position", [403, 429]) +@pytest.mark.parametrize("malformed", [None, 7, _IMAGEN_EDIT_IMAGE_KEY, [_IMAGEN_EDIT_IMAGE_KEY]]) +def test_imagen_image_edit_response_rejection_text_is_the_same_wherever_the_malformed_prediction_sits( + position: int, malformed: object +): + rejection_at_start = _imagen_edit_rejection_text({"predictions": [malformed]}) + + rejection_at_position = _imagen_edit_rejection_text({"predictions": [*[{}] * position, malformed]}) + + assert rejection_at_position == rejection_at_start and str(position) not in rejection_at_position + + +@pytest.mark.parametrize("body", [[], [{_IMAGEN_EDIT_IMAGE_KEY: "a"}], "predictions", 7, 2.5, True, None]) +def test_imagen_image_edit_response_rejects_a_predict_body_that_is_not_an_object_with_attribute_error(body: object): + with pytest.raises(AttributeError, match="object has no attribute 'get'"): + _imagen_edit_response_images(body) + + +def test_imagen_image_edit_returns_the_prediction_images_of_the_predict_endpoint(): + images = _edited_imagen_images( + { + "predictions": [ + {_IMAGEN_EDIT_IMAGE_KEY: "first"}, + {"mimeType": "image/png"}, + {_IMAGEN_EDIT_IMAGE_KEY: "second"}, + ] + } + ) + + assert images == [("first", None), ("second", None)] + + +@pytest.mark.parametrize( + "payload", + [ + {"predictions": None}, + {"predictions": 7}, + {"predictions": [None]}, + {"predictions": ["vertex-secret " + _IMAGEN_EDIT_IMAGE_KEY]}, + {"predictions": [*[{}] * 403, [_IMAGEN_EDIT_IMAGE_KEY, "vertex-secret"]]}, + ], +) +def test_imagen_image_edit_maps_malformed_predictions_to_a_connection_error_without_echoing_them(payload: object): + with pytest.raises(litellm.APIConnectionError) as mapped: + _edited_imagen_images(payload) + + assert (type(mapped.value), mapped.value.status_code, "vertex-secret" in str(mapped.value)) == ( + litellm.APIConnectionError, + 500, + False, + ) diff --git a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index 43f6b47e9bd..5ac16283094 100644 --- a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -8,8 +8,9 @@ import pytest from google.oauth2.credentials import Credentials from pydantic import ValidationError +import litellm from litellm.litellm_core_utils.litellm_logging import Logging - +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.vertex_ai.image_generation import ( get_vertex_ai_image_generation_config, ) @@ -715,6 +716,293 @@ def test_gemini_image_generation_response_rejects_non_object_usage_metadata_with assert "input_value" not in str(exc_info.value) +@pytest.mark.parametrize( + "payload", + [ + {}, + {"candidates": []}, + {"candidates": ""}, + {"candidates": {}}, + {"candidates": [{}]}, + {"candidates": [{"finishReason": "SAFETY"}]}, + {"candidates": [{"content": {}}]}, + {"candidates": [{"content": {"parts": ""}}]}, + {"candidates": [{"content": {"parts": "plain text"}}]}, + {"candidates": [{"content": {"parts": {"text": "only text"}}}]}, + {"candidates": [{"content": {"parts": [{"text": "only text"}, "plain text", ["text"], [], ""]}}]}, + {"candidates": [{"content": {"parts": [{"inlineData": {"mimeType": "image/png"}}]}}]}, + {"candidates": [{"content": {"parts": [{"inlineData": "no image"}, {"inlineData": ["mimeType"]}]}}]}, + {"candidates": [{"content": {"parts": [{"inlineData": ""}, {"inlineData": []}, {"inlineData": {}}]}}]}, + {"candidates": [{"content": {"parts": [{"inline_data": {"data": "snake-case-is-not-read"}}]}}]}, + ], +) +def test_gemini_image_generation_response_without_inline_image_data_has_no_images(payload: object): + assert _transform_gemini_response(payload).data == [] + + +def test_gemini_image_generation_response_keeps_image_order_and_thought_signatures(): + response = _transform_gemini_response( + { + "candidates": [ + { + "content": { + "parts": [ + {"text": "caption"}, + {"inlineData": {"mimeType": "image/png", "data": "first"}, "thoughtSignature": "sig-1"}, + "plain text", + ] + } + }, + { + "content": { + "parts": [ + {"inlineData": {"data": "second"}}, + {"inlineData": {"data": None}, "thoughtSignature": ""}, + {"inlineData": {"data": "fourth"}, "thoughtSignature": {"nested": ["sig"]}}, + ] + } + }, + ] + } + ) + + assert [(image.b64_json, image.provider_specific_fields) for image in response.data or []] == [ + ("first", {"thought_signature": "sig-1"}), + ("second", None), + (None, None), + ("fourth", {"thought_signature": {"nested": ["sig"]}}), + ] + + +@pytest.mark.parametrize( + "payload", + [ + ["vertex-secret"], + "vertex-secret", + {"candidates": None}, + {"candidates": 7}, + {"candidates": "vertex-secret"}, + {"candidates": ["vertex-secret"]}, + {"candidates": [{"content": None}]}, + {"candidates": [{"content": "vertex-secret"}]}, + {"candidates": [{"content": ["vertex-secret"]}]}, + {"candidates": [{"content": {"parts": None}}]}, + {"candidates": [{"content": {"parts": [None]}}]}, + {"candidates": [{"content": {"parts": [7]}}]}, + {"candidates": [{"content": {"parts": ["inlineData vertex-secret"]}}]}, + {"candidates": [{"content": {"parts": [["inlineData", "vertex-secret"]]}}]}, + {"candidates": [{"content": {"parts": {"inlineData": "vertex-secret"}}}]}, + {"candidates": [{"content": {"parts": [{"inlineData": None}]}}]}, + {"candidates": [{"content": {"parts": [{"inlineData": 7}]}}]}, + {"candidates": [{"content": {"parts": [{"inlineData": "data vertex-secret"}]}}]}, + {"candidates": [{"content": {"parts": [{"inlineData": ["data", "vertex-secret"]}]}}]}, + ], +) +def test_gemini_image_generation_response_rejects_malformed_candidates_without_echoing_them(payload: object): + with pytest.raises(ValidationError) as exc_info: + _transform_gemini_response(payload) + + assert "vertex-secret" not in str(exc_info.value) + + +def _gemini_rejection_text(payload: object) -> str: + with pytest.raises(ValidationError) as exc_info: + _transform_gemini_response(payload) + + return str(exc_info.value) + + +@pytest.mark.parametrize("position", [403, 429]) +@pytest.mark.parametrize("malformed_entry", [None, 7, "inlineData"]) +def test_gemini_image_generation_response_rejection_text_is_the_same_wherever_the_malformed_entry_sits( + position: int, malformed_entry: object +): + late_parts = [*[{"text": "caption"}] * position, malformed_entry] + late_candidates = [*[{"finishReason": "STOP"}] * position, malformed_entry] + + assert _gemini_rejection_text({"candidates": [{"content": {"parts": late_parts}}]}) == _gemini_rejection_text( + {"candidates": [{"content": {"parts": [malformed_entry]}}]} + ) + assert _gemini_rejection_text({"candidates": late_candidates}) == _gemini_rejection_text( + {"candidates": [malformed_entry]} + ) + + +_IMAGEN_IMAGE_KEY = "bytesBase64Encoded" + + +def _imagen_predict_response(payload: object) -> httpx.Response: + return httpx.Response(200, content=json.dumps(payload).encode(), headers={"content-type": "application/json"}) + + +def _transform_imagen_response(payload: object, model_response: ImageResponse) -> ImageResponse: + return VertexAIImagenImageGenerationConfig().transform_image_generation_response( + model="imagegeneration@006", + raw_response=_imagen_predict_response(payload), + model_response=model_response, + logging_obj=Logging( + model="imagegeneration@006", + messages=[], + stream=False, + call_type="image_generation", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="vertex-imagen-generation-test", + function_id="vertex-imagen-generation-test", + ), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def _imagen_rejection_text(payload: object) -> str: + with pytest.raises(ValidationError) as exc_info: + _transform_imagen_response(payload, ImageResponse()) + + return str(exc_info.value) + + +def _generate_imagen_image(payload: object) -> ImageResponse: + transport = httpx.MockTransport(lambda request: _imagen_predict_response(payload)) + return litellm.image_generation( + model="vertex_ai/imagegeneration@006", + prompt="a cat", + api_base="https://vertex.invalid/v1/predict", + client=HTTPHandler(client=httpx.Client(transport=transport)), + ) + + +def test_imagen_image_generation_response_appends_prediction_images_in_order_to_the_given_response(): + model_response = ImageResponse(data=[{"b64_json": "seeded", "url": None}]) + seeded_image = model_response.data[0] + + result = _transform_imagen_response( + { + "predictions": [ + {_IMAGEN_IMAGE_KEY: "first", "mimeType": "image/png"}, + {"raiFilteredReason": "filtered"}, + "plain text", + [], + {_IMAGEN_IMAGE_KEY: "second"}, + ], + "deployedModelId": "123", + }, + model_response, + ) + + assert result is model_response + assert result.data[0] is seeded_image + assert [(image.b64_json, image.url) for image in result.data] == [ + ("seeded", None), + ("first", None), + ("second", None), + ] + + +@pytest.mark.parametrize( + "payload", + [ + {}, + {"deployedModelId": "123"}, + {"predictions": []}, + {"predictions": ""}, + {"predictions": "plain text"}, + {"predictions": _IMAGEN_IMAGE_KEY}, + {"predictions": {}}, + {"predictions": {"mimeType": "image/png", "prompt": "a cat"}}, + {"predictions": [{}]}, + {"predictions": [{"mimeType": "image/png"}, {"raiFilteredReason": "filtered"}]}, + {"predictions": ["plain text", ""]}, + {"predictions": [["mimeType", "image/png"], []]}, + {"predictions": [[[_IMAGEN_IMAGE_KEY]]]}, + {"predictions": [{"bytes_base64_encoded": "snake-case-is-not-read"}]}, + ], +) +def test_imagen_image_generation_response_without_prediction_image_bytes_has_no_images(payload: object): + assert _transform_imagen_response(payload, ImageResponse()).data == [] + + +@pytest.mark.parametrize("b64_value", [None, "", "aGVsbG8=", "1.5", "true"]) +def test_imagen_image_generation_response_keeps_the_image_bytes_value_exactly_as_sent(b64_value: object): + result = _transform_imagen_response({"predictions": [{_IMAGEN_IMAGE_KEY: b64_value}]}, ImageResponse()) + + assert [repr(image.b64_json) for image in result.data] == [repr(b64_value)] + + +@pytest.mark.parametrize("b64_value", [5, 1.5, True, ["aGVsbG8="], {"data": "aGVsbG8="}]) +def test_imagen_image_generation_response_rejects_non_string_image_bytes_naming_b64_json(b64_value: object): + with pytest.raises(ValidationError, match="b64_json"): + _transform_imagen_response({"predictions": [{_IMAGEN_IMAGE_KEY: b64_value}]}, ImageResponse()) + + +@pytest.mark.parametrize( + "payload", + [ + ["vertex-secret"], + "vertex-secret", + 7, + None, + {"predictions": None}, + {"predictions": 7}, + {"predictions": True}, + {"predictions": [None]}, + {"predictions": [7]}, + {"predictions": [False]}, + {"predictions": [f"{_IMAGEN_IMAGE_KEY} vertex-secret"]}, + {"predictions": [[_IMAGEN_IMAGE_KEY, "vertex-secret"]]}, + {"predictions": {_IMAGEN_IMAGE_KEY: "vertex-secret"}}, + {"predictions": [{_IMAGEN_IMAGE_KEY: "kept"}, None]}, + ], +) +def test_imagen_image_generation_response_rejects_malformed_predictions_without_echoing_them(payload: object): + rejection_text = _imagen_rejection_text(payload) + + assert "vertex-secret" not in rejection_text + assert "input_value" not in rejection_text + + +@pytest.mark.parametrize("position", [403, 429]) +@pytest.mark.parametrize("malformed_entry", [None, 7, _IMAGEN_IMAGE_KEY, [_IMAGEN_IMAGE_KEY]]) +def test_imagen_image_generation_response_rejection_text_is_the_same_wherever_the_malformed_prediction_sits( + position: int, malformed_entry: object +): + late_predictions = [*[{_IMAGEN_IMAGE_KEY: "kept"}] * position, malformed_entry] + + assert _imagen_rejection_text({"predictions": late_predictions}) == _imagen_rejection_text( + {"predictions": [malformed_entry]} + ) + + +def test_imagen_image_generation_returns_the_prediction_images_of_the_predict_endpoint(): + result = _generate_imagen_image( + {"predictions": [{_IMAGEN_IMAGE_KEY: "first"}, {"mimeType": "image/png"}, {_IMAGEN_IMAGE_KEY: "second"}]} + ) + + assert [(image.b64_json, image.url) for image in result.data] == [("first", None), ("second", None)] + + +@pytest.mark.parametrize( + "payload", + [ + ["vertex-secret"], + "vertex-secret", + {"predictions": None}, + {"predictions": [*[{_IMAGEN_IMAGE_KEY: "kept"}] * 403, f"{_IMAGEN_IMAGE_KEY} vertex-secret"]}, + {"predictions": [*[{_IMAGEN_IMAGE_KEY: "kept"}] * 429, None]}, + ], +) +def test_imagen_image_generation_maps_a_malformed_predict_body_to_a_connection_error_without_echoing_it( + payload: object, +): + with pytest.raises(litellm.APIConnectionError) as exc_info: + _generate_imagen_image(payload) + + assert type(exc_info.value) is litellm.APIConnectionError + assert exc_info.value.status_code == 500 + assert "vertex-secret" not in str(exc_info.value) + + def test_gemini_image_generation_validation_uses_supplied_vertex_token() -> None: config: Final = VertexAIGeminiImageGenerationConfig() credentials_json: Final = json.dumps({"type": "authorized_user", "client_id": "test-client"}) diff --git a/tests/unit/llms/vertex_ai/text_to_speech/test_text_to_speech_handler.py b/tests/unit/llms/vertex_ai/text_to_speech/test_text_to_speech_handler.py new file mode 100644 index 00000000000..4a57316d5b3 --- /dev/null +++ b/tests/unit/llms/vertex_ai/text_to_speech/test_text_to_speech_handler.py @@ -0,0 +1,155 @@ +import base64 +import binascii +import json +from datetime import datetime +from typing import Final + +import google.auth +import google.auth.credentials +import httpx +import pytest +import respx +from pydantic import ValidationError + +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.vertex_ai.text_to_speech.text_to_speech_handler import VertexTextToSpeechAPI + +SYNTHESIZE_URL: Final = "https://texttospeech.googleapis.com/v1/text:synthesize" +AUDIO: Final = b"hello audio" +ENCODED_AUDIO: Final = base64.b64encode(AUDIO).decode() +BOTH_CALL_STYLES: Final = pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) + + +class _StaticGoogleCredentials(google.auth.credentials.Credentials): + def refresh(self, request: object) -> None: + self.token = "test-token" + + +@pytest.fixture +def synthesize_endpoint(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> respx.Route: + monkeypatch.setattr(google.auth, "default", lambda scopes: (_StaticGoogleCredentials(), "test-project")) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + return respx_mock.post(SYNTHESIZE_URL) + + +async def _synthesized_audio(run_async: bool, endpoint: respx.Route, payload: object) -> bytes: + endpoint.mock(return_value=httpx.Response(200, content=json.dumps(payload).encode())) + pending: Final = VertexTextToSpeechAPI().audio_speech( + logging_obj=Logging( + model="vertex_ai/", + messages=[], + stream=False, + call_type="speech", + start_time=datetime(2026, 1, 1), + litellm_call_id="speech-call", + function_id="speech-function", + ), + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials=None, + api_base=None, + timeout=10.0, + model="vertex_ai/", + input="hello world", + _is_async=run_async, + ) + speech: Final = await pending if run_async else pending + return speech.content + + +@BOTH_CALL_STYLES +@pytest.mark.parametrize( + ("payload", "audio"), + [ + ({"audioContent": ENCODED_AUDIO}, AUDIO), + ({"audioContent": ENCODED_AUDIO, "timepoints": [], "audioConfig": {"audioEncoding": "MP3"}}, AUDIO), + ({"audioContent": ENCODED_AUDIO[:8] + "\n" + ENCODED_AUDIO[8:] + "\n"}, AUDIO), + ({"audioContent": ""}, b""), + ], +) +async def test_vertex_speech_returns_the_decoded_audio_content( + run_async: bool, synthesize_endpoint: respx.Route, payload: object, audio: bytes +) -> None: + assert await _synthesized_audio(run_async, synthesize_endpoint, payload) == audio + + +@BOTH_CALL_STYLES +async def test_vertex_speech_sends_the_default_voice_request_with_the_google_token( + run_async: bool, synthesize_endpoint: respx.Route +) -> None: + await _synthesized_audio(run_async, synthesize_endpoint, {"audioContent": ENCODED_AUDIO}) + + sent: Final = synthesize_endpoint.calls.last.request + assert (json.loads(sent.content), sent.headers["authorization"], sent.headers["x-goog-user-project"]) == ( + { + "input": {"text": "hello world"}, + "voice": {"languageCode": "en-US", "name": "en-US-Studio-O"}, + "audioConfig": {"audioEncoding": "LINEAR16", "speakingRate": "1"}, + }, + "Bearer test-token", + "test-project", + ) + + +@BOTH_CALL_STYLES +@pytest.mark.parametrize( + "body", + [[], [{"audioContent": ENCODED_AUDIO}], "audioContent", 7, 403, 2.5, True, None, ["vertex-secret"]], +) +async def test_vertex_speech_rejects_a_synthesize_body_that_is_not_an_object_without_echoing_it( + run_async: bool, synthesize_endpoint: respx.Route, body: object +) -> None: + with pytest.raises(ValidationError) as rejection: + await _synthesized_audio(run_async, synthesize_endpoint, body) + + rejection_text: Final = str(rejection.value) + assert ( + [error["type"] for error in rejection.value.errors()], + "vertex-secret" in rejection_text, + "403" in rejection_text, + ) == (["dict_type"], False, False), rejection_text + + +@BOTH_CALL_STYLES +@pytest.mark.parametrize("audio_content", [None, 5, 403, 1.5, True, [ENCODED_AUDIO], ["vertex-secret"], {"data": "x"}]) +async def test_vertex_speech_rejects_audio_content_that_is_not_a_string_without_echoing_it( + run_async: bool, synthesize_endpoint: respx.Route, audio_content: object +) -> None: + with pytest.raises(ValidationError) as rejection: + await _synthesized_audio(run_async, synthesize_endpoint, {"audioContent": audio_content}) + + rejection_text: Final = str(rejection.value) + assert ( + [error["type"] for error in rejection.value.errors()], + "vertex-secret" in rejection_text, + "403" in rejection_text, + ) == (["string_type"], False, False), rejection_text + + +@BOTH_CALL_STYLES +@pytest.mark.parametrize( + "payload", + [{}, {"audiocontent": ENCODED_AUDIO}, {"error": {"code": 403, "message": "denied", "status": "PERMISSION_DENIED"}}], +) +async def test_vertex_speech_reports_a_synthesize_body_without_audio_content_as_a_key_error( + run_async: bool, synthesize_endpoint: respx.Route, payload: object +) -> None: + with pytest.raises(KeyError, match="audioContent"): + await _synthesized_audio(run_async, synthesize_endpoint, payload) + + +@BOTH_CALL_STYLES +async def test_vertex_speech_reports_badly_padded_audio_content_with_the_base64_error( + run_async: bool, synthesize_endpoint: respx.Route +) -> None: + with pytest.raises(binascii.Error, match="Incorrect padding"): + await _synthesized_audio(run_async, synthesize_endpoint, {"audioContent": "abc"}) + + +@BOTH_CALL_STYLES +async def test_vertex_speech_reports_non_ascii_audio_content_with_the_base64_error( + run_async: bool, synthesize_endpoint: respx.Route +) -> None: + with pytest.raises(ValueError, match="string argument should contain only ASCII characters"): + await _synthesized_audio(run_async, synthesize_endpoint, {"audioContent": "é中"}) diff --git a/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py index 8ca0a4a4335..dc6f35fcfe5 100644 --- a/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py +++ b/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py @@ -11,6 +11,7 @@ import responses from pydantic import ValidationError import litellm +from litellm.litellm_core_utils.exception_mapping_utils import exception_type from litellm.llms.vertex_ai.text_to_speech.transformation import ( VertexAILyriaTextToSpeechConfig, VertexAITextToSpeechConfig, @@ -761,3 +762,170 @@ async def test_aspeech_vertex_ai_sends_ssml_input( "voice": _UK_VOICE, "audioConfig": _UK_AUDIO_CONFIG, } + + +_WAV_AUDIO_B64: Final = "UklGRiQAAABXQVZFZm10IA==" +_WAV_AUDIO: Final = b"RIFF$\x00\x00\x00WAVEfmt " +_UNLABELED_AUDIO_B64: Final = "bHlyaWEtMy1hdWRpbw==" +_UNLABELED_AUDIO: Final = b"lyria-3-audio" + + +@pytest.mark.parametrize( + ("model", "payload", "expected_audio", "expected_mime_type"), + [ + pytest.param( + "lyria-002", + { + "predictions": [ + {"audioContent": None, "bytesBase64Encoded": _WAV_AUDIO_B64, "mimeType": None}, + "trailing entry that is never read", + ], + "deployedModelId": "123", + }, + _WAV_AUDIO, + "audio/wav", + id="predict-null-fields-and-unread-trailing-prediction", + ), + pytest.param( + "lyria-002", + {"predictions": [{"audioContent": _WAV_AUDIO_B64, "mimeType": ""}]}, + _WAV_AUDIO, + "audio/wav", + id="predict-empty-mime-type-is-sniffed", + ), + pytest.param( + "lyria-002", + {"predictions": [{"audioContent": _WAV_AUDIO_B64, "mimeType": "audio/x-custom"}]}, + _WAV_AUDIO, + "audio/x-custom", + id="predict-declared-mime-type-wins-over-sniffing", + ), + pytest.param( + "lyria-002", + {"predictions": [{"audioContent": "", "bytesBase64Encoded": ""}]}, + b"", + "audio/mpeg", + id="predict-empty-audio-strings-yield-empty-audio", + ), + pytest.param( + "lyria-3-pro-preview", + { + "steps": [ + {"type": "thought", "content": "free text that is never read as content parts"}, + {"type": "model_output", "content": None}, + {"type": "model_output", "content": [{"type": "audio", "data": _WAV_AUDIO_B64, "mime_type": None}]}, + ] + }, + _WAV_AUDIO, + "audio/wav", + id="interactions-unread-step-content-and-null-fields", + ), + pytest.param( + "lyria-3-pro-preview", + {"steps": None, "outputs": [{"type": "audio", "data": _UNLABELED_AUDIO_B64, "mime_type": "audio/mpeg"}]}, + _UNLABELED_AUDIO, + "audio/mpeg", + id="interactions-null-steps-fall-back-to-outputs", + ), + pytest.param( + "lyria-3-clip-preview", + { + "outputs": [ + {"type": "audio", "data": 7, "mime_type": 5}, + {"type": "audio", "data": _UNLABELED_AUDIO_B64, "mime_type": "audio/mpeg"}, + {"type": "audio", "data": ""}, + ] + }, + _UNLABELED_AUDIO, + "audio/mpeg", + id="interactions-last-non-empty-audio-part-wins-over-an-earlier-malformed-one", + ), + ], +) +def test_lyria_transform_response_tolerates_optional_and_unread_provider_fields( + model: str, payload: dict[str, object], expected_audio: bytes, expected_mime_type: str +): + response: Final = VertexAILyriaTextToSpeechConfig().transform_text_to_speech_response( + model=model, + raw_response=httpx.Response(200, json=payload), + logging_obj=MagicMock(), + ) + + assert response.content == expected_audio + assert response.response.headers["content-type"] == expected_mime_type + + +@pytest.mark.parametrize( + ("model", "payload"), + [ + ("lyria-002", {}), + ("lyria-002", {"predictions": None}), + ("lyria-002", {"predictions": []}), + ("lyria-002", {"predictions": [{}]}), + ("lyria-002", {"predictions": [{"audioContent": "", "bytesBase64Encoded": None}]}), + ("lyria-3-pro-preview", {}), + ("lyria-3-pro-preview", {"steps": [], "outputs": None}), + ("lyria-3-pro-preview", {"steps": [{"type": "model_output", "content": []}]}), + ("lyria-3-pro-preview", {"outputs": [{"type": "audio", "data": ""}, {"type": "text", "data": "ignored"}]}), + ], +) +def test_lyria_transform_response_without_audio_reports_it_missing(model: str, payload: dict[str, object]): + with pytest.raises(ValueError, match=f"No generated audio found in Vertex AI {model} response"): + VertexAILyriaTextToSpeechConfig().transform_text_to_speech_response( + model=model, + raw_response=httpx.Response(200, json=payload), + logging_obj=MagicMock(), + ) + + +@pytest.mark.parametrize( + ("model", "payload"), + [ + ("lyria-002", ["not", "an", "object"]), + ("lyria-002", {"predictions": {"audioContent": _WAV_AUDIO_B64}}), + ("lyria-002", {"predictions": ["not-an-object"]}), + ("lyria-002", {"predictions": [{"audioContent": 7}]}), + ("lyria-002", {"predictions": [{"audioContent": _WAV_AUDIO_B64, "mimeType": ["audio/wav"]}]}), + ("lyria-3-pro-preview", ["not", "an", "object"]), + ("lyria-3-pro-preview", {"steps": ["not-an-object"]}), + ("lyria-3-pro-preview", {"steps": [{"type": "model_output", "content": ["not-an-object"]}]}), + ("lyria-3-pro-preview", {"outputs": [{"type": "audio", "data": [_WAV_AUDIO_B64]}]}), + ("lyria-3-pro-preview", {"outputs": [{"type": "audio", "data": _WAV_AUDIO_B64, "mime_type": 5}]}), + ], +) +def test_lyria_transform_response_rejects_malformed_payloads_without_echoing_them(model: str, payload: object): + with pytest.raises(ValidationError) as exc_info: + VertexAILyriaTextToSpeechConfig().transform_text_to_speech_response( + model=model, + raw_response=httpx.Response(200, json=payload), + logging_obj=MagicMock(), + ) + + assert "input_value" not in str(exc_info.value) + + +@pytest.mark.parametrize( + "payload", + [ + pytest.param({"steps": [{"type": "thought"}] * 403 + ["not-an-object"]}, id="step"), + pytest.param({"outputs": [{"type": "text"}] * 403 + [None]}, id="output"), + pytest.param( + {"steps": [{"type": "model_output", "content": [{"type": "text"}] * 403 + [7]}]}, + id="content-part", + ), + ], +) +def test_lyria_malformed_item_position_does_not_change_the_mapped_error(payload: dict[str, object]): + with pytest.raises(ValidationError) as exc_info: + VertexAILyriaTextToSpeechConfig().transform_text_to_speech_response( + model="lyria-3-pro-preview", + raw_response=httpx.Response(200, json=payload), + logging_obj=MagicMock(), + ) + + with pytest.raises(litellm.APIConnectionError): + exception_type( + model="vertex_ai/lyria-3-pro-preview", + custom_llm_provider="vertex_ai", + original_exception=exc_info.value, + ) diff --git a/tests/unit/llms/vertex_ai/vertex_embeddings/__init__.py b/tests/unit/llms/vertex_ai/vertex_embeddings/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/vertex_embeddings/test_embedding_handler.py b/tests/unit/llms/vertex_ai/vertex_embeddings/test_embedding_handler.py new file mode 100644 index 00000000000..4ceae8d407a --- /dev/null +++ b/tests/unit/llms/vertex_ai/vertex_embeddings/test_embedding_handler.py @@ -0,0 +1,289 @@ +import json +from datetime import datetime +from typing import Final + +import google.auth +import google.auth.credentials +import httpx +import pytest +import respx +from pydantic import ValidationError + +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.vertex_ai.vertex_embeddings.embedding_handler import VertexEmbedding +from litellm.types.utils import EmbeddingResponse + +GECKO: Final = "textembedding-gecko" +FINE_TUNED_ENDPOINT: Final = "1234567890" +BGE: Final = "bge-small-en-v1.5" +PUBLIC_CALL_PROJECT: Final = "vertex-embedding-handler-tests" +BOTH_CALL_STYLES: Final = pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +EVERY_MODEL_FAMILY: Final = pytest.mark.parametrize("model", [GECKO, FINE_TUNED_ENDPOINT, BGE]) +BODIES_THAT_ARE_NOT_OBJECTS: Final = [ + [], + [{"predictions": []}], + ["predictions"], + "predictions", + "", + 7, + 403, + 2.5, + True, + None, + ["vertex-secret"], +] + + +class _StaticGoogleCredentials(google.auth.credentials.Credentials): + def refresh(self, request: object) -> None: + self.token = "test-token" + + +@pytest.fixture +def predict_endpoint(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> respx.Route: + monkeypatch.setattr(google.auth, "default", lambda scopes: (_StaticGoogleCredentials(), "test-project")) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + return respx_mock.post(host="us-central1-aiplatform.googleapis.com") + + +def _predict_response(payload: object) -> httpx.Response: + return httpx.Response(200, content=json.dumps(payload).encode()) + + +def _embedding_logging() -> Logging: + return Logging( + model=GECKO, + messages=[], + stream=False, + call_type="embedding", + start_time=datetime(2026, 1, 1), + litellm_call_id="embedding-call", + function_id="embedding-function", + ) + + +def _gecko_prediction(values: object, token_count: int) -> dict[str, object]: + return {"embeddings": {"values": values, "statistics": {"token_count": token_count, "truncated": False}}} + + +def _embedding_rows(vectors: list[object]) -> list[dict[str, object]]: + return [{"object": "embedding", "index": index, "embedding": vector} for index, vector in enumerate(vectors)] + + +async def _embedding_response( + run_async: bool, endpoint: respx.Route, model: str, payload: object, logging_obj: Logging +) -> EmbeddingResponse: + endpoint.mock(return_value=_predict_response(payload)) + pending: Final = VertexEmbedding().embedding( + model=model, + input=["hello", "world"], + print_verbose=None, + model_response=EmbeddingResponse(), + optional_params={}, + logging_obj=logging_obj, + custom_llm_provider="vertex_ai", + timeout=10.0, + aembedding=run_async, + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials=None, + litellm_params={}, + ) + return await pending if run_async else pending + + +@BOTH_CALL_STYLES +@pytest.mark.parametrize( + ("model", "payload", "prompt_tokens"), + [ + ( + GECKO, + { + "predictions": [_gecko_prediction([0.1, 0.2, 0.3], 3), _gecko_prediction([1, 2.0, 3], 5)], + "metadata": {"billableCharacterCount": 10}, + }, + 8, + ), + (FINE_TUNED_ENDPOINT, {"predictions": [[[0.1, 0.2, 0.3]], [[1, 2.0, 3]]], "deployedModelId": "42"}, 0), + (BGE, {"predictions": [[0.1, 0.2, 0.3], [1, 2.0, 3]], "deployedModelId": "42"}, 0), + ], +) +async def test_vertex_embedding_returns_each_prediction_vector_exactly_as_sent( + run_async: bool, predict_endpoint: respx.Route, model: str, payload: object, prompt_tokens: int +) -> None: + response: Final = await _embedding_response(run_async, predict_endpoint, model, payload, _embedding_logging()) + + assert ( + json.dumps(response.data), + response.object, + response.model, + response.usage.prompt_tokens, + response.usage.completion_tokens, + response.usage.total_tokens, + ) == ( + json.dumps(_embedding_rows([[0.1, 0.2, 0.3], [1, 2.0, 3]])), + "list", + model, + prompt_tokens, + 0, + prompt_tokens, + ) + + +@BOTH_CALL_STYLES +@pytest.mark.parametrize("values", [[], "not-a-vector", None, [0.5, "x", None, True, {"a": 1}]]) +async def test_vertex_embedding_keeps_prediction_values_that_are_not_float_vectors( + run_async: bool, predict_endpoint: respx.Route, values: object +) -> None: + response: Final = await _embedding_response( + run_async, predict_endpoint, GECKO, {"predictions": [_gecko_prediction(values, 3)]}, _embedding_logging() + ) + + assert json.dumps(response.data) == json.dumps(_embedding_rows([values])) + + +@BOTH_CALL_STYLES +async def test_vertex_embedding_returns_every_vector_of_a_large_batch_in_order( + run_async: bool, predict_endpoint: respx.Route +) -> None: + vectors: Final[list[object]] = [[index, index + 0.5] for index in range(500)] + + response: Final = await _embedding_response( + run_async, predict_endpoint, BGE, {"predictions": vectors}, _embedding_logging() + ) + + assert json.dumps(response.data) == json.dumps(_embedding_rows(vectors)) + + +@BOTH_CALL_STYLES +async def test_vertex_embedding_logs_the_predict_body_as_json_text( + run_async: bool, predict_endpoint: respx.Route +) -> None: + payload: Final = {"predictions": [[0.1, 0.2]], "deployedModelId": "42"} + logging_obj: Final = _embedding_logging() + + await _embedding_response(run_async, predict_endpoint, BGE, payload, logging_obj) + + assert logging_obj.model_call_details["original_response"] == json.dumps(payload) + + +@BOTH_CALL_STYLES +@EVERY_MODEL_FAMILY +@pytest.mark.parametrize("body", BODIES_THAT_ARE_NOT_OBJECTS) +async def test_vertex_embedding_rejects_a_predict_body_that_is_not_an_object_without_echoing_it( + run_async: bool, predict_endpoint: respx.Route, model: str, body: object +) -> None: + with pytest.raises(ValidationError) as rejection: + await _embedding_response(run_async, predict_endpoint, model, body, _embedding_logging()) + + rejection_text: Final = str(rejection.value) + assert ( + [error["type"] for error in rejection.value.errors()], + "vertex-secret" in rejection_text, + "403" in rejection_text, + ) == (["dict_type"], False, False), rejection_text + + +@BOTH_CALL_STYLES +async def test_vertex_embedding_logs_a_predict_body_that_is_not_an_object_before_rejecting_it( + run_async: bool, predict_endpoint: respx.Route +) -> None: + logging_obj: Final = _embedding_logging() + + with pytest.raises(ValidationError): + await _embedding_response(run_async, predict_endpoint, GECKO, ["not", "an", "object"], logging_obj) + + assert logging_obj.model_call_details["original_response"] == ["not", "an", "object"] + + +@BOTH_CALL_STYLES +@EVERY_MODEL_FAMILY +@pytest.mark.parametrize( + "payload", + [{}, {"error": {"code": 403, "message": "denied", "status": "PERMISSION_DENIED"}}], +) +async def test_vertex_embedding_reports_a_predict_object_without_predictions_as_a_key_error( + run_async: bool, predict_endpoint: respx.Route, model: str, payload: object +) -> None: + with pytest.raises(KeyError, match="predictions"): + await _embedding_response(run_async, predict_endpoint, model, payload, _embedding_logging()) + + +def test_vertex_embedding_call_returns_the_prediction_vectors(predict_endpoint: respx.Route) -> None: + predict_endpoint.mock( + return_value=_predict_response({"predictions": [_gecko_prediction([0.1, 2], 3), _gecko_prediction([3.5], 4)]}) + ) + + response: Final = litellm.embedding( + model=f"vertex_ai/{GECKO}", + input=["hello", "world"], + vertex_project=PUBLIC_CALL_PROJECT, + vertex_location="us-central1", + ) + + assert (json.dumps(response.data), response.usage.prompt_tokens) == ( + json.dumps(_embedding_rows([[0.1, 2], [3.5]])), + 7, + ) + + +async def test_vertex_async_embedding_call_returns_the_prediction_vectors(predict_endpoint: respx.Route) -> None: + predict_endpoint.mock( + return_value=_predict_response({"predictions": [_gecko_prediction([0.1, 2], 3), _gecko_prediction([3.5], 4)]}) + ) + + response: Final = await litellm.aembedding( + model=f"vertex_ai/{GECKO}", + input=["hello", "world"], + vertex_project=PUBLIC_CALL_PROJECT, + vertex_location="us-central1", + ) + + assert (json.dumps(response.data), response.usage.prompt_tokens) == ( + json.dumps(_embedding_rows([[0.1, 2], [3.5]])), + 7, + ) + + +MALFORMED_PREDICT_BODIES: Final = [ + ["vertex-secret"], + "predictions", + 403, + None, + {"predictions": [_gecko_prediction([0.5], 1) for _ in range(403)] + [None]}, +] + + +@pytest.mark.parametrize("body", MALFORMED_PREDICT_BODIES) +def test_vertex_embedding_call_maps_a_malformed_predict_body_to_a_connection_error_without_echoing_it( + predict_endpoint: respx.Route, body: object +) -> None: + predict_endpoint.mock(return_value=_predict_response(body)) + + with pytest.raises(litellm.APIConnectionError) as failure: + litellm.embedding( + model=f"vertex_ai/{GECKO}", + input=["hello", "world"], + vertex_project=PUBLIC_CALL_PROJECT, + vertex_location="us-central1", + ) + + assert (failure.value.status_code, "vertex-secret" in str(failure.value)) == (500, False) + + +@pytest.mark.parametrize("body", MALFORMED_PREDICT_BODIES) +async def test_vertex_async_embedding_call_maps_a_malformed_predict_body_to_a_connection_error_without_echoing_it( + predict_endpoint: respx.Route, body: object +) -> None: + predict_endpoint.mock(return_value=_predict_response(body)) + + with pytest.raises(litellm.APIConnectionError) as failure: + await litellm.aembedding( + model=f"vertex_ai/{GECKO}", + input=["hello", "world"], + vertex_project=PUBLIC_CALL_PROJECT, + vertex_location="us-central1", + ) + + assert (failure.value.status_code, "vertex-secret" in str(failure.value)) == (500, False) diff --git a/tests/unit/llms/voyage/embedding/test_transformation_contextual.py b/tests/unit/llms/voyage/embedding/test_transformation_contextual.py index 9fd630283bd..915fe1ea48d 100644 --- a/tests/unit/llms/voyage/embedding/test_transformation_contextual.py +++ b/tests/unit/llms/voyage/embedding/test_transformation_contextual.py @@ -1,12 +1,18 @@ import asyncio +import datetime import importlib import json import os from unittest.mock import MagicMock, patch +import httpx import pytest +from pydantic import ValidationError import litellm +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.voyage.embedding.transformation_contextual import VoyageContextualEmbeddingConfig +from litellm.types.utils import EmbeddingResponse from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -473,3 +479,86 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), "cohere_key": getattr(litellm, "cohere_key", None), } + + +EMBEDDING_ROW = {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]} + + +def _transform(payload: object) -> EmbeddingResponse: + logging_obj = Logging( + model="voyage/voyage-context-3", + messages=[], + stream=False, + call_type="embedding", + start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-id", + function_id="function-id", + ) + return VoyageContextualEmbeddingConfig().transform_embedding_response( + model="voyage-context-3", + raw_response=httpx.Response(200, json=payload), + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + api_key="sk-test", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + +def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]: + assert response.usage is not None + return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) + + +def test_voyage_contextual_embedding_response_carries_provider_fields_and_usage(): + response = _transform( + { + "object": "list", + "model": "voyage-context-3", + "data": [EMBEDDING_ROW], + "usage": {"prompt_tokens": 8, "total_tokens": 9}, + } + ) + + assert response.model == "voyage-context-3" + assert response.object == "list" + assert response.data == [EMBEDDING_ROW] + assert _token_counts(response) == (9, 0, 9) + + +@pytest.mark.parametrize( + ("payload", "expected"), + [ + ({"data": [EMBEDDING_ROW]}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 4}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 6}}, (6, 0, 6)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": None}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": ""}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": []}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": "12"}}, (12, 0, 12)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 12.0}}, (12, 0, 12)), + ( + {"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 3, "total_tokens": 5, "completion_tokens": 2}}, + (5, 0, 5), + ), + ], +) +def test_voyage_contextual_embedding_usage_tolerates_missing_null_and_loosely_typed_counts(payload, expected): + response = _transform(payload) + + assert response.data == [EMBEDDING_ROW] + assert _token_counts(response) == expected + + +@pytest.mark.parametrize("usage", [None, [], "tokens", 3]) +def test_voyage_contextual_embedding_usage_that_is_not_an_object_is_rejected(usage): + with pytest.raises(ValidationError): + _transform({"data": [EMBEDDING_ROW], "usage": usage}) + + +@pytest.mark.parametrize("usage", [{"total_tokens": "many"}, {"total_tokens": 1.5}, {"total_tokens": [1]}]) +def test_voyage_contextual_embedding_token_count_that_is_not_an_integer_is_rejected(usage): + with pytest.raises(ValidationError): + _transform({"data": [EMBEDDING_ROW], "usage": usage}) diff --git a/tests/unit/llms/voyage/embedding/test_transformation_multimodal.py b/tests/unit/llms/voyage/embedding/test_transformation_multimodal.py new file mode 100644 index 00000000000..6be8e313a84 --- /dev/null +++ b/tests/unit/llms/voyage/embedding/test_transformation_multimodal.py @@ -0,0 +1,91 @@ +import datetime + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.voyage.embedding.transformation_multimodal import VoyageMultimodalEmbeddingConfig +from litellm.types.utils import EmbeddingResponse + +EMBEDDING_ROW = {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]} + + +def _transform(payload: object) -> EmbeddingResponse: + logging_obj = Logging( + model="voyage/voyage-multimodal-3", + messages=[], + stream=False, + call_type="embedding", + start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-id", + function_id="function-id", + ) + return VoyageMultimodalEmbeddingConfig().transform_embedding_response( + model="voyage-multimodal-3", + raw_response=httpx.Response(200, json=payload), + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + api_key="sk-test", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + +def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]: + assert response.usage is not None + return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) + + +def test_voyage_multimodal_embedding_response_carries_provider_fields_and_usage(): + response = _transform( + { + "object": "list", + "model": "voyage-multimodal-3", + "data": [EMBEDDING_ROW], + "usage": {"prompt_tokens": 8, "total_tokens": 9}, + } + ) + + assert response.model == "voyage-multimodal-3" + assert response.object == "list" + assert response.data == [EMBEDDING_ROW] + assert _token_counts(response) == (9, 0, 9) + + +@pytest.mark.parametrize( + ("payload", "expected"), + [ + ({"data": [EMBEDDING_ROW]}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 4}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 6}}, (6, 0, 6)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": None}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": ""}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": []}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": "12"}}, (12, 0, 12)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 12.0}}, (12, 0, 12)), + ( + {"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 3, "total_tokens": 5, "completion_tokens": 2}}, + (5, 0, 5), + ), + ], +) +def test_voyage_multimodal_embedding_usage_tolerates_missing_null_and_loosely_typed_counts(payload, expected): + response = _transform(payload) + + assert response.data == [EMBEDDING_ROW] + assert _token_counts(response) == expected + + +@pytest.mark.parametrize("usage", [None, [], "tokens", 3]) +def test_voyage_multimodal_embedding_usage_that_is_not_an_object_is_rejected(usage): + with pytest.raises(ValidationError): + _transform({"data": [EMBEDDING_ROW], "usage": usage}) + + +@pytest.mark.parametrize("usage", [{"total_tokens": "many"}, {"total_tokens": 1.5}, {"total_tokens": [1]}]) +def test_voyage_multimodal_embedding_token_count_that_is_not_an_integer_is_rejected(usage): + with pytest.raises(ValidationError): + _transform({"data": [EMBEDDING_ROW], "usage": usage}) diff --git a/tests/unit/llms/voyage/embedding/test_voyage_embedding_transformation.py b/tests/unit/llms/voyage/embedding/test_voyage_embedding_transformation.py new file mode 100644 index 00000000000..41bc812ec4a --- /dev/null +++ b/tests/unit/llms/voyage/embedding/test_voyage_embedding_transformation.py @@ -0,0 +1,91 @@ +import datetime + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.voyage.embedding.transformation import VoyageEmbeddingConfig +from litellm.types.utils import EmbeddingResponse + +EMBEDDING_ROW = {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]} + + +def _transform(payload: object) -> EmbeddingResponse: + logging_obj = Logging( + model="voyage/voyage-3", + messages=[], + stream=False, + call_type="embedding", + start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-id", + function_id="function-id", + ) + return VoyageEmbeddingConfig().transform_embedding_response( + model="voyage-3", + raw_response=httpx.Response(200, json=payload), + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + api_key="sk-test", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + +def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]: + assert response.usage is not None + return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) + + +def test_voyage_embedding_response_carries_provider_fields_and_usage(): + response = _transform( + { + "object": "list", + "model": "voyage-3", + "data": [EMBEDDING_ROW], + "usage": {"prompt_tokens": 8, "total_tokens": 9}, + } + ) + + assert response.model == "voyage-3" + assert response.object == "list" + assert response.data == [EMBEDDING_ROW] + assert _token_counts(response) == (9, 0, 9) + + +@pytest.mark.parametrize( + ("payload", "expected"), + [ + ({"data": [EMBEDDING_ROW]}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 4}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 6}}, (6, 0, 6)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": None}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": ""}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": []}}, (0, 0, 0)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": "12"}}, (12, 0, 12)), + ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 12.0}}, (12, 0, 12)), + ( + {"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 3, "total_tokens": 5, "completion_tokens": 2}}, + (5, 0, 5), + ), + ], +) +def test_voyage_embedding_usage_tolerates_missing_null_and_loosely_typed_counts(payload, expected): + response = _transform(payload) + + assert response.data == [EMBEDDING_ROW] + assert _token_counts(response) == expected + + +@pytest.mark.parametrize("usage", [None, [], "tokens", 3]) +def test_voyage_embedding_usage_that_is_not_an_object_is_rejected(usage): + with pytest.raises(ValidationError): + _transform({"data": [EMBEDDING_ROW], "usage": usage}) + + +@pytest.mark.parametrize("usage", [{"total_tokens": "many"}, {"total_tokens": 1.5}, {"total_tokens": [1]}]) +def test_voyage_embedding_token_count_that_is_not_an_integer_is_rejected(usage): + with pytest.raises(ValidationError): + _transform({"data": [EMBEDDING_ROW], "usage": usage}) diff --git a/tests/unit/llms/watsonx/completion/__init__.py b/tests/unit/llms/watsonx/completion/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/watsonx/completion/test_watsonx_completion_transformation.py b/tests/unit/llms/watsonx/completion/test_watsonx_completion_transformation.py new file mode 100644 index 00000000000..4ec197dcd46 --- /dev/null +++ b/tests/unit/llms/watsonx/completion/test_watsonx_completion_transformation.py @@ -0,0 +1,60 @@ +import pytest + +from litellm.llms.watsonx.completion.transformation import IBMWatsonXAIConfig + + +@pytest.mark.parametrize( + ("supplied", "stored"), + [ + pytest.param({}, {"kwargs": {}}, id="nothing-supplied"), + pytest.param( + {"decoding_method": "", "max_new_tokens": 0, "include_stop_sequences": False, "stream": False}, + { + "decoding_method": "", + "max_new_tokens": 0, + "include_stop_sequences": False, + "stream": False, + "kwargs": {}, + }, + id="falsy-values-are-kept", + ), + pytest.param( + {"top_k": 40, "top_p": None, "random_seed": None}, {"top_k": 40, "kwargs": {}}, id="none-is-skipped" + ), + pytest.param( + {"temperature": 0.2, "project_id": "proj-1", "space_id": None}, + {"temperature": 0.2, "kwargs": {"project_id": "proj-1", "space_id": None}}, + id="unknown-keywords-are-grouped-under-kwargs", + ), + ], +) +def test_init_stores_each_supplied_value_except_none_as_class_level_config( + supplied: dict[str, object], stored: dict[str, object] +) -> None: + class ScopedWatsonXConfig(IBMWatsonXAIConfig): + pass + + ScopedWatsonXConfig(**supplied) + + assert ScopedWatsonXConfig.get_config() == stored + + +def test_init_stores_the_supplied_object_itself_rather_than_a_copy() -> None: + class ScopedWatsonXConfig(IBMWatsonXAIConfig): + pass + + stop_sequences = ["}", ")"] + + ScopedWatsonXConfig(stop_sequences=stop_sequences) + + assert ScopedWatsonXConfig.get_config()["stop_sequences"] is stop_sequences + + +def test_init_called_again_with_none_keeps_the_value_stored_earlier() -> None: + class ScopedWatsonXConfig(IBMWatsonXAIConfig): + pass + + ScopedWatsonXConfig(top_k=40) + ScopedWatsonXConfig(top_k=None, top_p=0.5) + + assert ScopedWatsonXConfig.get_config() == {"top_k": 40, "top_p": 0.5, "kwargs": {}} diff --git a/tests/unit/llms/xai/test_xai_common_utils.py b/tests/unit/llms/xai/test_xai_common_utils.py new file mode 100644 index 00000000000..42a61062229 --- /dev/null +++ b/tests/unit/llms/xai/test_xai_common_utils.py @@ -0,0 +1,124 @@ +import re + +import httpx +import pytest +import respx +from pydantic import ValidationError + +import litellm +from litellm.llms.xai.common_utils import XAIModelInfo + +MODELS_URL = "https://api.x.ai/v1/models" +MANY_MODELS = [{"id": f"grok-{index}"} for index in range(500)] + + +@pytest.fixture(autouse=True) +def _default_xai_api_base(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("XAI_API_BASE", raising=False) + + +def _get_models(payload: object) -> list[str]: + with respx.mock: + respx.get(MODELS_URL).mock(return_value=httpx.Response(200, json=payload)) + return XAIModelInfo().get_models(api_key="xai-test-key") + + +def test_get_models_prefixes_every_listed_model_id_in_provider_order(): + names = _get_models( + { + "object": "list", + "data": [ + {"id": "grok-4", "object": "model", "created": 1752019200, "owned_by": "xai"}, + {"id": "grok-3-mini", "object": "model", "created": 1743724800, "owned_by": "xai"}, + {"id": ""}, + ], + } + ) + + assert names == ["xai/grok-4", "xai/grok-3-mini", "xai/"] + assert [type(name) for name in names] == [str, str, str] + + +@pytest.mark.parametrize("data", [[], "", {}]) +def test_get_models_returns_no_names_for_an_empty_data_container(data: object): + assert _get_models({"data": data}) == [] + + +def test_get_models_sends_the_api_key_as_bearer_token(): + with respx.mock: + route = respx.get(MODELS_URL).mock(return_value=httpx.Response(200, json={"data": [{"id": "grok-4"}]})) + XAIModelInfo().get_models(api_key="xai-test-key") + + assert route.calls[0].request.headers["authorization"] == "Bearer xai-test-key" + + +@pytest.mark.parametrize( + ("payload", "missing_key"), + [ + ({"object": "list"}, "data"), + ({"data": [{"object": "model"}]}, "id"), + ({"data": [*MANY_MODELS[:403], {"object": "model"}, *MANY_MODELS[404:]]}, "id"), + ], +) +def test_get_models_raises_key_error_naming_the_field_the_provider_left_out(payload: object, missing_key: str): + with pytest.raises(KeyError) as exc_info: + _get_models(payload) + + assert exc_info.value.args == (missing_key,) + + +@pytest.mark.parametrize( + ("payload", "error_type"), + [ + ([{"id": "grok-4"}], "dict_type"), + ("data", "dict_type"), + ({"data": None}, "iterable_type"), + ({"data": 5}, "iterable_type"), + ({"data": "grok-4"}, "dict_type"), + ({"data": {"grok-4": {"id": "grok-4"}}}, "dict_type"), + ({"data": ["grok-4"]}, "dict_type"), + ({"data": [None]}, "dict_type"), + ({"data": [{"id": None}]}, "string_type"), + ({"data": [{"id": 12}]}, "string_type"), + ({"data": [{"id": ["grok-4"]}]}, "string_type"), + ({"data": [*MANY_MODELS[:403], "grok-403", *MANY_MODELS[404:]]}, "dict_type"), + ({"data": [*MANY_MODELS[:429], {"id": 429}, *MANY_MODELS[430:]]}, "string_type"), + ], +) +def test_get_models_rejects_a_malformed_listing_without_echoing_position_or_input(payload: object, error_type: str): + with pytest.raises(ValidationError) as exc_info: + _get_models(payload) + + assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [(error_type, ())] + assert "403" not in str(exc_info.value) + assert "429" not in str(exc_info.value) + assert "grok" not in str(exc_info.value) + + +def test_get_models_reports_the_upstream_status_and_body_on_an_http_error(): + with respx.mock: + respx.get(MODELS_URL).mock(return_value=httpx.Response(401, text="bad key")) + with pytest.raises( + Exception, match=re.escape("Failed to fetch models from XAI. Status code: 401, Response: bad key") + ): + XAIModelInfo().get_models(api_key="xai-test-key") + + +@pytest.mark.parametrize( + ("api_key", "payload", "expected"), + [ + ("xai-listing-key", {"data": [{"id": "grok-4"}, {"id": "grok-3"}]}, ["xai/grok-4", "xai/grok-3"]), + ("xai-bad-id-key", {"data": [{"id": "grok-4"}, {"id": 7}]}, []), + ("xai-null-data-key", {"data": None}, []), + ("xai-missing-data-key", {"object": "list"}, []), + ("xai-list-body-key", [{"id": "grok-4"}], []), + ], +) +def test_get_valid_models_lists_xai_models_and_falls_back_to_empty_on_a_malformed_listing( + api_key: str, payload: object, expected: list[str] +): + with respx.mock: + respx.get(MODELS_URL).mock(return_value=httpx.Response(200, json=payload)) + models = litellm.get_valid_models(check_provider_endpoint=True, custom_llm_provider="xai", api_key=api_key) + + assert models == expected diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index a5d0d0a3ecc..d2f2fec073c 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -13,8 +13,11 @@ from dataclasses import dataclass from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx +import litellm from litellm.proxy._types import UserAPIKeyAuth from litellm.types.agents import AgentCaller @@ -2689,3 +2692,59 @@ def test_forwarding_headers_minted_bearer_replaces_a_forwarded_authorization_of_ ) assert merged == {"X-Custom": "kept", "Authorization": "Bearer minted-token"} + + +_UPSTREAM_STATUS_ONLY_ERROR: Final = { + "jsonrpc": "2.0", + "id": "req-1", + "error": {"code": -32603, "message": "Not Found"}, +} + + +@pytest.mark.parametrize( + ("upstream_body", "relayed_event"), + [ + ( + b'{"jsonrpc":"2.0","id":"up-1","error":{"code":-32001,"message":"Task not found"}}', + {"jsonrpc": "2.0", "id": "up-1", "error": {"code": -32001, "message": "Task not found"}}, + ), + (b'{"error":"quota exceeded"}', {"error": "quota exceeded", "id": "req-1"}), + (b'{"jsonrpc":"2.0","id":"up-1","result":{"ok":true}}', _UPSTREAM_STATUS_ONLY_ERROR), + (b'["error"]', _UPSTREAM_STATUS_ONLY_ERROR), + (b'"error"', _UPSTREAM_STATUS_ONLY_ERROR), + (b"null", _UPSTREAM_STATUS_ONLY_ERROR), + (b"7", _UPSTREAM_STATUS_ONLY_ERROR), + (b"\xff\xfe", _UPSTREAM_STATUS_ONLY_ERROR), + (b"", _UPSTREAM_STATUS_ONLY_ERROR), + (b'{"error":null}', {"error": None, "id": "req-1"}), + ], +) +@pytest.mark.asyncio +async def test_failed_upstream_stream_relays_its_body_only_when_that_is_a_json_object_with_an_error( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + upstream_body: bytes, + relayed_event: Mapping[str, object], +): + from litellm.proxy.agent_endpoints.a2a_endpoints import _a2a_sse_event_source + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream: Final = respx_mock.post("http://backend-agent:10001").mock( + return_value=httpx.Response(404, content=upstream_body) + ) + + events: Final = [ + event + async for event in _a2a_sse_event_source( + "http://backend-agent:10001", + {"jsonrpc": "2.0", "id": "req-1", "method": "tasks/resubscribe"}, + request_id="req-1", + ) + ] + + assert json.dumps(events, sort_keys=True) == json.dumps([relayed_event], sort_keys=True), upstream_body + assert json.loads(upstream.calls.last.request.content) == { + "jsonrpc": "2.0", + "id": "req-1", + "method": "tasks/resubscribe", + } diff --git a/tests/unit/proxy/client/cli/test_auth_commands.py b/tests/unit/proxy/client/cli/test_auth_commands.py index b4a6aaf706e..27edf9f3eb3 100644 --- a/tests/unit/proxy/client/cli/test_auth_commands.py +++ b/tests/unit/proxy/client/cli/test_auth_commands.py @@ -2075,3 +2075,94 @@ def test_login_shows_the_error_detail_only_when_the_proxy_answers_with_a_json_ob "Authentication failed: Starting CLI login failed: HTTP 429 from https://test.example.com/sso/cli/start" f"{expected_detail}\n" ) + + +@pytest.mark.parametrize( + ("obj", "expected_exit_code", "expected_stdout", "expected_stderr"), + [ + ({}, 0, "sk-stored\n", ""), + ({"base_url": "https://test.example.com///", "base_url_explicit": True}, 0, "sk-stored\n", ""), + ( + {"base_url": "https://other.example.com", "base_url_explicit": True}, + 1, + "", + "Not authenticated for this server. Run 'lite login'.\n", + ), + ({"base_url": "https://other.example.com", "base_url_explicit": False}, 0, "sk-stored\n", ""), + ( + {"base_url": "https://test.example.com", "api_key": "sk-context", "api_key_from_token_file": True}, + 0, + "sk-context\n", + "", + ), + ( + {"base_url": "https://test.example.com", "api_key": "sk-context", "api_key_from_token_file": False}, + 0, + "sk-stored\n", + "", + ), + ( + {"base_url": "https://other.example.com", "api_key": "sk-context", "api_key_from_token_file": True}, + 0, + "sk-stored\n", + "", + ), + ], +) +def test_print_token_reads_the_connection_from_the_context_object( + isolated_home, obj, expected_exit_code, expected_stdout, expected_stderr +): + _write_token_file(isolated_home, key="sk-stored") + + result = CliRunner().invoke(print_token, obj=obj) + + assert result.exit_code == expected_exit_code + assert result.stdout == expected_stdout + assert result.stderr == expected_stderr + + +def test_print_token_without_a_context_object_reports_no_login_when_nothing_is_stored(isolated_home): + result = CliRunner().invoke(print_token) + + assert result.exit_code == 1 + assert result.stdout == "" + assert result.stderr == "Not authenticated. Run 'lite login'.\n" + + +def test_print_token_without_a_context_object_fails_reading_it_once_a_login_is_stored(isolated_home): + _write_token_file(isolated_home, key="sk-stored") + + result = CliRunner().invoke(print_token) + + assert result.exit_code == 1 + assert type(result.exception) is AttributeError + assert result.stdout == "" + assert result.stderr == "" + + +@pytest.mark.parametrize( + "obj", + [ + {}, + {"api_key": "sk-context", "base_url_explicit": True}, + ], +) +@responses.activate +def test_login_without_a_base_url_in_the_context_fails_before_any_request(obj): + result = CliRunner().invoke(login, obj=obj) + + assert result.exit_code == 1 + assert type(result.exception) is KeyError + assert result.exception.args == ("base_url",) + assert result.output == "" + assert len(responses.calls) == 0 + + +@responses.activate +def test_login_without_a_context_object_fails_before_any_request(): + result = CliRunner().invoke(login) + + assert result.exit_code == 1 + assert type(result.exception) is TypeError + assert result.output == "" + assert len(responses.calls) == 0 diff --git a/tests/unit/proxy/client/cli/test_global_options.py b/tests/unit/proxy/client/cli/test_global_options.py index d46cc2ad120..96b200b3983 100644 --- a/tests/unit/proxy/client/cli/test_global_options.py +++ b/tests/unit/proxy/client/cli/test_global_options.py @@ -266,3 +266,106 @@ def test_lite_entrypoint_prints_nothing_on_stderr(monkeypatch, capsys, requests_ assert "LiteLLM Proxy Server Version: 1.2.3" in captured.out assert f"LiteLLM Proxy CLI Version: {litellm_version}" in captured.out assert captured.err == "" + + +@pytest.mark.parametrize( + ("args", "config", "expected_context", "expected_authorization"), + [ + ( + ["--base-url", "https://flag-proxy.example.com/", "--api-key", "sk-flag"], + {}, + { + "base_url": "https://flag-proxy.example.com", + "api_key": "sk-flag", + "api_key_from_token_file": False, + "base_url_explicit": True, + }, + "Bearer sk-flag", + ), + ( + [], + {}, + { + "base_url": "http://localhost:4000", + "api_key": None, + "api_key_from_token_file": True, + "base_url_explicit": False, + }, + None, + ), + ( + ["--api-key", ""], + {"base_url": "https://config-proxy.example.com"}, + { + "base_url": "https://config-proxy.example.com", + "api_key": "", + "api_key_from_token_file": False, + "base_url_explicit": True, + }, + None, + ), + ], +) +def test_cli_stores_the_resolved_connection_on_the_supplied_context_object( + cli_runner, isolated_home, requests_mock, args, config, expected_context, expected_authorization +): + _write_config_file(isolated_home, config) + requests_mock.get(f"{expected_context['base_url']}/health/readiness", json={"litellm_version": "1.2.3"}) + context: Final[dict[str, object]] = {} + + result = cli_runner.invoke(cli, [*args, "version"], obj=context) + + assert result.exit_code == 0 + assert json.dumps(context) == json.dumps(expected_context) + assert "LiteLLM Proxy Server Version: 1.2.3" in result.output + assert requests_mock.last_request.headers.get("Authorization") == expected_authorization + + +@pytest.mark.parametrize( + ("context", "expected_lines", "expected_authorization"), + [ + ( + {"base_url": "http://localhost:4000", "api_key": "sk-context"}, + ["LiteLLM Proxy Server URL: http://localhost:4000", "LiteLLM Proxy Server Version: 1.2.3"], + "Bearer sk-context", + ), + ( + {"base_url": "http://localhost:4000"}, + ["LiteLLM Proxy Server URL: http://localhost:4000", "LiteLLM Proxy Server Version: 1.2.3"], + None, + ), + ( + {"base_url": "http://localhost:4000", "api_key": None, "api_key_from_token_file": True}, + ["LiteLLM Proxy Server URL: http://localhost:4000", "LiteLLM Proxy Server Version: 1.2.3"], + None, + ), + ], +) +def test_version_command_reads_the_connection_from_the_context_object( + cli_runner, requests_mock, context, expected_lines, expected_authorization +): + requests_mock.get("http://localhost:4000/health/readiness", json={"litellm_version": "1.2.3"}) + + result = cli_runner.invoke(cli.commands["version"], [], obj=context) + + assert result.exit_code == 0 + assert result.output.splitlines() == [f"LiteLLM Proxy CLI Version: {litellm_version}", *expected_lines] + assert requests_mock.last_request.headers.get("Authorization") == expected_authorization + + +def test_version_command_without_a_base_url_reports_the_failed_lookup(cli_runner): + result = cli_runner.invoke(cli.commands["version"], [], obj={"api_key": "sk-context"}) + + lines: Final = result.output.splitlines() + assert result.exit_code == 0 + assert len(lines) == 2 + assert lines[0] == f"LiteLLM Proxy CLI Version: {litellm_version}" + assert lines[1].startswith("Could not retrieve server version: ") + + +def test_version_command_without_a_context_object_fails_reading_it(cli_runner): + result = cli_runner.invoke(cli.commands["version"], []) + + assert result.exit_code == 1 + assert type(result.exception) is AttributeError + assert result.output == "" diff --git a/tests/unit/proxy/client/cli/test_keys_commands.py b/tests/unit/proxy/client/cli/test_keys_commands.py index 5cc0fb70881..7053e7d36fb 100644 --- a/tests/unit/proxy/client/cli/test_keys_commands.py +++ b/tests/unit/proxy/client/cli/test_keys_commands.py @@ -2,14 +2,16 @@ import json import os from unittest.mock import patch -import requests - - - import pytest +import requests +import responses from click.testing import CliRunner from litellm.proxy.client.cli import cli +from litellm.proxy.client.cli.commands.keys import keys + +DESTINATION = "https://destination.example.com" +SOURCE = "https://source.example.com" @pytest.fixture @@ -538,3 +540,84 @@ def test_keys_import_with_all_key_properties(mock_keys_client, cli_runner): budget_id="budget-456", config={"max_tokens": 1000}, ) + + +def _sent_requests(): + return [ + (call.request.method, call.request.url, call.request.headers.get("Authorization")) for call in responses.calls + ] + + +@pytest.mark.parametrize(("api_key", "authorization"), [("sk-destination", "Bearer sk-destination"), (None, None)]) +@pytest.mark.parametrize( + ("args", "method", "path", "sent_body", "payload"), + [ + ( + ["list", "--format", "json"], + "GET", + "/key/list?return_full_object=true&include_team_keys=false", + None, + {"keys": [{"token": "abc123", "key_alias": "alias1"}], "total_count": 1}, + ), + (["generate", "--key-alias", "ci"], "POST", "/key/generate", {"key_alias": "ci"}, {"key": "sk-new"}), + ( + ["delete", "--keys", "sk-old"], + "POST", + "/key/delete", + {"keys": ["sk-old"], "key_aliases": None}, + {"deleted_keys": ["sk-old"]}, + ), + ], +) +@responses.activate +def test_keys_command_calls_the_server_stored_on_the_cli_context( + cli_runner, args, method, path, sent_body, payload, api_key, authorization +): + responses.add(method, f"{DESTINATION}{path.partition('?')[0]}", json=payload) + + result = cli_runner.invoke(keys, args, obj={"base_url": DESTINATION, "api_key": api_key}) + + assert result.exit_code == 0, result.output + assert json.loads(result.output) == payload + assert _sent_requests() == [(method, f"{DESTINATION}{path}", authorization)] + assert (responses.calls[0].request.body and json.loads(responses.calls[0].request.body)) == sent_body + + +@pytest.mark.parametrize(("api_key", "authorization"), [("sk-destination", "Bearer sk-destination"), (None, None)]) +@responses.activate +def test_keys_import_reads_the_source_and_writes_to_the_server_stored_on_the_cli_context( + cli_runner, api_key, authorization +): + responses.get(f"{SOURCE}/key/list", json={"keys": [{"key_alias": "imported", "user_id": "u1"}]}) + responses.post(f"{DESTINATION}/key/generate", json={"key": "sk-imported"}) + + result = cli_runner.invoke( + keys, + ["import", "--source-base-url", SOURCE, "--source-api-key", "sk-source"], + obj={"base_url": DESTINATION, "api_key": api_key}, + ) + + assert result.exit_code == 0, result.output + assert "Successfully imported: 1" in result.output + assert _sent_requests() == [ + ("GET", f"{SOURCE}/key/list?page=1&size=100&return_full_object=true", "Bearer sk-source"), + ("POST", f"{DESTINATION}/key/generate", authorization), + ] + assert json.loads(responses.calls[1].request.body) == {"key_alias": "imported", "user_id": "u1"} + + +@pytest.mark.parametrize( + "args", + [ + ["list"], + ["generate", "--key-alias", "ci"], + ["delete", "--keys", "sk-old"], + ["import", "--source-base-url", SOURCE, "--source-api-key", "sk-source"], + ], +) +@responses.activate +def test_keys_command_without_cli_context_values_fails_before_any_request(cli_runner, args): + result = cli_runner.invoke(keys, args) + + assert result.exit_code == 1 + assert _sent_requests() == [] diff --git a/tests/unit/proxy/client/test_http_commands.py b/tests/unit/proxy/client/test_http_commands.py index 04894248ff2..b4dd8070134 100644 --- a/tests/unit/proxy/client/test_http_commands.py +++ b/tests/unit/proxy/client/test_http_commands.py @@ -118,3 +118,122 @@ def test_request_invalid_header(runner): ) assert result.exit_code == 2 # Click error code for invalid parameter assert "Invalid header format" in result.output + + +@pytest.mark.parametrize( + "raw_body", + [ + '{"count": "12", "ratio": 1, "whole": 1.0, "enabled": "true", "flag": true, "missing": null}', + '{"data": [{"id": 1, "tags": ["a", 2, null, false]}, "two", null], "nested": {"deep": {"n": "007"}}}', + '[{"id": "1"}, 2, "3", null]', + '"plain string"', + "12", + "null", + "{}", + ], +) +@responses.activate +def test_request_prints_the_server_json_unchanged(runner, raw_body): + responses.add( + responses.GET, + "http://localhost:4000/models", + body=raw_body, + content_type="application/json", + status=200, + ) + result = runner.invoke( + http, + ["request", "GET", "/models"], + obj={"base_url": "http://localhost:4000", "api_key": "sk-test-key"}, + ) + assert result.exit_code == 0 + assert json.dumps(json.loads(result.stdout)) == json.dumps(json.loads(raw_body)) + + +@pytest.mark.parametrize( + "raw_body", + [ + '{"error": {"message": "model not found", "code": "404", "status": 404, "retry": false}}', + '["first", {"detail": null}]', + '"denied"', + ], +) +@responses.activate +def test_request_http_error_prints_the_json_error_body_unchanged(runner, raw_body): + responses.add( + responses.GET, + "http://localhost:4000/models", + body=raw_body, + content_type="application/json", + status=404, + ) + result = runner.invoke( + http, + ["request", "GET", "/models"], + obj={"base_url": "http://localhost:4000", "api_key": "sk-test-key"}, + ) + assert result.exit_code == 1 + assert result.stderr.splitlines()[:1] == ["Error: HTTP 404"] + assert json.dumps(json.loads(result.stdout)) == json.dumps(json.loads(raw_body)) + + +@responses.activate +def test_request_http_error_with_a_non_json_body_echoes_the_text(runner): + responses.add( + responses.GET, + "http://localhost:4000/models", + body="upstream exploded", + status=502, + ) + result = runner.invoke( + http, + ["request", "GET", "/models"], + obj={"base_url": "http://localhost:4000", "api_key": "sk-test-key"}, + ) + assert result.exit_code == 1 + assert result.stdout == "" + assert result.stderr.splitlines()[:2] == ["Error: HTTP 502", "upstream exploded"] + + +@pytest.mark.parametrize( + ("api_key", "expected_authorization"), + [ + ("sk-test-key", "Bearer sk-test-key"), + (None, None), + ("", None), + ], +) +@responses.activate +def test_request_sends_the_context_api_key_to_the_context_base_url(runner, api_key, expected_authorization): + responses.add(responses.GET, "http://proxy.internal:4100/models", json={"models": []}, status=200) + result = runner.invoke( + http, + ["request", "GET", "/models"], + obj={"base_url": "http://proxy.internal:4100/", "api_key": api_key}, + ) + assert result.exit_code == 0 + assert len(responses.calls) == 1 + assert responses.calls[0].request.url == "http://proxy.internal:4100/models" + assert responses.calls[0].request.headers.get("Authorization") == expected_authorization + + +@pytest.mark.parametrize( + ("obj", "missing_key"), + [ + ({"api_key": "sk-test-key"}, "base_url"), + ({"base_url": "http://localhost:4000"}, "api_key"), + ({}, "base_url"), + ], +) +def test_request_with_an_incomplete_context_fails_on_the_first_missing_value(runner, obj, missing_key): + result = runner.invoke(http, ["request", "GET", "/models"], obj=obj) + assert result.exit_code == 1 + assert type(result.exception) is KeyError + assert result.exception.args == (missing_key,) + + +def test_request_without_a_context_object_fails_instead_of_assuming_a_default_server(runner): + result = runner.invoke(http, ["request", "GET", "/models"], obj=None) + assert result.exit_code == 1 + assert type(result.exception) is TypeError + assert result.output == "" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py b/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py index 67ce4046ce0..80bde77afdc 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py @@ -3,10 +3,12 @@ Test DynamoAI Guardrails integration """ import importlib +import json import os from unittest.mock import AsyncMock, MagicMock, patch import pytest +import respx import litellm from litellm.caching.caching import DualCache @@ -118,6 +120,89 @@ async def test_dynamoai_allows_content_with_none_action(): assert result == request_data +_ANALYZE_URL = "https://api.dynamo.ai/v1/moderation/analyze/" +_BLOCK_VERDICT = { + "finalAction": "BLOCK", + "appliedPolicies": [ + { + "policy": {"id": "policy-123", "name": "Toxicity Policy", "method": "TOXICITY"}, + "outputs": {"action": "BLOCK", "message": "Content contains toxic language"}, + } + ], +} + + +def _post_call_guardrail() -> DynamoAIGuardrails: + return DynamoAIGuardrails( + guardrail_name="test-dynamoai", + api_key="test-api-key", + api_base="https://api.dynamo.ai", + event_hook="post_call", + default_on=True, + ) + + +def _assistant_reply(content: str | None) -> litellm.ModelResponse: + return litellm.ModelResponse( + choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": content}}] + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +async def test_post_call_hook_sends_the_assistant_text_to_dynamoai_and_allows_a_clean_reply( + respx_mock: respx.MockRouter, +) -> None: + route = respx_mock.post(_ANALYZE_URL).respond(json={"finalAction": "NONE", "appliedPolicies": []}) + request_data = {"model": "gpt-5.5", "messages": [{"role": "user", "content": "hi"}]} + + result = await _post_call_guardrail().async_post_call_success_hook( + data=request_data, user_api_key_dict=UserAPIKeyAuth(), response=_assistant_reply("Hello, how are you?") + ) + + assert result is None + sent = route.calls.last.request + assert json.loads(sent.content) == {"messages": [{"role": "assistant", "content": "Hello, how are you?"}]} + assert sent.headers["Authorization"] == "Bearer test-api-key" + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +async def test_post_call_hook_raises_when_dynamoai_blocks_the_assistant_text(respx_mock: respx.MockRouter) -> None: + respx_mock.post(_ANALYZE_URL).respond(json=_BLOCK_VERDICT) + + with pytest.raises(ValueError, match="Guardrail failed") as blocked: + await _post_call_guardrail().async_post_call_success_hook( + data={"model": "gpt-5.5", "messages": [{"role": "user", "content": "hi"}]}, + user_api_key_dict=UserAPIKeyAuth(), + response=_assistant_reply("This is harmful content"), + ) + + assert str(blocked.value) == ( + "Guardrail failed: 1 violation(s) detected\n\n" + "- TOXICITY POLICY:\n" + " Action: BLOCK\n" + " Method: TOXICITY\n" + " Message: Content contains toxic language\n" + " Policy ID: policy-123" + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +async def test_post_call_hook_skips_dynamoai_when_the_reply_has_no_text(respx_mock: respx.MockRouter) -> None: + route = respx_mock.post(_ANALYZE_URL).respond(json=_BLOCK_VERDICT) + + result = await _post_call_guardrail().async_post_call_success_hook( + data={"model": "gpt-5.5", "messages": [{"role": "user", "content": "hi"}]}, + user_api_key_dict=UserAPIKeyAuth(), + response=_assistant_reply(None), + ) + + assert result is None + assert route.called is False + + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py b/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py b/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py new file mode 100644 index 00000000000..4a591c08bf3 --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py @@ -0,0 +1,124 @@ +import json + +import pytest +import respx + +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.ibm_guardrails.ibm_detector import IBMGuardrailDetector + +pytestmark = pytest.mark.usefixtures("httpx_transport") + +_BASE_URL = "https://ibm-detector.example.test" +_DETECTOR_SERVER_URL = f"{_BASE_URL}/api/v1/text/contents" +_ORCHESTRATOR_URL = f"{_BASE_URL}/api/v2/text/detection/content" +_NESTED_PARAMS: dict[str, object] = { + "threshold": 0.5, + "max_hits": 3, + "strict": True, + "labels": ["hap", None], + "regex": {"email": ".+@.+"}, +} +_DETECTOR_PARAMS_CASES = [ + pytest.param(None, {}, id="omitted"), + pytest.param({}, {}, id="empty"), + pytest.param(_NESTED_PARAMS, _NESTED_PARAMS, id="nested-mixed-types"), +] +_HAP_DETECTION = { + "start": 0, + "end": 5, + "text": "hello", + "detection": "HAP", + "detection_type": "hap", + "score": 0.97, +} + + +def _guardrail(*, is_detector_server: bool, detector_params: dict[str, object] | None) -> IBMGuardrailDetector: + return IBMGuardrailDetector( + guardrail_name="ibm-guard", + auth_token="ibm-token", + base_url=_BASE_URL, + detector_id="hap-detector", + is_detector_server=is_detector_server, + detector_params=detector_params, + event_hook="pre_call", + default_on=True, + ) + + +def _request_data() -> dict[str, object]: + return {"messages": [{"role": "user", "content": "hello there"}]} + + +def test_constructor_keeps_the_configured_detector_params_object() -> None: + guardrail = _guardrail(is_detector_server=True, detector_params=_NESTED_PARAMS) + + assert guardrail.detector_params is _NESTED_PARAMS + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("detector_params", "sent_params"), _DETECTOR_PARAMS_CASES) +async def test_detector_server_request_carries_configured_detector_params( + respx_mock: respx.MockRouter, detector_params: dict[str, object] | None, sent_params: dict[str, object] +) -> None: + route = respx_mock.post(_DETECTOR_SERVER_URL).respond(json=[[]]) + data = _request_data() + + returned = await _guardrail(is_detector_server=True, detector_params=detector_params).async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type="completion" + ) + + sent = route.calls.last.request + assert returned is data + assert repr(json.loads(sent.content)) == repr({"contents": ["hello there"], "detector_params": sent_params}) + assert sent.headers["detector-id"] == "hap-detector" + assert sent.headers["Authorization"] == "Bearer ibm-token" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("detector_params", "sent_params"), _DETECTOR_PARAMS_CASES) +async def test_orchestrator_request_keys_detector_params_by_detector_id( + respx_mock: respx.MockRouter, detector_params: dict[str, object] | None, sent_params: dict[str, object] +) -> None: + route = respx_mock.post(_ORCHESTRATOR_URL).respond(json={"detections": []}) + data = _request_data() + + returned = await _guardrail(is_detector_server=False, detector_params=detector_params).async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type="completion" + ) + + sent = route.calls.last.request + assert returned is data + assert repr(json.loads(sent.content)) == repr( + {"content": "hello there", "detectors": {"hap-detector": sent_params}} + ) + assert "detector-id" not in sent.headers + + +@pytest.mark.asyncio +async def test_detector_server_detection_blocks_the_request_with_a_violation_summary( + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post(_DETECTOR_SERVER_URL).respond(json=[[_HAP_DETECTION]]) + + with pytest.raises(ValueError, match=r"IBM Guardrail Detector failed: 1 violation\(s\) detected") as blocked: + await _guardrail(is_detector_server=True, detector_params=_NESTED_PARAMS).async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=_request_data(), call_type="completion" + ) + + assert " - HAP (score: 0.970)\n Text: 'hello'" in str(blocked.value) + + +@pytest.mark.asyncio +async def test_orchestrator_detection_blocks_the_request_and_names_the_detector( + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post(_ORCHESTRATOR_URL).respond(json={"detections": [_HAP_DETECTION]}) + + with pytest.raises(ValueError, match=r"IBM Guardrail Detector failed: 1 violation\(s\) detected") as blocked: + await _guardrail(is_detector_server=False, detector_params=_NESTED_PARAMS).async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=_request_data(), call_type="completion" + ) + + assert "- HAP (detector: hap-detector, score: 0.970)\n Text: 'hello'" in str(blocked.value) diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 24fb60e5389..973d6bc9f7b 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import litellm import pytest +from pydantic import ValidationError from litellm.integrations.custom_logger import CustomLogger @@ -2864,6 +2865,51 @@ def test_build_complete_streaming_response(all_chunks): assert result.usage.total_tokens == 266 +_USABLE_MESSAGE_START: Final = 'data: {"type": "message_start", "message": {"id": "msg_real", "model": "claude-real"}}' + + +@pytest.mark.parametrize( + "unusable_line", + [ + "data: [1, 2]", + 'data: "message_start"', + "data: 7", + "data: null", + "data: not json", + 'data: {"type": "message_start"}', + 'data: {"type": "message_start", "message": null}', + 'data: {"type": "message_start", "message": {}}', + 'data: {"type": "message_start", "message": ""}', + 'data: {"type": "message_start", "message": []}', + 'data: {"type": "message_start", "message": {"id": 7, "model": ["claude"]}}', + 'data: {"type": "message_start", "message": {"id": "", "model": null}}', + 'data: {"type": "ping", "message": "not an object"}', + ], +) +def test_model_and_response_id_come_from_the_first_message_start_that_carries_them(unusable_line: str): + all_chunks: Final = [f"event: message_start\n{unusable_line}\n\n".encode(), _USABLE_MESSAGE_START] + + extracted: Final = ( + AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks(all_chunks), + AnthropicPassthroughLoggingHandler._extract_response_id_from_anthropic_chunks(all_chunks), + ) + + assert extracted == ("claude-real", "msg_real"), unusable_line + + +@pytest.mark.parametrize("message", ['"upstream text"', "7", "[1, 2]", "true"]) +def test_message_start_whose_message_is_not_an_object_fails_without_echoing_the_upstream_value(message: str): + all_chunks: Final = [f'data: {{"type": "message_start", "message": {message}}}', _USABLE_MESSAGE_START] + + with pytest.raises(ValidationError) as exc_info: + AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks(all_chunks) + + assert str(exc_info.value).splitlines()[:2] == [ + "1 validation error for dict[str,any]", + " Input should be a valid dictionary [type=dict_type]", + ], message + + @pytest.mark.parametrize( "end_user_id", [{"litellm_metadata": {"user": "test"}}, {"metadata": {"user_id": "test"}}], diff --git a/tests/unit/proxy/policy_engine/test_policy_registry.py b/tests/unit/proxy/policy_engine/test_policy_registry.py index db7775b6238..1d50230a7a9 100644 --- a/tests/unit/proxy/policy_engine/test_policy_registry.py +++ b/tests/unit/proxy/policy_engine/test_policy_registry.py @@ -152,3 +152,134 @@ async def test_update_policy_in_db_rejects_a_malformed_pipeline_before_writing( f"Error updating policy in DB: 1 validation error for GuardrailPipeline\n{rejected_field}\n" ) assert table.writes == [] + + +_VERSION_CONTENT = { + "inherit": "base-policy", + "description": "blocks pii", + "guardrails_add": ["pii-guard"], + "guardrails_remove": ["legacy-guard"], + "condition": {"model": "gpt-4"}, + "pipeline": _STORED_PIPELINE, +} + +_CHANGED_VERSION_CONTENT = { + "inherit": None, + "description": "", + "guardrails_add": ["pii-guard", "toxicity-guard"], + "guardrails_remove": [], + "condition": {"model": "gpt-4o"}, + "pipeline": None, +} + + +class _VersionsTable: + def __init__(self, rows: Mapping[str, SimpleNamespace]) -> None: + self.rows = rows + self.lookups: list[Mapping[str, object]] = [] + + async def find_unique(self, where: Mapping[str, object]) -> SimpleNamespace | None: + self.lookups.append(where) + return self.rows.get(str(where["policy_id"])) + + +def _version_row(policy_id: str, content: Mapping[str, object]) -> SimpleNamespace: + return SimpleNamespace( + policy_id=policy_id, + policy_name="pii-policy", + version_number=1, + version_status="draft", + parent_version_id=None, + is_latest=True, + published_at=None, + production_at=None, + created_at=_NOW, + updated_at=_NOW, + created_by=None, + updated_by=None, + **{**_VERSION_CONTENT, **content}, + ) + + +def _versions_client(table: _VersionsTable) -> SimpleNamespace: + return SimpleNamespace(db=SimpleNamespace(litellm_policytable=table)) + + +@pytest.mark.parametrize("field", list(_VERSION_CONTENT)) +@pytest.mark.asyncio +async def test_compare_versions_reports_exactly_the_one_field_that_differs(field: str): + value_a = _VERSION_CONTENT[field] + value_b = _CHANGED_VERSION_CONTENT[field] + table = _VersionsTable( + { + "version-a": _version_row("version-a", {}), + "version-b": _version_row("version-b", {field: value_b}), + } + ) + + result = await PolicyRegistry().compare_versions("version-a", "version-b", _versions_client(table)) + + assert result.field_diffs == {field: {"version_a": value_a, "version_b": value_b}} + assert [type(value) for value in result.field_diffs[field].values()] == [type(value_a), type(value_b)] + assert (result.version_a.policy_id, result.version_b.policy_id) == ("version-a", "version-b") + assert table.lookups == [{"policy_id": "version-a"}, {"policy_id": "version-b"}] + + +@pytest.mark.asyncio +async def test_compare_versions_lists_every_differing_field_in_content_order_with_each_side(): + table = _VersionsTable( + { + "version-a": _version_row("version-a", _CHANGED_VERSION_CONTENT), + "version-b": _version_row("version-b", {}), + } + ) + + result = await PolicyRegistry().compare_versions("version-a", "version-b", _versions_client(table)) + + assert list(result.field_diffs) == list(_VERSION_CONTENT) + assert result.field_diffs == { + field: {"version_a": _CHANGED_VERSION_CONTENT[field], "version_b": _VERSION_CONTENT[field]} + for field in _VERSION_CONTENT + } + assert result.field_diffs["inherit"]["version_a"] is None + assert result.field_diffs["pipeline"]["version_a"] is None + + +@pytest.mark.asyncio +async def test_compare_versions_ignores_version_metadata_and_treats_unset_guardrail_lists_as_empty(): + row_a = _version_row("version-a", {"guardrails_add": None, "guardrails_remove": None}) + row_b = _version_row("version-b", {"guardrails_add": [], "guardrails_remove": []}) + row_b.policy_name = "renamed-policy" + row_b.version_number = 7 + row_b.version_status = "production" + row_b.parent_version_id = "version-a" + row_b.is_latest = False + row_b.created_by = "admin" + table = _VersionsTable({"version-a": row_a, "version-b": row_b}) + + result = await PolicyRegistry().compare_versions("version-a", "version-b", _versions_client(table)) + + assert result.field_diffs == {} + assert (result.version_a.version_number, result.version_b.version_number) == (1, 7) + assert (result.version_a.guardrails_add, result.version_b.guardrails_add) == ([], []) + + +@pytest.mark.parametrize( + ("stored_versions", "missing_version"), + [ + pytest.param(("version-b",), "version-a", id="first-missing"), + pytest.param(("version-a",), "version-b", id="second-missing"), + pytest.param((), "version-a", id="both-missing-names-the-first"), + ], +) +@pytest.mark.asyncio +async def test_compare_versions_names_the_missing_version_after_looking_both_up( + stored_versions: tuple[str, ...], missing_version: str +): + table = _VersionsTable({policy_id: _version_row(policy_id, {}) for policy_id in stored_versions}) + + with pytest.raises(Exception, match=f"^Error comparing versions: Policy {missing_version} not found$") as raised: + await PolicyRegistry().compare_versions("version-a", "version-b", _versions_client(table)) + + assert type(raised.value) is Exception + assert table.lookups == [{"policy_id": "version-a"}, {"policy_id": "version-b"}] diff --git a/tests/unit/proxy/response_polling/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py index 81ca0114a8d..979ab499f3b 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -22,6 +22,7 @@ from unittest.mock import AsyncMock, Mock, patch import pytest from fastapi import Request +from pydantic import ValidationError from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler @@ -1842,3 +1843,47 @@ class TestEdgeCases: stored = json.loads(call_args.kwargs["value"]) assert stored["output"] == [] + + +@pytest.mark.parametrize( + ("request_data", "expected_metadata"), + [ + ( + {"model": "gpt-4o", "input": "Hello", "metadata": {"trace": "t-1", "tags": ["a", None], "attempt": 2}}, + {"trace": "t-1", "tags": ["a", None], "attempt": 2}, + ), + ({"model": "gpt-4o", "metadata": {}}, {}), + ({"model": "gpt-4o", "metadata": None}, None), + ({"model": "gpt-4o"}, {}), + ({"id": "resp_from_client", "status": "completed", "output": ["x"], "usage": {"total_tokens": 9}}, {}), + ], +) +async def test_create_initial_state_carries_only_the_request_metadata_into_the_queued_response( + request_data, expected_metadata +): + handler = ResponsePollingHandler(redis_cache=None) + + response = await handler.create_initial_state(polling_id="litellm_poll_abc", request_data=request_data) + + assert repr(response.metadata) == repr(expected_metadata) + assert type(response.created_at) is int + assert response.model_dump(exclude={"created_at", "metadata"}, exclude_none=True) == { + "id": "litellm_poll_abc", + "object": "response", + "status": "queued", + "output": [], + } + assert response.model_fields_set == {"id", "object", "status", "created_at", "output", "metadata", "usage"} + + +@pytest.mark.parametrize("metadata", ["trace-1", ["trace", "t-1"], 7, True, b"trace"]) +async def test_create_initial_state_rejects_request_metadata_that_is_not_an_object(metadata): + handler = ResponsePollingHandler(redis_cache=None) + + with pytest.raises(ValidationError) as rejected: + await handler.create_initial_state(polling_id="litellm_poll_abc", request_data={"metadata": metadata}) + + assert rejected.value.title == "ResponsesAPIResponse" + assert [(error["loc"], error["type"], repr(error["input"])) for error in rejected.value.errors()] == [ + (("metadata",), "dict_type", repr(metadata)) + ] diff --git a/tests/unit/proxy/test_caching_routes.py b/tests/unit/proxy/test_caching_routes.py index 6f80c281565..8a59793495a 100644 --- a/tests/unit/proxy/test_caching_routes.py +++ b/tests/unit/proxy/test_caching_routes.py @@ -2,11 +2,14 @@ import json import pytest from fastapi.testclient import TestClient +from pydantic import ValidationError import litellm from litellm.caching import RedisCache +from litellm.caching.caching import Cache +from litellm.proxy.caching_routes import _extract_cache_params from litellm.proxy.proxy_server import app client = TestClient(app) @@ -334,3 +337,63 @@ def test_cache_redis_info_no_cache(): # Restore original cache litellm.cache = original_cache + + +def _use_redis_backed_cache(monkeypatch, backend_attributes): + cache = Cache(type="redis", host="localhost", port=6379, password="hello") + vars(cache.cache).update(backend_attributes) + monkeypatch.setattr(litellm, "cache", cache) + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("host", "redis.internal"), + ("port", 6379), + ("port", "6379"), + ("namespace", "team-a"), + ("redis_version", 7.2), + ("redis_version", 7), + ("redis_version", "7.2.4"), + ], +) +def test_extract_cache_params_keeps_each_backend_value_and_its_type(monkeypatch, field, value): + _use_redis_backed_cache(monkeypatch, {field: value}) + + params = _extract_cache_params() + + assert sorted(params) == ["host", "namespace", "port", "redis_kwargs", "redis_version"] + assert repr(params[field]) == repr(value) + + +def test_extract_cache_params_reads_connection_settings_and_masks_the_password(monkeypatch): + _use_redis_backed_cache(monkeypatch, {}) + + redis_kwargs = _extract_cache_params()["redis_kwargs"] + + assert redis_kwargs["host"] == "localhost" + assert repr(redis_kwargs["port"]) == "6379" + assert "hello" not in json.dumps(redis_kwargs) + + +def test_extract_cache_params_reports_every_malformed_backend_field(monkeypatch): + _use_redis_backed_cache(monkeypatch, {"host": 7, "namespace": 5, "redis_version": ["7"]}) + + with pytest.raises(ValidationError) as exc_info: + _extract_cache_params() + + assert exc_info.value.title == "HealthCheckCacheParams" + assert sorted({error["loc"][0] for error in exc_info.value.errors()}) == ["host", "namespace", "redis_version"] + + +def test_extract_cache_params_is_empty_when_no_cache_is_configured(monkeypatch): + monkeypatch.setattr(litellm, "cache", None) + + assert _extract_cache_params() == {} + + +@pytest.mark.parametrize("attribute_name", [1, None, (1, 2), b"host"]) +def test_extract_cache_params_is_empty_when_a_backend_attribute_name_is_not_a_string(monkeypatch, attribute_name): + _use_redis_backed_cache(monkeypatch, {attribute_name: "redis.internal"}) + + assert _extract_cache_params() == {}, attribute_name diff --git a/tests/unit/proxy/test_plugin_routes.py b/tests/unit/proxy/test_plugin_routes.py index c8d2939385d..7d3fbda14fd 100644 --- a/tests/unit/proxy/test_plugin_routes.py +++ b/tests/unit/proxy/test_plugin_routes.py @@ -15,7 +15,10 @@ import asyncio from unittest.mock import MagicMock import pytest +import respx +from starlette.requests import Request +import litellm from litellm.proxy._types import ( ConfigGeneralSettings, LitellmUserRoles, @@ -230,3 +233,90 @@ def test_configured_custom_key_header_is_stripped() -> None: assert "x-my-tenant-key" in _request_strip_headers() finally: proxy_server.general_settings = original + + +def _plugin_proxy_request(headers: dict[str, str]) -> Request: + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": b'{"q": 1}', "more_body": False} + + return Request( + { + "type": "http", + "method": "POST", + "path": "/plugin-proxy/p/items", + "query_string": b"page=2", + "headers": [(name.lower().encode(), value.encode()) for name, value in headers.items()], + }, + receive, + ) + + +@pytest.mark.parametrize( + ("user_id", "expected_identity_headers"), + [ + ("admin-1", {"x-litellm-user-id": "admin-1", "x-litellm-user-role": "proxy_admin"}), + (None, {"x-litellm-user-role": "proxy_admin"}), + ("", {"x-litellm-user-role": "proxy_admin"}), + ], +) +def test_plugin_proxy_forwards_caller_identity_and_replaces_the_caller_credential( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + user_id: str | None, + expected_identity_headers: dict[str, str], +) -> None: + from litellm.proxy.plugin_routes import plugin_proxy + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + register_plugins_from_config( + {"plugins": [{"name": "p", "url": "http://plugin.example", "plugin_key": "sk-plugin"}]} + ) + upstream = respx_mock.post("http://plugin.example/items?page=2").respond(text="plugin says hi") + + response = asyncio.run( + plugin_proxy( + "p", + "items", + _plugin_proxy_request({"Authorization": "Bearer sk-caller", "X-Trace-Id": "t-1"}), + UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN, user_id=user_id), + ) + ) + register_plugins_from_config({}) + litellm.in_memory_llm_clients_cache.flush_cache() + + forwarded = upstream.calls[0].request + assert response.status_code == 200 + assert response.body == b"plugin says hi" + assert forwarded.content == b'{"q": 1}' + assert forwarded.headers["authorization"] == "Bearer sk-plugin" + assert forwarded.headers["x-trace-id"] == "t-1" + assert { + name: value for name, value in forwarded.headers.items() if name.startswith("x-litellm-user-") + } == expected_identity_headers + + +def test_plugin_proxy_refuses_non_admin_callers_without_contacting_the_plugin( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + from litellm.proxy.plugin_routes import plugin_proxy + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + register_plugins_from_config({"plugins": [{"name": "p", "url": "http://plugin.example"}]}) + upstream = respx_mock.post("http://plugin.example/items?page=2").respond(text="plugin says hi") + + response = asyncio.run( + plugin_proxy( + "p", + "items", + _plugin_proxy_request({"Authorization": "Bearer sk-caller"}), + UserAPIKeyAuth(api_key="sk-user", user_role=LitellmUserRoles.INTERNAL_USER, user_id="user-1"), + ) + ) + register_plugins_from_config({}) + + assert response.status_code == 403 + assert response.body == b"Plugin proxy access requires proxy_admin role." + assert not upstream.called diff --git a/tests/unit/proxy/types_utils/test_utils.py b/tests/unit/proxy/types_utils/test_utils.py new file mode 100644 index 00000000000..eb7a57ba37e --- /dev/null +++ b/tests/unit/proxy/types_utils/test_utils.py @@ -0,0 +1,250 @@ +import io +import re +import sys +import tempfile +from collections.abc import Iterator, Mapping +from pathlib import Path + +import boto3 +import pytest + +from litellm.proxy.types_utils.utils import get_instance_fn + +_PLUGIN_SOURCE = """ +class Handler: + label = "plugin-handler" + + +handler = Handler() +alias = handler +nothing = None +zero = 0 +disabled = False +empty = [] +label = "plugin" +location = __file__ + + +def build(): + return "built" +""" + +_DYNAMIC_PLUGIN_SOURCE = """ +def __getattr__(name): + if name.startswith("make_"): + return name.removeprefix("make_").upper() + if name == "needs_dependency": + raise ImportError("optional dependency missing") + raise LookupError(f"plugin cannot provide {name}") +""" + +_WARNING_PLUGIN_TEMPLATE = """ +import warnings + + +def __getattr__(name): + warnings.warn(f"{{name}} is deprecated", DeprecationWarning, stacklevel={stacklevel}) + return name.upper() +""" + +_PLAIN_ATTRIBUTES = [ + pytest.param("nothing", "NoneType", "None", id="none"), + pytest.param("zero", "int", "0", id="zero"), + pytest.param("disabled", "bool", "False", id="false"), + pytest.param("empty", "list", "[]", id="empty-list"), + pytest.param("label", "str", "'plugin'", id="string"), +] + +_INSTALLED_PLUGIN = "types_utils_installed_plugin" + + +class _FakeS3: + def __init__(self, objects: Mapping[tuple[str, str], bytes]) -> None: + self.objects = objects + self.requests: list[tuple[str, str]] = [] + + def get_object(self, Bucket: str, Key: str) -> Mapping[str, io.BytesIO]: + self.requests.append((Bucket, Key)) + return {"Body": io.BytesIO(self.objects[(Bucket, Key)])} + + +@pytest.fixture +def local_config(tmp_path: Path) -> str: + (tmp_path / "local_plugin.py").write_text(_PLUGIN_SOURCE) + (tmp_path / "dynamic_plugin.py").write_text(_DYNAMIC_PLUGIN_SOURCE) + return str(tmp_path / "config.yaml") + + +@pytest.fixture +def installed_plugin(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + site = tmp_path / "site" + site.mkdir() + (site / f"{_INSTALLED_PLUGIN}.py").write_text(_PLUGIN_SOURCE) + monkeypatch.syspath_prepend(str(site)) + yield _INSTALLED_PLUGIN + sys.modules.pop(_INSTALLED_PLUGIN, None) + + +@pytest.fixture +def downloads(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + directory = tmp_path / "downloads" + directory.mkdir() + monkeypatch.setattr(tempfile, "tempdir", str(directory)) + return directory + + +@pytest.fixture +def plugin_bucket(downloads: Path, monkeypatch: pytest.MonkeyPatch) -> _FakeS3: + bucket = _FakeS3( + { + ("plugins-bucket", "loggers/remote_plugin.py"): _PLUGIN_SOURCE.encode(), + ("plugins-bucket", "loggers/warning_plugin.py"): _WARNING_PLUGIN_TEMPLATE.format(stacklevel=4).encode(), + } + ) + + def fake_client( + service_name: str, + aws_access_key_id: str, + aws_secret_access_key: str, + aws_session_token: str | None, + ) -> _FakeS3: + assert service_name == "s3" + return bucket + + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "unit-test") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "unit-test") + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + monkeypatch.setattr(boto3, "client", fake_client) + return bucket + + +@pytest.mark.parametrize(("attribute", "expected_type", "expected_repr"), _PLAIN_ATTRIBUTES) +def test_get_instance_fn_returns_the_attribute_of_the_module_beside_the_config_unchanged( + local_config: str, attribute: str, expected_type: str, expected_repr: str +): + result = get_instance_fn(f"local_plugin.{attribute}", config_file_path=local_config) + + assert (type(result).__name__, repr(result)) == (expected_type, expected_repr) + + +def test_get_instance_fn_returns_instances_and_callables_defined_in_the_local_module(local_config: str): + handler = get_instance_fn("local_plugin.handler", config_file_path=local_config) + build = get_instance_fn("local_plugin.build", config_file_path=local_config) + + assert (type(handler).__name__, handler.label) == ("Handler", "plugin-handler") + assert build() == "built" + + +def test_get_instance_fn_surfaces_a_name_the_local_module_lacks_as_attribute_error(local_config: str): + with pytest.raises(AttributeError) as raised: + get_instance_fn("local_plugin.missing", config_file_path=local_config) + + assert type(raised.value) is AttributeError + assert (raised.value.obj.__name__, raised.value.name) == ("local_plugin", "missing") + assert raised.value.__cause__ is None + + +@pytest.mark.parametrize( + ("attribute", "expected"), + [ + pytest.param("make_logger", "LOGGER", id="computed"), + pytest.param("make_", "", id="computed-empty"), + ], +) +def test_get_instance_fn_returns_what_a_module_level_getattr_computes(local_config: str, attribute: str, expected: str): + result = get_instance_fn(f"dynamic_plugin.{attribute}", config_file_path=local_config) + + assert (type(result), result) == (str, expected) + + +def test_get_instance_fn_lets_an_error_from_a_module_level_getattr_through_unchanged(local_config: str): + with pytest.raises(LookupError, match=r"^plugin cannot provide other$") as raised: + get_instance_fn("dynamic_plugin.other", config_file_path=local_config) + + assert type(raised.value) is LookupError + + +def test_get_instance_fn_rewords_an_import_error_raised_while_reading_the_attribute(local_config: str): + with pytest.raises(ImportError, match=r"^Could not import needs_dependency from dynamic_plugin$") as raised: + get_instance_fn("dynamic_plugin.needs_dependency", config_file_path=local_config) + + assert type(raised.value) is ImportError + assert (type(raised.value.__cause__), str(raised.value.__cause__)) == (ImportError, "optional dependency missing") + + +def test_get_instance_fn_returns_the_very_object_an_installed_module_holds(installed_plugin: str): + handler = get_instance_fn(f"{installed_plugin}.handler") + alias = get_instance_fn(f"{installed_plugin}.alias") + + assert handler is sys.modules[installed_plugin].handler + assert alias is handler + + +def test_get_instance_fn_surfaces_a_name_an_installed_module_lacks_as_attribute_error(installed_plugin: str): + with pytest.raises(AttributeError) as raised: + get_instance_fn(f"{installed_plugin}.missing") + + assert type(raised.value) is AttributeError + assert (raised.value.obj, raised.value.name) == (sys.modules[installed_plugin], "missing") + assert raised.value.__cause__ is None + + +@pytest.mark.parametrize(("attribute", "expected_type", "expected_repr"), _PLAIN_ATTRIBUTES) +def test_get_instance_fn_loads_the_attribute_from_the_bucket_object_unchanged( + plugin_bucket: _FakeS3, attribute: str, expected_type: str, expected_repr: str +): + result = get_instance_fn( + f"s3://plugins-bucket/loggers/remote_plugin.{attribute}", config_file_path="/etc/litellm/config.yaml" + ) + + assert (type(result).__name__, repr(result)) == (expected_type, expected_repr) + assert plugin_bucket.requests == [("plugins-bucket", "loggers/remote_plugin.py")] + + +def test_get_instance_fn_runs_the_bucket_module_from_a_temporary_file_and_removes_it( + plugin_bucket: _FakeS3, downloads: Path +): + location = get_instance_fn( + "s3://plugins-bucket/loggers/remote_plugin.location", config_file_path="/etc/litellm/config.yaml" + ) + + assert (Path(location).parent, Path(location).suffix) == (downloads, ".py") + assert not Path(location).exists() + + +def test_get_instance_fn_reports_a_name_the_bucket_module_lacks_as_a_failed_load(plugin_bucket: _FakeS3): + prefix = "Failed to load custom logger from s3://plugins-bucket/loggers/remote_plugin.missing: " + + with pytest.raises(ImportError, match=f"^{re.escape(prefix)}") as raised: + get_instance_fn( + "s3://plugins-bucket/loggers/remote_plugin.missing", config_file_path="/etc/litellm/config.yaml" + ) + + assert type(raised.value) is ImportError + assert type(raised.value.__cause__) is AttributeError + assert str(raised.value) == prefix + str(raised.value.__cause__) + assert (raised.value.__cause__.obj.__name__, raised.value.__cause__.name) == ("loggers/remote_plugin", "missing") + assert plugin_bucket.requests == [("plugins-bucket", "loggers/remote_plugin.py")] + + +def test_get_instance_fn_attributes_a_stacklevel_three_warning_from_a_local_plugin_to_its_own_caller(tmp_path: Path): + (tmp_path / "warning_plugin.py").write_text(_WARNING_PLUGIN_TEMPLATE.format(stacklevel=3)) + + with pytest.warns(DeprecationWarning, match=r"^old_handler is deprecated$") as recorded: + result = get_instance_fn("warning_plugin.old_handler", config_file_path=str(tmp_path / "config.yaml")) + + assert result == "OLD_HANDLER" + assert [Path(w.filename) for w in recorded if str(w.message) == "old_handler is deprecated"] == [Path(__file__)] + + +def test_get_instance_fn_attributes_a_stacklevel_four_warning_from_a_bucket_plugin_to_its_own_caller( + plugin_bucket: _FakeS3, +): + with pytest.warns(DeprecationWarning, match=r"^old_handler is deprecated$") as recorded: + result = get_instance_fn( + "s3://plugins-bucket/loggers/warning_plugin.old_handler", config_file_path="/etc/litellm/config.yaml" + ) + + assert result == "OLD_HANDLER" + assert [Path(w.filename) for w in recorded if str(w.message) == "old_handler is deprecated"] == [Path(__file__)] + assert plugin_bucket.requests == [("plugins-bucket", "loggers/warning_plugin.py")] diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_utils.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_utils.py new file mode 100644 index 00000000000..5b71fcb8037 --- /dev/null +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_utils.py @@ -0,0 +1,94 @@ +import json + +import pytest + +import litellm +from litellm.proxy.vector_store_endpoints.utils import get_litellm_managed_vector_store +from litellm.vector_stores.vector_store_registry import VectorStoreRegistry + + +def _registered_store(litellm_params: object) -> dict[str, object]: + return { + "vector_store_id": "vs_1", + "custom_llm_provider": "bedrock", + "vector_store_name": "kb", + "vector_store_metadata": '{"source": "db"}', + "team_id": None, + "litellm_params": litellm_params, + } + + +@pytest.mark.parametrize( + ("stored_params", "expected_params"), + [ + ( + '{"api_base": "https://example.com", "embedding": {"api_key": "sk-1", "dims": [1, null]}}', + {"api_base": "https://example.com", "embedding": {"api_key": "sk-1", "dims": [1, None]}}, + ), + ( + '{"top_k": 12, "score_threshold": 12.0, "api_version": "12", "stream": true, "verify": false,' + ' "region": null, "max_tokens": 123456789012345678901234567890}', + { + "top_k": 12, + "score_threshold": 12.0, + "api_version": "12", + "stream": True, + "verify": False, + "region": None, + "max_tokens": 123456789012345678901234567890, + }, + ), + ('{"z": 1, "a": 2, "a": 3}', {"z": 1, "a": 3}), + ("{}", {}), + ], +) +async def test_registry_store_with_json_object_string_params_resolves_to_decoded_params( + monkeypatch: pytest.MonkeyPatch, stored_params: str, expected_params: dict[str, object] +): + store = _registered_store(stored_params) + monkeypatch.setattr(litellm, "vector_store_registry", VectorStoreRegistry([store])) + + resolved = await get_litellm_managed_vector_store("vs_1") + + assert resolved == {**store, "litellm_params": expected_params} + assert json.dumps(resolved["litellm_params"]) == json.dumps(expected_params) + assert store["litellm_params"] == stored_params + + +@pytest.mark.parametrize( + "stored_params", + [ + "[1, 2]", + '[["api_base", "https://example.com"]]', + '"{\\"api_base\\": \\"https://example.com\\"}"', + "12", + "12.5", + "true", + "null", + "{not json", + "", + ], +) +async def test_registry_store_with_non_object_or_invalid_json_string_params_resolves_to_empty_params( + monkeypatch: pytest.MonkeyPatch, stored_params: str +): + store = _registered_store(stored_params) + monkeypatch.setattr(litellm, "vector_store_registry", VectorStoreRegistry([store])) + + resolved = await get_litellm_managed_vector_store("vs_1") + + assert resolved == {**store, "litellm_params": {}} + assert store["litellm_params"] == stored_params + + +@pytest.mark.parametrize("stored_params", [{"api_base": "https://example.com"}, {}, None]) +async def test_registry_store_with_non_string_params_is_returned_as_is( + monkeypatch: pytest.MonkeyPatch, stored_params: object +): + store = _registered_store(stored_params) + monkeypatch.setattr(litellm, "vector_store_registry", VectorStoreRegistry([store])) + + resolved = await get_litellm_managed_vector_store("vs_1") + + assert resolved is store + assert resolved["litellm_params"] is stored_params diff --git a/tests/unit/secret_managers/test_google_secret_manager.py b/tests/unit/secret_managers/test_google_secret_manager.py index ed93aa2861d..2051be08f7c 100644 --- a/tests/unit/secret_managers/test_google_secret_manager.py +++ b/tests/unit/secret_managers/test_google_secret_manager.py @@ -1,12 +1,21 @@ import base64 +import binascii +import json from dataclasses import dataclass from typing import Final +import google.auth +import httpx import litellm +import litellm.proxy.proxy_server import pytest import respx +from google.auth.credentials import Credentials +from pydantic import ValidationError from litellm.secret_managers.google_secret_manager import GoogleSecretManager +from litellm.secret_managers.main import get_secret +from litellm.types.secret_managers.main import KeyManagementSettings @dataclass(frozen=True, slots=True) @@ -81,3 +90,237 @@ def test_google_secret_manager_returns_cached_values_without_http( assert missing_value is None assert cached_value == "lite-llm" assert upstream.calls.call_count == 0 + + +PROJECT_ID: Final = "unit-test-project" +SECRET_NAME: Final = "UNIT_TEST_GOOGLE_SECRET" +ACCESS_URL: Final = ( + f"https://secretmanager.googleapis.com/v1/projects/{PROJECT_ID}/secrets/{SECRET_NAME}/versions/latest:access" +) +SECRET_VALUE: Final = "s3cr3t-value" +ENCODED_SECRET: Final = base64.b64encode(SECRET_VALUE.encode()).decode("ascii") + + +class _StaticCredentials(Credentials): + def refresh(self, request: object) -> None: + self.token = "unit-test-token" + + +def _default_credentials(scopes: object = None) -> tuple[Credentials, str]: + return _StaticCredentials(), PROJECT_ID + + +def _access_response(body: object) -> httpx.Response: + return httpx.Response(200, content=json.dumps(body).encode()) + + +@pytest.fixture +def manager(monkeypatch: pytest.MonkeyPatch) -> GoogleSecretManager: + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + monkeypatch.setattr(litellm, "secret_manager_client", None) + monkeypatch.setattr(litellm, "_key_management_system", None) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings()) + monkeypatch.setattr(litellm.vertex_chat_completion, "_credentials_project_mapping", {}) + monkeypatch.setattr(google.auth, "default", _default_credentials) + monkeypatch.setenv("GOOGLE_SECRET_MANAGER_PROJECT_ID", PROJECT_ID) + monkeypatch.delenv("GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER", raising=False) + monkeypatch.delenv("GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL", raising=False) + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_MOCK", raising=False) + monkeypatch.delenv(SECRET_NAME, raising=False) + return GoogleSecretManager() + + +@pytest.mark.parametrize( + ("body", "expected"), + [ + ({"payload": {"data": ENCODED_SECRET}}, SECRET_VALUE), + ( + { + "name": f"projects/123456/secrets/{SECRET_NAME}/versions/7", + "payload": {"data": ENCODED_SECRET, "dataCrc32c": "2932599711"}, + }, + SECRET_VALUE, + ), + ({"payload": {"data": ENCODED_SECRET, "labels": {"nested": [1, None]}}, "replication": [1, 2]}, SECRET_VALUE), + ({"payload": {"data": "czNj\ncjN0 LXZhbHVl"}}, SECRET_VALUE), + ({"payload": {"data": base64.b64encode("clé ☃".encode()).decode("ascii")}}, "clé ☃"), + ({"payload": {"data": ""}}, ""), + ], +) +@respx.mock +def test_secret_is_decoded_from_the_access_response(manager: GoogleSecretManager, body: object, expected: str) -> None: + route: Final = respx.get(ACCESS_URL).mock(return_value=_access_response(body)) + + secret: Final = manager.get_secret_from_google_secret_manager(SECRET_NAME) + + assert secret == expected + assert type(secret) is str + assert manager.cache.cache_dict[SECRET_NAME] is secret + assert route.calls.last.request.headers["Authorization"] == "Bearer unit-test-token" + + +@pytest.mark.parametrize( + "body", + [ + {}, + {"name": f"projects/123456/secrets/{SECRET_NAME}/versions/7"}, + {"payload": {}}, + {"payload": {"dataCrc32c": "2932599711"}}, + {"payload": {"data": None}}, + ], +) +@respx.mock +def test_response_without_secret_data_reports_not_found_and_caches_the_miss( + manager: GoogleSecretManager, body: object +) -> None: + respx.get(ACCESS_URL).mock(return_value=_access_response(body)) + + with pytest.raises(ValueError, match=rf"^secret {SECRET_NAME} not found in Google Secret Manager$") as raised: + manager.get_secret_from_google_secret_manager(SECRET_NAME) + + assert type(raised.value) is ValueError + assert manager.cache.cache_dict[SECRET_NAME] is None + assert manager.get_secret_from_google_secret_manager(SECRET_NAME) is None + + +@pytest.mark.parametrize( + "body", + [ + [], + [{"payload": {"data": ENCODED_SECRET}}], + [["payload", {"data": ENCODED_SECRET}]], + ENCODED_SECRET, + 7, + None, + True, + {"payload": None}, + {"payload": []}, + {"payload": [["data", ENCODED_SECRET]]}, + {"payload": ENCODED_SECRET}, + {"payload": ""}, + {"payload": 0}, + {"payload": False}, + ], +) +@respx.mock +def test_response_that_is_not_a_json_object_is_rejected_without_caching( + manager: GoogleSecretManager, body: object +) -> None: + respx.get(ACCESS_URL).mock(return_value=_access_response(body)) + + with pytest.raises(ValidationError) as raised: + manager.get_secret_from_google_secret_manager(SECRET_NAME) + + assert [error["type"] for error in raised.value.errors()] == ["dict_type"] + assert ENCODED_SECRET not in str(raised.value) + assert SECRET_NAME not in manager.cache.cache_dict + + +@pytest.mark.parametrize("data", [5, 0, 1.5, True, False, [], [ENCODED_SECRET], {}, {"value": ENCODED_SECRET}]) +@respx.mock +def test_secret_data_that_is_not_text_is_rejected_without_caching(manager: GoogleSecretManager, data: object) -> None: + respx.get(ACCESS_URL).mock(return_value=_access_response({"payload": {"data": data}})) + + with pytest.raises(ValidationError) as raised: + manager.get_secret_from_google_secret_manager(SECRET_NAME) + + assert [error["type"] for error in raised.value.errors()] == ["string_type"] + assert ENCODED_SECRET not in str(raised.value) + assert SECRET_NAME not in manager.cache.cache_dict + + +@pytest.mark.parametrize( + "body", + [ + [{"payload": {"data": ENCODED_SECRET}}] * 403 + [7], + {"payload": [{"data": ENCODED_SECRET}] * 403 + [None]}, + {"payload": {"data": [ENCODED_SECRET] * 403 + [429]}}, + {"payload": {"data": {str(position): "Request timed out" for position in range(403, 430)}}}, + ], +) +@respx.mock +def test_rejection_of_a_long_response_names_no_position_or_content_from_it( + manager: GoogleSecretManager, body: object +) -> None: + respx.get(ACCESS_URL).mock(return_value=_access_response(body)) + + with pytest.raises(ValidationError) as raised: + manager.get_secret_from_google_secret_manager(SECRET_NAME) + + message: Final = str(raised.value) + assert [error["loc"] for error in raised.value.errors()] == [()] + assert [fragment for fragment in ("403", "429", "Request timed out", ENCODED_SECRET) if fragment in message] == [] + assert SECRET_NAME not in manager.cache.cache_dict + + +@pytest.mark.parametrize( + ("data", "error_class"), + [ + ("abc", binascii.Error), + (base64.b64encode(b"\xff\xfe").decode("ascii"), UnicodeDecodeError), + ("éééé", ValueError), + ], +) +@respx.mock +def test_undecodable_secret_text_keeps_raising_the_decoder_error( + manager: GoogleSecretManager, data: str, error_class: type[Exception] +) -> None: + respx.get(ACCESS_URL).mock(return_value=_access_response({"payload": {"data": data}})) + + with pytest.raises(error_class) as raised: + manager.get_secret_from_google_secret_manager(SECRET_NAME) + + assert type(raised.value) is error_class + assert SECRET_NAME not in manager.cache.cache_dict + + +@pytest.mark.usefixtures("manager") +@respx.mock +def test_get_secret_reads_through_the_configured_manager() -> None: + respx.get(ACCESS_URL).mock(return_value=_access_response({"payload": {"data": ENCODED_SECRET}})) + + assert get_secret(SECRET_NAME) == SECRET_VALUE + + +@pytest.mark.usefixtures("manager") +@pytest.mark.parametrize( + "body", + [ + [], + {"payload": None}, + {"payload": ENCODED_SECRET}, + {"payload": {"data": 5}}, + {"payload": {"data": [ENCODED_SECRET]}}, + {"payload": {}}, + [{"payload": {"data": ENCODED_SECRET}}] * 403 + [7], + {"payload": {"data": [ENCODED_SECRET] * 403 + [429]}}, + {"payload": "Request timed out"}, + ], +) +@respx.mock +def test_get_secret_falls_back_to_the_environment_when_the_response_is_unusable( + monkeypatch: pytest.MonkeyPatch, body: object +) -> None: + monkeypatch.setenv(SECRET_NAME, "from-environment") + respx.get(ACCESS_URL).mock(return_value=_access_response(body)) + + assert get_secret(SECRET_NAME) == "from-environment" + + +@pytest.mark.usefixtures("manager") +@pytest.mark.parametrize( + "body", + [ + [{"payload": {"data": ENCODED_SECRET}}] * 403 + [7], + {"payload": None}, + {"payload": {"data": [ENCODED_SECRET] * 403 + [429]}}, + {"payload": {"data": 429}}, + {"payload": {}}, + ], +) +@respx.mock +def test_get_secret_returns_none_when_the_response_is_unusable_and_the_environment_has_no_value(body: object) -> None: + respx.get(ACCESS_URL).mock(return_value=_access_response(body)) + + assert get_secret(SECRET_NAME) is None From b77108a4091fe847e2eac7d70c3339fba52ebd06 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 13 Oct 2026 18:38:46 -0400 Subject: [PATCH 020/325] feat(ui): simplify LiteAdmin on the homepage (#46301) * feat(ui): simplify LiteAdmin on the homepage * fix(ui): preserve contrast on LiteAdmin primary actions --- .../_components/HomePage.integration.test.tsx | 2 +- .../(dashboard)/home/_components/HomePage.tsx | 9 +- .../liteadmin/LiteAdmin.integration.test.tsx | 57 ++++++--- .../src/components/liteadmin/LiteAdmin.tsx | 118 ++++++++++++------ .../liteadmin/LiteAdminHomeComposer.tsx | 92 ++++++++++++++ 5 files changed, 224 insertions(+), 54 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/liteadmin/LiteAdminHomeComposer.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx index 598fbf74efc..b957393be64 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.integration.test.tsx @@ -150,7 +150,7 @@ describe("HomePage", () => { , ); const entry = screen.getByRole("region", { name: "Ask LiteAdmin" }); - expect(await within(entry).findByPlaceholderText("Ask LiteAdmin…")).toBeInTheDocument(); + expect(await within(entry).findByPlaceholderText("What would you like to do?")).toBeInTheDocument(); expect(within(entry).getByRole("button", { name: "Send message" })).toBeDisabled(); expect(await screen.findByText("$501.74")).toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.tsx index b0c63bb5dc0..22eb0a07f32 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/home/_components/HomePage.tsx @@ -29,11 +29,14 @@ export default function HomePage() { return ( - + Home {!isViewOnly && ( - - diff --git a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx index 2bc586cd1bb..690721dbbd9 100644 --- a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx @@ -184,7 +184,7 @@ async function openWidget() { } function send(text: string) { - fireEvent.change(screen.getByPlaceholderText("Ask LiteAdmin…"), { target: { value: text } }); + fireEvent.change(screen.getByRole("textbox"), { target: { value: text } }); fireEvent.click(screen.getByRole("button", { name: "Send message" })); } @@ -209,7 +209,7 @@ describe("LiteAdmin in the gateway", () => { renderWidget(undefined, true); const home = await screen.findByRole("region", { name: "Ask LiteAdmin" }); await selectModel(); - const input = within(home).getByPlaceholderText("Ask LiteAdmin…"); + const input = within(home).getByPlaceholderText("What would you like to do?"); fireEvent.change(input, { target: { value: "First question" } }); if (method === "keyboard") fireEvent.keyDown(input, { key: "Enter" }); else fireEvent.click(within(home).getByRole("button", { name: "Send message" })); @@ -235,6 +235,31 @@ describe("LiteAdmin in the gateway", () => { ]); }); + it("prefills each Home suggestion without sending until the edited draft is submitted", async () => { + const requests = gateway([answer("Budget details received.")]); + renderWidget(undefined, true); + await selectModel(); + const input = screen.getByRole("textbox", { name: "Message LiteAdmin" }); + for (const label of ["Create a key", "Add a model", "Set a budget"]) { + fireEvent.click(screen.getByRole("button", { name: label, exact: true })); + expect(input).toHaveValue(label); + expect(input).toHaveFocus(); + expect(screen.queryByRole("complementary", { name: "LiteAdmin" })).not.toBeInTheDocument(); + expect(requests.filter((request) => request.url.endsWith("/chat/completions"))).toHaveLength(0); + } + fireEvent.change(input, { target: { value: "Set a budget of $10 for my team" } }); + fireEvent.keyDown(input, { key: "Enter", shiftKey: true }); + fireEvent.keyDown(input, { key: "Enter", isComposing: true }); + expect(requests.filter((request) => request.url.endsWith("/chat/completions"))).toHaveLength(0); + fireEvent.keyDown(input, { key: "Enter" }); + expect(await screen.findByText("Budget details received.")).toBeInTheDocument(); + const completions = requests.filter((request) => request.url.endsWith("/chat/completions")); + expect(completions).toHaveLength(1); + expect(completions[0].body.model).toBe("chat-model"); + expect(completions[0].body.messages).toContainEqual({ role: "user", content: "Set a budget of $10 for my team" }); + expect(requests.filter((request) => request.url.endsWith("/key/generate"))).toHaveLength(0); + }); + it("preserves the draft and chosen model when opening, closing and leaving Home", async () => { const requests = gateway([answer("Draft received.")], { models: [ @@ -244,12 +269,14 @@ describe("LiteAdmin in the gateway", () => { }); const view = renderWidget(undefined, true); await selectModel(); - fireEvent.change(screen.getByPlaceholderText("Ask LiteAdmin…"), { target: { value: "Keep this draft" } }); - fireEvent.click(screen.getByRole("button", { name: "LiteAdmin" })); + fireEvent.change(screen.getByPlaceholderText("What would you like to do?"), { + target: { value: "Keep this draft" }, + }); + fireEvent.click(screen.getByRole("button", { name: "Recent chats" })); expect(screen.getByPlaceholderText("Ask LiteAdmin…")).toHaveValue("Keep this draft"); fireEvent.click(screen.getByRole("button", { name: "Close LiteAdmin" })); const home = screen.getByRole("region", { name: "Ask LiteAdmin" }); - expect(within(home).getByPlaceholderText("Ask LiteAdmin…")).toHaveValue("Keep this draft"); + expect(within(home).getByPlaceholderText("What would you like to do?")).toHaveValue("Keep this draft"); view.refresh(false); expect(home).not.toBeInTheDocument(); fireEvent.click(screen.getByRole("button", { name: "LiteAdmin" })); @@ -293,16 +320,16 @@ describe("LiteAdmin in the gateway", () => { fireEvent.change(screen.getByPlaceholderText("Ask LiteAdmin…"), { target: { value: "First draft" } }); fireEvent.click(screen.getByRole("button", { name: "Close LiteAdmin" })); const home = screen.getByRole("region", { name: "Ask LiteAdmin" }); - expect(within(home).getByPlaceholderText("Ask LiteAdmin…")).toHaveValue("First draft"); - expect(screen.getAllByPlaceholderText("Ask LiteAdmin…")).toHaveLength(1); + expect(within(home).getByPlaceholderText("What would you like to do?")).toHaveValue("First draft"); + expect(screen.getAllByPlaceholderText("What would you like to do?")).toHaveLength(1); - fireEvent.click(screen.getByRole("button", { name: "LiteAdmin" })); + fireEvent.click(screen.getByRole("button", { name: "Recent chats" })); fireEvent.click(screen.getByRole("tab", { name: "Chat 2" })); expect(screen.getByPlaceholderText("Ask LiteAdmin…")).toHaveValue("Second draft"); expect(screen.getByRole("combobox", { name: "LiteAdmin model" })).toHaveValue("other-chat-model"); expect(screen.getAllByLabelText("LiteAdmin conversation")).toHaveLength(1); fireEvent.click(screen.getByRole("button", { name: "Close LiteAdmin" })); - expect(within(home).getByPlaceholderText("Ask LiteAdmin…")).toHaveValue("Second draft"); + expect(within(home).getByPlaceholderText("What would you like to do?")).toHaveValue("Second draft"); fireEvent.click(within(home).getByRole("button", { name: "Send message" })); await screen.findByText("Second follow-up."); expect(screen.getByRole("tab", { name: "Chat 2" })).toHaveAttribute("aria-selected", "true"); @@ -531,13 +558,13 @@ describe("LiteAdmin in the gateway", () => { const { client } = renderWidget(undefined, true); const home = await screen.findByRole("region", { name: "Ask LiteAdmin" }); expect(await within(home).findByText("Add a chat model to your gateway to use LiteAdmin.")).toBeInTheDocument(); - expect(within(home).getByPlaceholderText("Ask LiteAdmin…")).toBeDisabled(); + expect(within(home).getByPlaceholderText("What would you like to do?")).toBeDisabled(); expect(within(home).getByRole("button", { name: "Send message" })).toBeDisabled(); options.models = [{ model_group: "chat-model", mode: "chat" }]; await act(async () => client.invalidateQueries({ queryKey: ["liteadmin-models"] })); await selectModel(); for (const text of [" ", "x".repeat(MAX_INPUT_LENGTH + 1)]) { - const input = within(home).getByPlaceholderText("Ask LiteAdmin…"); + const input = within(home).getByPlaceholderText("What would you like to do?"); fireEvent.change(input, { target: { value: text } }); fireEvent.keyDown(input, { key: "Enter" }); const submit = within(home).getByRole("button", { name: "Send message" }); @@ -566,15 +593,17 @@ describe("LiteAdmin in the gateway", () => { await act(async () => settingsReady.resolve()); const consent = await within(home).findByRole("button", { name: "Use configured gateway" }); expect(within(home).getByText(EXTERNAL_INFERENCE)).toBeInTheDocument(); - expect(within(home).queryByPlaceholderText("Ask LiteAdmin…")).not.toBeInTheDocument(); + expect(within(home).queryByPlaceholderText("What would you like to do?")).not.toBeInTheDocument(); fireEvent.click(consent); await selectModel(); - fireEvent.change(within(home).getByPlaceholderText("Ask LiteAdmin…"), { target: { value: "Old gateway draft" } }); + fireEvent.change(within(home).getByPlaceholderText("What would you like to do?"), { + target: { value: "Old gateway draft" }, + }); settings.target = `${EXTERNAL_INFERENCE}/changed`; await act(async () => client.invalidateQueries({ queryKey: ["proxySettings"] })); fireEvent.click(await within(home).findByRole("button", { name: "Use configured gateway" })); await selectModel(); - expect(within(home).getByPlaceholderText("Ask LiteAdmin…")).toHaveValue(""); + expect(within(home).getByPlaceholderText("What would you like to do?")).toHaveValue(""); expect(home).toBeVisible(); expect(requests.some((request) => request.url.endsWith("/chat/completions"))).toBe(false); }); diff --git a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.tsx b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.tsx index bf296d47746..350379588c1 100644 --- a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.tsx +++ b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.tsx @@ -8,12 +8,13 @@ import { useLayoutEffect, useRef, useState, + type ComponentProps, type ReactNode, type RefObject, } from "react"; import { createPortal } from "react-dom"; import { useQuery, type UseQueryResult } from "@tanstack/react-query"; -import { Plus, Sparkles, X } from "lucide-react"; +import { History, Plus, Sparkles, X } from "lucide-react"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useDisableLiteAdmin } from "@/app/(dashboard)/hooks/useDisableLiteAdmin"; import { useProxySettingsQuery } from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; @@ -32,6 +33,7 @@ import { cn } from "@/lib/cva.config"; import { isProxyAdminRole } from "@/utils/roles"; import { getPreferredLiteAdminModel, MAX_INPUT_LENGTH, resolveInferenceTarget } from "./agent"; import { LiteAdminConversation } from "./LiteAdminConversation"; +import { LiteAdminHomeComposer } from "./LiteAdminHomeComposer"; import { useLiteAdmin, type ConversationEntry, type LiteAdminSession } from "./useLiteAdmin"; type ManagementSession = Omit; @@ -100,17 +102,18 @@ export function LiteAdminHome() { if (!state) return null; const { open, show, setHomeContainer } = state; return ( -