From 0ac0362b42890d5d40698dbd5c206065b0015e5f Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Fri, 18 Sep 2026 21:01:16 +0000 Subject: [PATCH 01/47] fix(proxy): enforce virtual key budgets for JEV test routing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../auto_router_endpoints.py | 2 +- .../test_auto_router_endpoints.py | 76 ++++++++++++++++++- 2 files changed, 73 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 200ed6c3bf3..f79425d2e97 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -319,7 +319,7 @@ async def _authorize_models_this_test_can_call( its calls through the proxy. Team and member budgets are already enforced on every route. """ models: Final = _models_this_test_can_call(config) - if not models: + if not models and config.classifier_type != "jev": return from litellm.proxy.proxy_server import proxy_logging_obj diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 067f30c2fd7..36130137c64 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -3,23 +3,33 @@ Unit tests for auto router management endpoints """ from collections.abc import Mapping, Sequence +from functools import partial from pathlib import Path from typing import Final +from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException, Request from pydantic import ValidationError +from litellm.proxy import proxy_server from litellm.proxy._types import ( LitellmUserRoles, ProxyErrorTypes, ProxyException, UserAPIKeyAuth, ) +from litellm.proxy.management_endpoints import auto_router_endpoints from litellm.proxy.management_endpoints.auto_router_endpoints import ( preview_auto_router_routing, ) from litellm.router import Router +from litellm.router_strategy.complexity_router import ComplexityRouter +from litellm.router_strategy.complexity_router.jev_classifier import ( + JevChoiceAnswer, + JevClassifierClient, + JevSystemOneResponse, +) from litellm.types.management_endpoints.auto_router_endpoints import ( AutoRouterBenchmarksResponse, AutoRouterRoutingTestRequest, @@ -422,8 +432,67 @@ async def test_a_key_over_its_budget_cannot_run_a_classifier_config(monkeypatch: assert calls == [] +@pytest.mark.parametrize( + "max_budget, spend, denied", + ( + pytest.param(0.0, 0.0, True, id="zero-budget"), + pytest.param(1.0, 1.0, True, id="budget-reached"), + pytest.param(1.0, 2.0, True, id="budget-exceeded"), + pytest.param(1.0, 0.5, False, id="budget-remaining"), + pytest.param(None, 2.0, False, id="unlimited"), + ), +) @pytest.mark.asyncio -async def test_a_heuristic_config_does_not_need_a_budget(monkeypatch: pytest.MonkeyPatch): +async def test_jev_test_routing_enforces_key_budget_before_provider_invocation( + monkeypatch: pytest.MonkeyPatch, max_budget: float | None, spend: float, denied: bool +) -> None: + client: Final = AsyncMock(spec=JevClassifierClient) + client.evaluate.return_value = JevSystemOneResponse( + model="jev-test", + answers={ + "tier": JevChoiceAnswer(type="choice", choice="SIMPLE", probabilities={"SIMPLE": 1.0}, confidence=1.0) + }, + ) + monkeypatch.setattr(proxy_server, "llm_router", _router()) + monkeypatch.setattr(auto_router_endpoints, "ComplexityRouter", partial(ComplexityRouter, jev_client=client)) + actor: Final = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-jev-budget-test", + user_id="admin", + models=["cheap-model"], + max_budget=max_budget, + spend=spend, + ) + request: Final = _request( + "what is 2+2", + classifier_type="jev", + jev_classifier_config={"model": "jev-test"}, + ) + + if denied: + with pytest.raises(ProxyException) as exc_info: + await preview_auto_router_routing(http_request=ROUTING_HTTP_REQUEST, data=request, user_api_key_dict=actor) + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + assert exc_info.value.code == "400" + assert exc_info.value.param is None + assert "Budget has been exceeded!" in exc_info.value.message + client.evaluate.assert_not_called() + return + + response: Final = await preview_auto_router_routing( + http_request=ROUTING_HTTP_REQUEST, data=request, user_api_key_dict=actor + ) + assert response.routed_model == "cheap-model" + assert response.routing_decision["cause"] == "jev_classifier" + assert response.routing_decision["classifier_model"] == "typesafe/jev-test" + client.evaluate.assert_awaited_once() + + +@pytest.mark.parametrize("max_budget, spend", ((0.0, 0.0), (1.0, 2.0))) +@pytest.mark.asyncio +async def test_a_heuristic_config_does_not_need_a_budget( + monkeypatch: pytest.MonkeyPatch, max_budget: float, spend: float +): import litellm.proxy.proxy_server as proxy_server monkeypatch.setattr(proxy_server, "llm_router", _router()) @@ -435,8 +504,8 @@ async def test_a_heuristic_config_does_not_need_a_budget(monkeypatch: pytest.Mon user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-broke", user_id="admin", - max_budget=1.0, - spend=2.0, + max_budget=max_budget, + spend=spend, models=["cheap-model"], ), ) @@ -851,7 +920,6 @@ class TestAutoRouterBenchmarks: # --------------------------------------------------------------------------- from datetime import datetime, timedelta, timezone -from unittest.mock import AsyncMock, MagicMock from litellm.proxy.management_endpoints.auto_router_endpoints import ( get_shadow_eval_job, From b9e5bb3abb0f2f0ddc06dcaf9edb63563cac2a2a Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Fri, 18 Sep 2026 22:03:37 +0000 Subject: [PATCH 02/47] test(proxy): allow JEV dependency in budget fixtures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/management_endpoints/test_auto_router_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 36130137c64..a5c93c41a84 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -459,7 +459,7 @@ async def test_jev_test_routing_enforces_key_budget_before_provider_invocation( user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-jev-budget-test", user_id="admin", - models=["cheap-model"], + models=["cheap-model", "typesafe/jev-test"], max_budget=max_budget, spend=spend, ) From 5f54f87d9887c13b75c42ffca5142ef2b7bcf2f1 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 16:40:10 +0000 Subject: [PATCH 03/47] feat(fal_ai): add Seedance video generation via fal queue API Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/videos/__init__.py | 3 + litellm/llms/fal_ai/videos/transformation.py | 512 ++++++++++++++++++ ...odel_prices_and_context_window_backup.json | 121 +++++ litellm/utils.py | 4 + model_prices_and_context_window.json | 121 +++++ .../test_fal_ai_video_transformation.py | 231 ++++++++ 6 files changed, 992 insertions(+) create mode 100644 litellm/llms/fal_ai/videos/__init__.py create mode 100644 litellm/llms/fal_ai/videos/transformation.py create mode 100644 tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py diff --git a/litellm/llms/fal_ai/videos/__init__.py b/litellm/llms/fal_ai/videos/__init__.py new file mode 100644 index 00000000000..c7e8f76c75b --- /dev/null +++ b/litellm/llms/fal_ai/videos/__init__.py @@ -0,0 +1,3 @@ +from litellm.llms.fal_ai.videos.transformation import FalAIVideoConfig + +__all__ = ("FalAIVideoConfig",) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py new file mode 100644 index 00000000000..f8ebf828d68 --- /dev/null +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -0,0 +1,512 @@ +import math +import time +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import httpx +from httpx._types import FileContent, RequestFiles +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.videos.transformation import BaseVideoConfig +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + _get_httpx_client, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # shared HTTP factory is private + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared HTTP factory lacks typed params +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.types.videos.main import ( + CharacterObject, + VideoCreateOptionalRequestParams, + VideoObject, +) +from litellm.types.videos.utils import ( + decode_video_id_with_provider, + encode_video_id_with_provider, +) + + +class FalAIVideoError(BaseLLMException): + pass + + +_ALLOWED_ASPECT_RATIOS: Final[frozenset[str]] = frozenset({"auto", "16:9", "9:16", "1:1", "4:3", "3:4", "21:9"}) +_ALLOWED_RESOLUTIONS: Final[frozenset[str]] = frozenset({"480p", "720p", "1080p", "4k"}) +_RESOLUTION_TIERS: Final[tuple[tuple[int, str], ...]] = ( + (480, "480p"), + (720, "720p"), + (1080, "1080p"), +) +_FAL_AI_PROVIDER: Final[str] = LlmProviders.FAL_AI.value + + +def _queue_request_base_path(model: str) -> str: + segments: Final[tuple[str, ...]] = tuple(model.split("/")) + segment_count: Final[int] = 3 if segments and segments[0] in frozenset(("workflows", "comfy")) else 2 + return "/".join(segments[:segment_count]) + + +def _duration_value(value: object) -> str | None: + if isinstance(value, str) and value == "auto": + return value + if isinstance(value, bool) or not isinstance(value, (int, float, str)): + return None + try: + return str(int(float(value))) + except (TypeError, ValueError): + return None + + +def _resolution_for_height(height: int) -> str: + return next((resolution for threshold, resolution in _RESOLUTION_TIERS if height <= threshold), "4k") + + +def _size_params(size: object) -> Mapping[str, str]: + if not isinstance(size, str): + return MappingProxyType({}) + if size in _ALLOWED_RESOLUTIONS: + return MappingProxyType({"resolution": size}) + if size.count("x") != 1: + return MappingProxyType({}) + width_text, height_text = size.split("x") + if not (width_text.isdigit() and height_text.isdigit()): + return MappingProxyType({}) + width: Final[int] = int(width_text) + height: Final[int] = int(height_text) + if width <= 0 or height <= 0: + return MappingProxyType({}) + reduced_gcd: Final[int] = math.gcd(width, height) + aspect_ratio: Final[str] = f"{width // reduced_gcd}:{height // reduced_gcd}" + resolution: Final[str] = _resolution_for_height(height) + if aspect_ratio in _ALLOWED_ASPECT_RATIOS: + return MappingProxyType({"resolution": resolution, "aspect_ratio": aspect_ratio}) + return MappingProxyType({"resolution": resolution}) + + +def _numeric_duration(value: object) -> float | None: + duration: Final[str | None] = _duration_value(value) + if duration is None or duration == "auto": + return None + return float(duration) + + +def _response_data(raw_response: httpx.Response) -> Mapping[str, object]: + return TypeAdapter(Mapping[str, object]).validate_python(raw_response.json()) + + +def _response_string(response_data: Mapping[str, object], key: str, default: str = "") -> str: + value: Final[object] = response_data.get(key) + return value if isinstance(value, str) else default + + +class FalAIVideoConfig(BaseVideoConfig): + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: API contract requires a list + return [ # mutable-ok: API contract requires a list + "model", + "prompt", + "input_reference", + "seconds", + "size", + "user", + "extra_headers", + ] + + def map_openai_params( + self, + video_create_optional_params: VideoCreateOptionalRequestParams, + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: BaseVideoConfig requires a mutable mapping + supported_params: Final[frozenset[str]] = frozenset(self.get_supported_openai_params(model)) + input_reference: Final[object] = video_create_optional_params.get("input_reference") + input_reference_params: Final[Mapping[str, str]] = ( + MappingProxyType({}) + if "input_reference" not in video_create_optional_params + else ( + MappingProxyType({"image_url": input_reference}) + if isinstance(input_reference, str) + else self._invalid_input_reference() + ) + ) + duration_params: Final[Mapping[str, str]] = ( + MappingProxyType({}) + if "seconds" not in video_create_optional_params + else self._duration_params(video_create_optional_params["seconds"]) + ) + size_params: Final[Mapping[str, str]] = ( + self._size_params(video_create_optional_params["size"]) + if "size" in video_create_optional_params + else MappingProxyType({}) + ) + user_params: Final[Mapping[str, str]] = ( + MappingProxyType({"end_user_id": user}) + if isinstance(user := video_create_optional_params.get("user"), str) + else MappingProxyType({}) + ) + return dict( # mutable-ok: BaseVideoConfig requires a mutable mapping + MappingProxyType( + { + **input_reference_params, + **duration_params, + **size_params, + **user_params, + **{ # mutable-ok: dynamic passthrough fields require a mapping + key: value for key, value in video_create_optional_params.items() if key not in supported_params + }, + } + ) + ) # mutable-ok: BaseVideoConfig requires a mutable mapping + + @staticmethod + def _invalid_input_reference() -> Mapping[str, str]: + raise ValueError("fal.ai needs a public image URL for input_reference") + + @staticmethod + def _duration_params(seconds: object) -> Mapping[str, str]: + duration: Final[str | None] = _duration_value(seconds) + if duration is None: + raise ValueError("fal.ai seconds must be a numeric value") + return MappingProxyType({"duration": duration}) + + @staticmethod + def _size_params(size: object) -> Mapping[str, str]: + return _size_params(size) + + def validate_environment( + self, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + model: str, + api_key: str | None = None, + litellm_params: GenericLiteLLMParams | None = None, + ) -> dict[str, str]: # mutable-ok: BaseVideoConfig requires mutable headers + final_api_key: Final[str | None] = ( + api_key + or (litellm_params.api_key if litellm_params is not None else None) + or get_secret_str("FAL_AI_API_KEY") + or get_secret_str("FAL_KEY") + ) + if not final_api_key: + raise ValueError("fal.ai API key is required") + return dict( # mutable-ok: BaseVideoConfig requires mutable headers + MappingProxyType( + { + **headers, + "Authorization": f"Key {final_api_key}", + "Content-Type": "application/json", + } + ) + ) # mutable-ok: BaseVideoConfig requires mutable headers + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict[str, object], # mutable-ok: BaseVideoConfig requires mutable parameters + ) -> str: + return (api_base or get_secret_str("FAL_AI_QUEUE_API_BASE") or "https://queue.fal.run").rstrip("/") + + def transform_video_create_request( + self, + model: str, + prompt: str, + api_base: str, + video_create_optional_request_params: dict[ # mutable-ok: BaseVideoConfig requires mutable parameters + str, object + ], # mutable-ok: BaseVideoConfig requires mutable parameters + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + ) -> tuple[dict[str, object], RequestFiles, str]: # mutable-ok: BaseVideoConfig requires mutable mappings + request_data: Final[dict[str, object]] = dict( # mutable-ok: HTTP JSON payload requires mutable data + MappingProxyType( + { + "prompt": prompt, + **{ # mutable-ok: dynamic request fields require a mapping + key: value for key, value in video_create_optional_request_params.items() if key != "model" + }, + } + ) + ) + return request_data, [], f"{api_base.rstrip('/')}/{model}" # mutable-ok: HTTP files payload requires a list + + def transform_video_create_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + request_data: Mapping[str, object] | None = None, + ) -> VideoObject: + response_data: Final[Mapping[str, object]] = _response_data(raw_response) + request_params: Final[Mapping[str, object]] = request_data or MappingProxyType({}) + request_id: Final[str] = _response_string(response_data, "request_id") + provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER + duration: Final[float | None] = _numeric_duration(request_params.get("duration")) + resolution: Final[object] = request_params.get("resolution") + seconds: Final[str | None] = _duration_value(request_params["duration"]) if duration is not None else None + size: Final[str | None] = resolution if isinstance(resolution, str) else None + usage: Final[dict[str, object]] = dict( # mutable-ok: VideoObject requires a mutable usage mapping + MappingProxyType( + { + key: value + for key, value in ( + ("duration_seconds", duration), + ("video_resolution", resolution if isinstance(resolution, str) else "720p"), + ) + if value is not None + } + ) + ) # mutable-ok: VideoObject requires a mutable usage mapping + video_object: Final[VideoObject] = VideoObject( + id=encode_video_id_with_provider(request_id, provider, model), + object="video", + status="queued", + created_at=int(time.time()), + model=model, + seconds=seconds, + size=size, + ) + video_object.usage = usage + return video_object + + def transform_video_status_retrieve_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + request_id, model_id = self._decode_video_id(video_id) + encoded_request_id: Final[str] = encode_url_path_segment(request_id, field_name="video_id") + return ( + f"{api_base.rstrip('/')}/{_queue_request_base_path(model_id)}/requests/{encoded_request_id}/status", + {}, # mutable-ok: BaseVideoConfig requires a mutable mapping + ) + + def transform_video_status_retrieve_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + ) -> VideoObject: + response_data: Final[Mapping[str, object]] = _response_data(raw_response) + raw_status: Final[str] = _response_string(response_data, "status", "IN_QUEUE") + status: Final[str] = MappingProxyType( + { + "IN_QUEUE": "queued", + "IN_PROGRESS": "in_progress", + "COMPLETED": "completed", + } + ).get(raw_status, "queued") + error_value: Final[object] = response_data.get("error") + error: Final[str | None] = error_value if isinstance(error_value, str) else None + provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER + return VideoObject( + id=encode_video_id_with_provider(_response_string(response_data, "request_id"), provider), + object="video", + status="failed" if error else status, + created_at=0, + error=( + {"code": "fal_error", "message": error} if error else None # mutable-ok: VideoObject requires a dict + ), # mutable-ok: VideoObject requires a dict + ) + + @staticmethod + def _decode_video_id(video_id: str) -> tuple[str, str]: + decoded: Final = decode_video_id_with_provider(video_id) + request_id: Final[str] = decoded.get("video_id", video_id) + model_id: Final[str | None] = decoded.get("model_id") + if not model_id: + raise ValueError("fal.ai video ids must be created through litellm with a model") + return request_id, model_id + + def transform_video_content_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + variant: str | None = None, + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseVideoConfig requires mutable mappings + request_id, model_id = self._decode_video_id(video_id) + encoded_request_id: Final[str] = encode_url_path_segment(request_id, field_name="video_id") + return ( + f"{api_base.rstrip('/')}/{_queue_request_base_path(model_id)}/requests/{encoded_request_id}", + {}, # mutable-ok: BaseVideoConfig requires a mutable mapping + ) + + @staticmethod + def _extract_video_url(response_data: Mapping[str, object]) -> str: + raw_video_data: Final[object] = response_data.get("video") + video_data: Final[Mapping[str, object] | None] = ( + TypeAdapter(Mapping[str, object]).validate_python(raw_video_data) + if isinstance(raw_video_data, Mapping) + else None + ) + if video_data is not None: + video_url: Final[object] = video_data.get("url") + if isinstance(video_url, str) and video_url: + return video_url + error_message: Final[str | None] = next( + (value for key in ("error", "detail") if isinstance(value := response_data.get(key), str)), + None, + ) + if error_message: + raise ValueError(f"fal.ai video result did not include a video URL: {error_message}") + raise ValueError("fal.ai video result did not include a video URL") + + def transform_video_content_response(self, raw_response: httpx.Response, logging_obj: object) -> bytes: + video_url: Final[str] = self._extract_video_url(_response_data(raw_response)) + httpx_client: Final[HTTPHandler] = _get_httpx_client() + video_response: Final[httpx.Response] = httpx_client.get( # pyright: ignore[reportUnknownMemberType] # HTTP handler stubs are untyped + video_url + ) + video_response.raise_for_status() + return video_response.content + + async def async_transform_video_content_response(self, raw_response: httpx.Response, logging_obj: object) -> bytes: + video_url: Final[str] = self._extract_video_url(_response_data(raw_response)) + async_httpx_client: Final[AsyncHTTPHandler] = get_async_httpx_client(llm_provider=LlmProviders.FAL_AI) + video_response: Final[httpx.Response] = await async_httpx_client.get( # pyright: ignore[reportUnknownMemberType] # HTTP handler stubs are untyped + video_url + ) + video_response.raise_for_status() + return video_response.content + + def transform_video_remix_request( + self, + video_id: str, + prompt: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + extra_body: Mapping[str, object] | None = None, + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video remix is not supported for fal.ai") + + def transform_video_remix_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + ) -> VideoObject: + raise NotImplementedError("video remix is not supported for fal.ai") + + def transform_video_list_request( + self, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + after: str | None = None, + limit: int | None = None, + order: str | None = None, + extra_query: Mapping[str, object] | None = None, + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video listing is not supported for fal.ai") + + def transform_video_list_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video listing is not supported for fal.ai") + + def transform_video_delete_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video delete is not supported for fal.ai") + + def transform_video_delete_response(self, raw_response: httpx.Response, logging_obj: object) -> VideoObject: + raise NotImplementedError("video delete is not supported for fal.ai") + + def transform_video_create_character_request( + self, + name: str, + video: object, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + ) -> tuple[str, list[object]]: # mutable-ok: BaseVideoConfig requires mutable lists + raise NotImplementedError("video character creation is not supported for fal.ai") + + def transform_video_create_character_response( + self, + raw_response: httpx.Response, + logging_obj: object, + ) -> CharacterObject: + raise NotImplementedError("video character creation is not supported for fal.ai") + + def transform_video_get_character_request( + self, + character_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video character retrieval is not supported for fal.ai") + + def transform_video_get_character_response( + self, + raw_response: httpx.Response, + logging_obj: object, + ) -> CharacterObject: + raise NotImplementedError("video character retrieval is not supported for fal.ai") + + def transform_video_edit_request( + self, + prompt: str, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + video_file: FileContent | None = None, + extra_body: Mapping[str, object] | None = None, + prefetched_source_data: Mapping[str, object] | None = None, + ) -> tuple[str, Mapping[str, object], RequestFiles | None]: + raise NotImplementedError("video edit is not supported for fal.ai") + + def transform_video_edit_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + request_data: Mapping[str, object] | None = None, + ) -> VideoObject: + raise NotImplementedError("video edit is not supported for fal.ai") + + def transform_video_extension_request( + self, + prompt: str, + video_id: str, + seconds: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + extra_body: Mapping[str, object] | None = None, + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + raise NotImplementedError("video extension is not supported for fal.ai") + + def transform_video_extension_response( + self, + raw_response: httpx.Response, + logging_obj: object, + custom_llm_provider: str | None = None, + ) -> VideoObject: + raise NotImplementedError("video extension is not supported for fal.ai") + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, str] | httpx.Headers, # mutable-ok: BaseLLMException requires mutable headers + ) -> BaseLLMException: + return FalAIVideoError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4dbf0337894..324eb6b2d66 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -22807,6 +22807,127 @@ "/v1/images/generations" ] }, + "fal_ai/bytedance/seedance-2.5/text-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/text-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.5/image-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/image-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.5/reference-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/reference-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/text-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/text-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/image-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/image-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/reference-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/reference-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, "fal_ai/fal-ai/ideogram/v3": { "litellm_provider": "fal_ai", "mode": "image_generation", diff --git a/litellm/utils.py b/litellm/utils.py index 48d13bc16af..3991cecdac6 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9403,6 +9403,10 @@ class ProviderConfigManager: from litellm.llms.runwayml.videos.transformation import RunwayMLVideoConfig return RunwayMLVideoConfig() + elif LlmProviders.FAL_AI == provider: + from litellm.llms.fal_ai.videos.transformation import FalAIVideoConfig + + return FalAIVideoConfig() elif LlmProviders.HOSTED_VLLM == provider: from litellm.llms.hosted_vllm.videos import get_hosted_vllm_video_config diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4dbf0337894..324eb6b2d66 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -22807,6 +22807,127 @@ "/v1/images/generations" ] }, + "fal_ai/bytedance/seedance-2.5/text-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/text-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.5/image-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/image-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.5/reference-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.473, + "output_cost_per_second_480p": 0.2205, + "output_cost_per_second_720p": 0.473, + "source": "https://fal.ai/models/bytedance/seedance-2.5/reference-to-video", + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/text-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/text-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/image-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/image-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, + "fal_ai/bytedance/seedance-2.0/reference-to-video": { + "litellm_provider": "fal_ai", + "mode": "video_generation", + "output_cost_per_second": 0.3034, + "output_cost_per_second_480p": 0.1346, + "output_cost_per_second_720p": 0.3034, + "output_cost_per_second_1080p": 0.682, + "output_cost_per_second_4k": 1.5552, + "source": "https://fal.ai/models/bytedance/seedance-2.0/reference-to-video", + "metadata": { + "comment": "fal bills $0.014 per 1k tokens (480p/720p/1080p) and $0.008 per 1k tokens (4k) with tokens = h*w*seconds*24/1024; 480p and 4k rates derived from that formula at 854x480 and 3840x2160" + }, + "supported_endpoints": [ + "/v1/videos" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ] + }, "fal_ai/fal-ai/ideogram/v3": { "litellm_provider": "fal_ai", "mode": "image_generation", diff --git a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py new file mode 100644 index 00000000000..d4b058ddcf6 --- /dev/null +++ b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py @@ -0,0 +1,231 @@ +from unittest.mock import Mock + +import httpx +import pytest + +import litellm +import litellm.llms.fal_ai.videos.transformation as fal_video_module +from litellm.cost_calculator import default_video_cost_calculator +from litellm.llms.fal_ai.videos.transformation import ( + FalAIVideoConfig, + FalAIVideoError, + _queue_request_base_path, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.types.videos.utils import decode_video_id_with_provider +from litellm.utils import ProviderConfigManager + +MODEL = "bytedance/seedance-2.5/text-to-video" + + +class TestFalAIVideoTransformation: + def setup_method(self): + self.config = FalAIVideoConfig() + self.logging_obj = Mock() + + def test_map_openai_params(self): + mapped = self.config.map_openai_params( + { + "seconds": "5", + "size": "1280x720", + "input_reference": "https://example.com/image.png", + "user": "user-123", + "generate_audio": False, + }, + MODEL, + False, + ) + + assert mapped == { + "duration": "5", + "resolution": "720p", + "aspect_ratio": "16:9", + "image_url": "https://example.com/image.png", + "end_user_id": "user-123", + "generate_audio": False, + } + + assert self.config.map_openai_params({"size": "1080x1080"}, MODEL, False) == { + "resolution": "1080p", + "aspect_ratio": "1:1", + } + assert self.config.map_openai_params({"size": "720p"}, MODEL, False) == {"resolution": "720p"} + + def test_map_openai_params_rejects_non_url_input_reference(self): + with pytest.raises(ValueError, match="public image URL"): + self.config.map_openai_params({"input_reference": b"image"}, MODEL, False) + + def test_transform_video_create_request(self): + body, files, url = self.config.transform_video_create_request( + model=MODEL, + prompt="A quiet ocean at sunrise", + api_base="https://queue.fal.run", + video_create_optional_request_params={ + "duration": "5", + "resolution": "480p", + "aspect_ratio": "16:9", + "generate_audio": False, + "model": MODEL, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert url == f"https://queue.fal.run/{MODEL}" + assert files == [] + assert body == { + "prompt": "A quiet ocean at sunrise", + "duration": "5", + "resolution": "480p", + "aspect_ratio": "16:9", + "generate_audio": False, + } + assert "model" not in body + + def test_transform_video_create_response_encodes_model_and_usage(self): + response = Mock(spec=httpx.Response) + response.json.return_value = {"request_id": "abc"} + + video = self.config.transform_video_create_response( + model=MODEL, + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + request_data={"duration": "5", "resolution": "480p"}, + ) + + decoded = decode_video_id_with_provider(video.id) + assert decoded["custom_llm_provider"] == "fal_ai" + assert decoded["model_id"] == MODEL + assert decoded["video_id"] == "abc" + assert video.status == "queued" + assert video.usage == {"duration_seconds": 5.0, "video_resolution": "480p"} + + auto_video = self.config.transform_video_create_response( + model=MODEL, + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + request_data={"duration": "auto"}, + ) + assert auto_video.usage == {"video_resolution": "720p"} + assert auto_video.seconds is None + assert auto_video.size is None + + def test_status_request_uses_queue_base_path(self): + response = Mock(spec=httpx.Response) + response.json.return_value = {"request_id": "abc"} + video = self.config.transform_video_create_response( + model=MODEL, + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + request_data={}, + ) + + url, params = self.config.transform_video_status_retrieve_request( + video_id=video.id, + api_base="https://queue.fal.run", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert url == "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status" + assert params == {} + assert _queue_request_base_path("workflows/owner/app/x") == "workflows/owner/app" + assert _queue_request_base_path("comfy/owner/app/x") == "comfy/owner/app" + + def test_status_request_rejects_unencoded_video_id(self): + with pytest.raises(ValueError, match="must be created through litellm"): + self.config.transform_video_status_retrieve_request( + video_id="abc", + api_base="https://queue.fal.run", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + @pytest.mark.parametrize( + ("response_data", "expected_status"), + [ + ({"request_id": "abc", "status": "IN_QUEUE"}, "queued"), + ({"request_id": "abc", "status": "IN_PROGRESS"}, "in_progress"), + ({"request_id": "abc", "status": "COMPLETED"}, "completed"), + ], + ) + def test_status_response_mapping(self, response_data, expected_status): + response = Mock(spec=httpx.Response) + response.json.return_value = response_data + + video = self.config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + ) + + assert video.status == expected_status + assert video.created_at == 0 + + def test_status_response_error(self): + response = Mock(spec=httpx.Response) + response.json.return_value = { + "request_id": "abc", + "status": "COMPLETED", + "error": "generation failed", + } + + video = self.config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + ) + + assert video.status == "failed" + assert video.error == {"code": "fal_error", "message": "generation failed"} + + def test_content_response_downloads_video_url(self, monkeypatch): + content_response = httpx.Response( + 200, + content=b"video-bytes", + request=httpx.Request("GET", "https://cdn.example.com/video.mp4"), + ) + + class FakeHTTPClient: + def get(self, url): + assert url == "https://cdn.example.com/video.mp4" + return content_response + + monkeypatch.setattr(fal_video_module, "_get_httpx_client", lambda: FakeHTTPClient()) + response = Mock(spec=httpx.Response) + response.json.return_value = {"video": {"url": "https://cdn.example.com/video.mp4"}} + + assert self.config.transform_video_content_response(response, self.logging_obj) == b"video-bytes" + + def test_content_response_rejects_missing_video(self): + response = Mock(spec=httpx.Response) + response.json.return_value = {"error": "generation failed"} + + with pytest.raises(ValueError, match="generation failed"): + self.config.transform_video_content_response(response, self.logging_obj) + + def test_provider_config_and_error_class(self): + provider_config = ProviderConfigManager.get_provider_video_config( + model=MODEL, + provider=LlmProviders.FAL_AI, + ) + assert isinstance(provider_config, FalAIVideoConfig) + assert isinstance(self.config.get_error_class("bad key", 401, {}), FalAIVideoError) + + def test_video_cost_uses_tiered_rows(self): + rows = { + model: row + for model, row in litellm.model_cost.items() + if row.get("litellm_provider") == "fal_ai" and row.get("mode") == "video_generation" + } + assert rows + for model, row in rows.items(): + assert default_video_cost_calculator(model, 5, "fal_ai", video_resolution="480p") == ( + 5 * row["output_cost_per_second_480p"] + ) + assert default_video_cost_calculator(model, 5, "fal_ai", video_resolution="720p") == ( + 5 * row["output_cost_per_second"] + ) From 0f4ce95492cde57b1f2eb29a1ea119dda0b13e24 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 16:45:53 +0000 Subject: [PATCH 04/47] refactor(fal_ai): simplify video config mappings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/videos/transformation.py | 178 +++++++++---------- 1 file changed, 80 insertions(+), 98 deletions(-) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index f8ebf828d68..f0142ddc3de 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -2,7 +2,7 @@ import math import time from collections.abc import Mapping from types import MappingProxyType -from typing import Final +from typing import Final, TypeAlias import httpx from httpx._types import FileContent, RequestFiles @@ -42,12 +42,25 @@ _RESOLUTION_TIERS: Final[tuple[tuple[int, str], ...]] = ( (720, "720p"), (1080, "1080p"), ) +_QUEUE_NAMESPACES: Final[frozenset[str]] = frozenset(("workflows", "comfy")) +_STATUS_MAP: Final[Mapping[str, str]] = MappingProxyType( + { + "IN_QUEUE": "queued", + "IN_PROGRESS": "in_progress", + "COMPLETED": "completed", + } +) _FAL_AI_PROVIDER: Final[str] = LlmProviders.FAL_AI.value +_SupportedParams: TypeAlias = list[str] +_VideoParams: TypeAlias = dict[str, object] +_VideoHeaders: TypeAlias = dict[str, str] +_VideoStringParams: TypeAlias = dict[str, str] +_VideoFiles: TypeAlias = list[object] def _queue_request_base_path(model: str) -> str: segments: Final[tuple[str, ...]] = tuple(model.split("/")) - segment_count: Final[int] = 3 if segments and segments[0] in frozenset(("workflows", "comfy")) else 2 + segment_count: Final[int] = 3 if segments and segments[0] in _QUEUE_NAMESPACES else 2 return "/".join(segments[:segment_count]) @@ -105,8 +118,8 @@ def _response_string(response_data: Mapping[str, object], key: str, default: str class FalAIVideoConfig(BaseVideoConfig): - def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: API contract requires a list - return [ # mutable-ok: API contract requires a list + def get_supported_openai_params(self, model: str) -> _SupportedParams: + supported_params: Final[_SupportedParams] = [ # mutable-ok: BaseVideoConfig requires a list "model", "prompt", "input_reference", @@ -115,23 +128,22 @@ class FalAIVideoConfig(BaseVideoConfig): "user", "extra_headers", ] + return supported_params def map_openai_params( self, video_create_optional_params: VideoCreateOptionalRequestParams, model: str, drop_params: bool, - ) -> dict[str, object]: # mutable-ok: BaseVideoConfig requires a mutable mapping + ) -> _VideoParams: supported_params: Final[frozenset[str]] = frozenset(self.get_supported_openai_params(model)) input_reference: Final[object] = video_create_optional_params.get("input_reference") + if "input_reference" in video_create_optional_params and not isinstance(input_reference, str): + raise ValueError("fal.ai needs a public image URL for input_reference") input_reference_params: Final[Mapping[str, str]] = ( MappingProxyType({}) - if "input_reference" not in video_create_optional_params - else ( - MappingProxyType({"image_url": input_reference}) - if isinstance(input_reference, str) - else self._invalid_input_reference() - ) + if not isinstance(input_reference, str) + else MappingProxyType({"image_url": input_reference}) ) duration_params: Final[Mapping[str, str]] = ( MappingProxyType({}) @@ -139,7 +151,7 @@ class FalAIVideoConfig(BaseVideoConfig): else self._duration_params(video_create_optional_params["seconds"]) ) size_params: Final[Mapping[str, str]] = ( - self._size_params(video_create_optional_params["size"]) + _size_params(video_create_optional_params["size"]) if "size" in video_create_optional_params else MappingProxyType({}) ) @@ -148,23 +160,16 @@ class FalAIVideoConfig(BaseVideoConfig): if isinstance(user := video_create_optional_params.get("user"), str) else MappingProxyType({}) ) - return dict( # mutable-ok: BaseVideoConfig requires a mutable mapping - MappingProxyType( - { - **input_reference_params, - **duration_params, - **size_params, - **user_params, - **{ # mutable-ok: dynamic passthrough fields require a mapping - key: value for key, value in video_create_optional_params.items() if key not in supported_params - }, - } - ) - ) # mutable-ok: BaseVideoConfig requires a mutable mapping - - @staticmethod - def _invalid_input_reference() -> Mapping[str, str]: - raise ValueError("fal.ai needs a public image URL for input_reference") + mapped_params: Final[_VideoParams] = { + **input_reference_params, + **duration_params, + **size_params, + **user_params, + **{ # mutable-ok: BaseVideoConfig requires a mutable parameter mapping + key: value for key, value in video_create_optional_params.items() if key not in supported_params + }, + } + return mapped_params @staticmethod def _duration_params(seconds: object) -> Mapping[str, str]: @@ -173,17 +178,13 @@ class FalAIVideoConfig(BaseVideoConfig): raise ValueError("fal.ai seconds must be a numeric value") return MappingProxyType({"duration": duration}) - @staticmethod - def _size_params(size: object) -> Mapping[str, str]: - return _size_params(size) - def validate_environment( self, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, model: str, api_key: str | None = None, litellm_params: GenericLiteLLMParams | None = None, - ) -> dict[str, str]: # mutable-ok: BaseVideoConfig requires mutable headers + ) -> _VideoHeaders: final_api_key: Final[str | None] = ( api_key or (litellm_params.api_key if litellm_params is not None else None) @@ -192,21 +193,18 @@ class FalAIVideoConfig(BaseVideoConfig): ) if not final_api_key: raise ValueError("fal.ai API key is required") - return dict( # mutable-ok: BaseVideoConfig requires mutable headers - MappingProxyType( - { - **headers, - "Authorization": f"Key {final_api_key}", - "Content-Type": "application/json", - } - ) - ) # mutable-ok: BaseVideoConfig requires mutable headers + validated_headers: Final[_VideoHeaders] = { + **headers, + "Authorization": f"Key {final_api_key}", + "Content-Type": "application/json", + } + return validated_headers def get_complete_url( self, model: str, api_base: str | None, - litellm_params: dict[str, object], # mutable-ok: BaseVideoConfig requires mutable parameters + litellm_params: _VideoParams, ) -> str: return (api_base or get_secret_str("FAL_AI_QUEUE_API_BASE") or "https://queue.fal.run").rstrip("/") @@ -215,22 +213,16 @@ class FalAIVideoConfig(BaseVideoConfig): model: str, prompt: str, api_base: str, - video_create_optional_request_params: dict[ # mutable-ok: BaseVideoConfig requires mutable parameters - str, object - ], # mutable-ok: BaseVideoConfig requires mutable parameters + video_create_optional_request_params: _VideoParams, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers - ) -> tuple[dict[str, object], RequestFiles, str]: # mutable-ok: BaseVideoConfig requires mutable mappings - request_data: Final[dict[str, object]] = dict( # mutable-ok: HTTP JSON payload requires mutable data - MappingProxyType( - { - "prompt": prompt, - **{ # mutable-ok: dynamic request fields require a mapping - key: value for key, value in video_create_optional_request_params.items() if key != "model" - }, - } - ) - ) + headers: _VideoHeaders, + ) -> tuple[_VideoParams, RequestFiles, str]: + request_data: Final[_VideoParams] = { + "prompt": prompt, + **{ # mutable-ok: HTTP JSON payload requires a mutable mapping + key: value for key, value in video_create_optional_request_params.items() if key != "model" + }, + } return request_data, [], f"{api_base.rstrip('/')}/{model}" # mutable-ok: HTTP files payload requires a list def transform_video_create_response( @@ -249,18 +241,14 @@ class FalAIVideoConfig(BaseVideoConfig): resolution: Final[object] = request_params.get("resolution") seconds: Final[str | None] = _duration_value(request_params["duration"]) if duration is not None else None size: Final[str | None] = resolution if isinstance(resolution, str) else None - usage: Final[dict[str, object]] = dict( # mutable-ok: VideoObject requires a mutable usage mapping - MappingProxyType( - { - key: value - for key, value in ( - ("duration_seconds", duration), - ("video_resolution", resolution if isinstance(resolution, str) else "720p"), - ) - if value is not None - } + usage: Final[_VideoParams] = { # mutable-ok: VideoObject requires a mutable usage mapping + key: value + for key, value in ( + ("duration_seconds", duration), + ("video_resolution", resolution if isinstance(resolution, str) else "720p"), ) - ) # mutable-ok: VideoObject requires a mutable usage mapping + if value is not None + } video_object: Final[VideoObject] = VideoObject( id=encode_video_id_with_provider(request_id, provider, model), object="video", @@ -278,8 +266,8 @@ class FalAIVideoConfig(BaseVideoConfig): video_id: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + headers: _VideoHeaders, + ) -> tuple[str, _VideoParams]: request_id, model_id = self._decode_video_id(video_id) encoded_request_id: Final[str] = encode_url_path_segment(request_id, field_name="video_id") return ( @@ -295,13 +283,7 @@ class FalAIVideoConfig(BaseVideoConfig): ) -> VideoObject: response_data: Final[Mapping[str, object]] = _response_data(raw_response) raw_status: Final[str] = _response_string(response_data, "status", "IN_QUEUE") - status: Final[str] = MappingProxyType( - { - "IN_QUEUE": "queued", - "IN_PROGRESS": "in_progress", - "COMPLETED": "completed", - } - ).get(raw_status, "queued") + status: Final[str] = _STATUS_MAP.get(raw_status, "queued") error_value: Final[object] = response_data.get("error") error: Final[str | None] = error_value if isinstance(error_value, str) else None provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER @@ -312,7 +294,7 @@ class FalAIVideoConfig(BaseVideoConfig): created_at=0, error=( {"code": "fal_error", "message": error} if error else None # mutable-ok: VideoObject requires a dict - ), # mutable-ok: VideoObject requires a dict + ), ) @staticmethod @@ -329,9 +311,9 @@ class FalAIVideoConfig(BaseVideoConfig): video_id: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, variant: str | None = None, - ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseVideoConfig requires mutable mappings + ) -> tuple[str, _VideoStringParams]: request_id, model_id = self._decode_video_id(video_id) encoded_request_id: Final[str] = encode_url_path_segment(request_id, field_name="video_id") return ( @@ -383,9 +365,9 @@ class FalAIVideoConfig(BaseVideoConfig): prompt: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, extra_body: Mapping[str, object] | None = None, - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + ) -> tuple[str, _VideoParams]: raise NotImplementedError("video remix is not supported for fal.ai") def transform_video_remix_response( @@ -400,12 +382,12 @@ class FalAIVideoConfig(BaseVideoConfig): self, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, after: str | None = None, limit: int | None = None, order: str | None = None, extra_query: Mapping[str, object] | None = None, - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + ) -> tuple[str, _VideoParams]: raise NotImplementedError("video listing is not supported for fal.ai") def transform_video_list_response( @@ -413,7 +395,7 @@ class FalAIVideoConfig(BaseVideoConfig): raw_response: httpx.Response, logging_obj: object, custom_llm_provider: str | None = None, - ) -> dict[str, str]: # mutable-ok: BaseVideoConfig requires mutable mappings + ) -> _VideoStringParams: raise NotImplementedError("video listing is not supported for fal.ai") def transform_video_delete_request( @@ -421,8 +403,8 @@ class FalAIVideoConfig(BaseVideoConfig): video_id: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + headers: _VideoHeaders, + ) -> tuple[str, _VideoParams]: raise NotImplementedError("video delete is not supported for fal.ai") def transform_video_delete_response(self, raw_response: httpx.Response, logging_obj: object) -> VideoObject: @@ -434,8 +416,8 @@ class FalAIVideoConfig(BaseVideoConfig): video: object, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers - ) -> tuple[str, list[object]]: # mutable-ok: BaseVideoConfig requires mutable lists + headers: _VideoHeaders, + ) -> tuple[str, _VideoFiles]: raise NotImplementedError("video character creation is not supported for fal.ai") def transform_video_create_character_response( @@ -450,8 +432,8 @@ class FalAIVideoConfig(BaseVideoConfig): character_id: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + headers: _VideoHeaders, + ) -> tuple[str, _VideoParams]: raise NotImplementedError("video character retrieval is not supported for fal.ai") def transform_video_get_character_response( @@ -467,7 +449,7 @@ class FalAIVideoConfig(BaseVideoConfig): video_id: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, video_file: FileContent | None = None, extra_body: Mapping[str, object] | None = None, prefetched_source_data: Mapping[str, object] | None = None, @@ -490,9 +472,9 @@ class FalAIVideoConfig(BaseVideoConfig): seconds: str, api_base: str, litellm_params: GenericLiteLLMParams, - headers: dict[str, str], # mutable-ok: BaseVideoConfig requires mutable headers + headers: _VideoHeaders, extra_body: Mapping[str, object] | None = None, - ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig requires mutable mappings + ) -> tuple[str, _VideoParams]: raise NotImplementedError("video extension is not supported for fal.ai") def transform_video_extension_response( @@ -507,6 +489,6 @@ class FalAIVideoConfig(BaseVideoConfig): self, error_message: str, status_code: int, - headers: dict[str, str] | httpx.Headers, # mutable-ok: BaseLLMException requires mutable headers + headers: _VideoHeaders | httpx.Headers, ) -> BaseLLMException: return FalAIVideoError(status_code=status_code, message=error_message, headers=headers) From 141548dcf3ec7588e86c2e304fca33be6c4cdc7e Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 16:54:11 +0000 Subject: [PATCH 05/47] fix(fal_ai): keep status ids pollable and size resolution by the short side Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/videos/transformation.py | 17 +++++++-- .../test_fal_ai_video_transformation.py | 37 +++++++++++++++++++ 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index f0142ddc3de..4ca1f918c90 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -75,8 +75,16 @@ def _duration_value(value: object) -> str | None: return None -def _resolution_for_height(height: int) -> str: - return next((resolution for threshold, resolution in _RESOLUTION_TIERS if height <= threshold), "4k") +def _resolution_for_short_side(short_side: int) -> str: + return next((resolution for threshold, resolution in _RESOLUTION_TIERS if short_side <= threshold), "4k") + + +def _model_path_from_queue_url(url: object) -> str | None: + if not isinstance(url, str) or not url: + return None + path: Final[str] = httpx.URL(url).path.strip("/") + model_path, separator, _ = path.partition("/requests/") + return model_path if separator and model_path else None def _size_params(size: object) -> Mapping[str, str]: @@ -95,7 +103,7 @@ def _size_params(size: object) -> Mapping[str, str]: return MappingProxyType({}) reduced_gcd: Final[int] = math.gcd(width, height) aspect_ratio: Final[str] = f"{width // reduced_gcd}:{height // reduced_gcd}" - resolution: Final[str] = _resolution_for_height(height) + resolution: Final[str] = _resolution_for_short_side(min(width, height)) if aspect_ratio in _ALLOWED_ASPECT_RATIOS: return MappingProxyType({"resolution": resolution, "aspect_ratio": aspect_ratio}) return MappingProxyType({"resolution": resolution}) @@ -287,8 +295,9 @@ class FalAIVideoConfig(BaseVideoConfig): error_value: Final[object] = response_data.get("error") error: Final[str | None] = error_value if isinstance(error_value, str) else None provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER + model_path: Final[str | None] = _model_path_from_queue_url(response_data.get("response_url")) return VideoObject( - id=encode_video_id_with_provider(_response_string(response_data, "request_id"), provider), + id=encode_video_id_with_provider(_response_string(response_data, "request_id"), provider, model_path), object="video", status="failed" if error else status, created_at=0, diff --git a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py index d4b058ddcf6..f367fa331d5 100644 --- a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py +++ b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py @@ -51,6 +51,14 @@ class TestFalAIVideoTransformation: "aspect_ratio": "1:1", } assert self.config.map_openai_params({"size": "720p"}, MODEL, False) == {"resolution": "720p"} + assert self.config.map_openai_params({"size": "720x1280"}, MODEL, False) == { + "resolution": "720p", + "aspect_ratio": "9:16", + } + assert self.config.map_openai_params({"size": "1080x1920"}, MODEL, False) == { + "resolution": "1080p", + "aspect_ratio": "9:16", + } def test_map_openai_params_rejects_non_url_input_reference(self): with pytest.raises(ValueError, match="public image URL"): @@ -165,6 +173,35 @@ class TestFalAIVideoTransformation: assert video.status == expected_status assert video.created_at == 0 + def test_status_response_id_stays_pollable(self): + response = Mock(spec=httpx.Response) + response.json.return_value = { + "request_id": "abc", + "status": "IN_PROGRESS", + "response_url": "https://queue.fal.run/bytedance/seedance-2.5/requests/abc", + } + + video = self.config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + ) + + status_url, _ = self.config.transform_video_status_retrieve_request( + video_id=video.id, + api_base="https://queue.fal.run", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + content_url, _ = self.config.transform_video_content_request( + video_id=video.id, + api_base="https://queue.fal.run", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert status_url == "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status" + assert content_url == "https://queue.fal.run/bytedance/seedance-2.5/requests/abc" + def test_status_response_error(self): response = Mock(spec=httpx.Response) response.json.return_value = { From c359ef763eae44523cace811bcb1999992250bde Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 17:01:09 +0000 Subject: [PATCH 06/47] fix(fal_ai): keep model in polled video ids and pick resolution from the short side Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/videos/transformation.py | 28 ++++++-- .../test_fal_ai_video_transformation.py | 64 +++++++++++-------- 2 files changed, 58 insertions(+), 34 deletions(-) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index 4ca1f918c90..766316a5b18 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -79,12 +79,22 @@ def _resolution_for_short_side(short_side: int) -> str: return next((resolution for threshold, resolution in _RESOLUTION_TIERS if short_side <= threshold), "4k") -def _model_path_from_queue_url(url: object) -> str | None: - if not isinstance(url, str) or not url: +def _model_path_from_request_url(raw_response: httpx.Response) -> str | None: + segments: Final[tuple[str, ...]] = tuple(segment for segment in raw_response.request.url.path.split("/") if segment) + if "requests" not in segments: return None - path: Final[str] = httpx.URL(url).path.strip("/") - model_path, separator, _ = path.partition("/requests/") - return model_path if separator and model_path else None + model_segments: Final[tuple[str, ...]] = segments[: segments.index("requests")] + segment_count: Final[int] = 3 if len(model_segments) >= 3 and model_segments[-3] in _QUEUE_NAMESPACES else 2 + return "/".join(model_segments[-segment_count:]) if len(model_segments) >= segment_count else None + + +def _request_id_from_request_url(raw_response: httpx.Response) -> str | None: + segments: Final[tuple[str, ...]] = tuple(segment for segment in raw_response.request.url.path.split("/") if segment) + if "requests" not in segments: + return None + request_index: Final[int] = segments.index("requests") + request_id_index: Final[int] = request_index + 1 + return segments[request_id_index] if len(segments) > request_id_index else None def _size_params(size: object) -> Mapping[str, str]: @@ -295,12 +305,16 @@ class FalAIVideoConfig(BaseVideoConfig): error_value: Final[object] = response_data.get("error") error: Final[str | None] = error_value if isinstance(error_value, str) else None provider: Final[str] = custom_llm_provider or _FAL_AI_PROVIDER - model_path: Final[str | None] = _model_path_from_queue_url(response_data.get("response_url")) + model_path: Final[str | None] = _model_path_from_request_url(raw_response) + request_id: Final[str] = _response_string(response_data, "request_id") or ( + _request_id_from_request_url(raw_response) or "" + ) return VideoObject( - id=encode_video_id_with_provider(_response_string(response_data, "request_id"), provider, model_path), + id=encode_video_id_with_provider(request_id, provider, model_path), object="video", status="failed" if error else status, created_at=0, + model=model_path, error=( {"code": "fal_error", "message": error} if error else None # mutable-ok: VideoObject requires a dict ), diff --git a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py index f367fa331d5..8e0c68e30bb 100644 --- a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py +++ b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py @@ -161,8 +161,8 @@ class TestFalAIVideoTransformation: ], ) def test_status_response_mapping(self, response_data, expected_status): - response = Mock(spec=httpx.Response) - response.json.return_value = response_data + status_url = "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status" + response = httpx.Response(200, json=response_data, request=httpx.Request("GET", status_url)) video = self.config.transform_video_status_retrieve_response( raw_response=response, @@ -172,43 +172,32 @@ class TestFalAIVideoTransformation: assert video.status == expected_status assert video.created_at == 0 + decoded = decode_video_id_with_provider(video.id) + assert decoded["model_id"] == "bytedance/seedance-2.5" + assert decoded["video_id"] == "abc" - def test_status_response_id_stays_pollable(self): - response = Mock(spec=httpx.Response) - response.json.return_value = { - "request_id": "abc", - "status": "IN_PROGRESS", - "response_url": "https://queue.fal.run/bytedance/seedance-2.5/requests/abc", - } - - video = self.config.transform_video_status_retrieve_response( - raw_response=response, - logging_obj=self.logging_obj, - custom_llm_provider="fal_ai", - ) - - status_url, _ = self.config.transform_video_status_retrieve_request( + poll_url, _ = self.config.transform_video_status_retrieve_request( video_id=video.id, api_base="https://queue.fal.run", litellm_params=GenericLiteLLMParams(), headers={}, ) - content_url, _ = self.config.transform_video_content_request( - video_id=video.id, - api_base="https://queue.fal.run", - litellm_params=GenericLiteLLMParams(), - headers={}, - ) - assert status_url == "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status" - assert content_url == "https://queue.fal.run/bytedance/seedance-2.5/requests/abc" + assert poll_url == status_url def test_status_response_error(self): - response = Mock(spec=httpx.Response) - response.json.return_value = { + response_data = { "request_id": "abc", "status": "COMPLETED", "error": "generation failed", } + response = httpx.Response( + 200, + json=response_data, + request=httpx.Request( + "GET", + "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status", + ), + ) video = self.config.transform_video_status_retrieve_response( raw_response=response, @@ -219,6 +208,27 @@ class TestFalAIVideoTransformation: assert video.status == "failed" assert video.error == {"code": "fal_error", "message": "generation failed"} + def test_status_response_uses_namespaced_request_url(self): + response = httpx.Response( + 200, + json={"status": "IN_PROGRESS"}, + request=httpx.Request( + "GET", + "https://example.com/proxy/workflows/owner/app/requests/xyz/status", + ), + ) + + video = self.config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=self.logging_obj, + custom_llm_provider="fal_ai", + ) + + decoded = decode_video_id_with_provider(video.id) + assert decoded["model_id"] == "workflows/owner/app" + assert decoded["video_id"] == "xyz" + assert video.model == "workflows/owner/app" + def test_content_response_downloads_video_url(self, monkeypatch): content_response = httpx.Response( 200, From aa5f0858f75b3e074264e0266f87b70a6cb70391 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 17:13:41 +0000 Subject: [PATCH 07/47] test(pricing): allow video endpoint and rates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index b40c10de428..d694c08510a 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -940,6 +940,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "/v1/audio/transcriptions", "/v1/audio/speech", "/v1/ocr", + "/v1/videos", "/vertex_ai/live", "/v1/listen", "/v1beta/interactions", @@ -1069,6 +1070,9 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): # Add any model IDs that should be exempt from the cost validation # Example: "expensive-model-id", "runwayml/seedance2", # 4K output is 150 credits/second = $1.50/second + "fal_ai/bytedance/seedance-2.0/text-to-video", + "fal_ai/bytedance/seedance-2.0/image-to-video", + "fal_ai/bytedance/seedance-2.0/reference-to-video", ] is_valid, violations = validate_model_cost_values(actual_json, exceptions) From e0b455e94e83dfedad0364aedc6b4cdfcb59acb0 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 17:27:44 +0000 Subject: [PATCH 08/47] fix(fal_ai): read only the documented FAL_AI_API_KEY env var Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fal_ai/videos/transformation.py | 5 ++--- .../test_fal_ai_video_transformation.py | 20 +++++++++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index 766316a5b18..98528c82f6b 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -207,10 +207,9 @@ class FalAIVideoConfig(BaseVideoConfig): api_key or (litellm_params.api_key if litellm_params is not None else None) or get_secret_str("FAL_AI_API_KEY") - or get_secret_str("FAL_KEY") ) if not final_api_key: - raise ValueError("fal.ai API key is required") + raise ValueError("FAL_AI_API_KEY is not set") validated_headers: Final[_VideoHeaders] = { **headers, "Authorization": f"Key {final_api_key}", @@ -224,7 +223,7 @@ class FalAIVideoConfig(BaseVideoConfig): api_base: str | None, litellm_params: _VideoParams, ) -> str: - return (api_base or get_secret_str("FAL_AI_QUEUE_API_BASE") or "https://queue.fal.run").rstrip("/") + return (api_base or "https://queue.fal.run").rstrip("/") def transform_video_create_request( self, diff --git a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py index 8e0c68e30bb..5e2e4532265 100644 --- a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py +++ b/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py @@ -91,6 +91,26 @@ class TestFalAIVideoTransformation: } assert "model" not in body + def test_get_complete_url_respects_api_base_override(self): + url = self.config.get_complete_url( + model=MODEL, + api_base="https://proxy.internal/", + litellm_params={}, + ) + + assert url == "https://proxy.internal" + + def test_validate_environment_requires_fal_ai_api_key(self, monkeypatch): + monkeypatch.setattr(fal_video_module, "get_secret_str", lambda _: None) + + with pytest.raises(ValueError, match="FAL_AI_API_KEY is not set"): + self.config.validate_environment( + headers={}, + model=MODEL, + api_key=None, + litellm_params=GenericLiteLLMParams(), + ) + def test_transform_video_create_response_encodes_model_and_usage(self): response = Mock(spec=httpx.Response) response.json.return_value = {"request_id": "abc"} From 3b0d32ec6f14979687d4bb76f51db5f9427a131a Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Sat, 19 Sep 2026 18:11:50 +0000 Subject: [PATCH 09/47] fix(proxy): reject throttled exhausted budgets in JEV previews Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../auto_router_endpoints.py | 8 +++ .../test_auto_router_endpoints.py | 49 +++++++++++++++++++ 2 files changed, 57 insertions(+) diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index f79425d2e97..c5a10cf5c80 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -345,6 +345,14 @@ async def _authorize_models_this_test_can_call( code=status.HTTP_400_BAD_REQUEST, ) from e + if config.classifier_type == "jev" and user_api_key_dict.budget_throttle_pct is not None: + raise ProxyException( + message="Budget has been exceeded! JEV Test Routing requires available budget.", + type=ProxyErrorTypes.budget_exceeded, + param=None, + code=status.HTTP_400_BAD_REQUEST, + ) + @router.post( "/auto_router/validate_complexity_router_config", diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index a5c93c41a84..eb9076a9d2d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -12,6 +12,7 @@ import pytest from fastapi import HTTPException, Request from pydantic import ValidationError +import litellm from litellm.proxy import proxy_server from litellm.proxy._types import ( LitellmUserRoles, @@ -488,6 +489,54 @@ async def test_jev_test_routing_enforces_key_budget_before_provider_invocation( client.evaluate.assert_awaited_once() +@pytest.mark.parametrize( + "max_budget, spend, denied", + ((0.0, 0.0, True), (1.0, 2.0, True), (1.0, 0.5, False), (None, 2.0, False)), +) +@pytest.mark.asyncio +async def test_jev_test_routing_hard_blocks_exhausted_throttle_enabled_keys( + monkeypatch: pytest.MonkeyPatch, max_budget: float | None, spend: float, denied: bool +) -> None: + client: Final = AsyncMock(spec=JevClassifierClient) + client.evaluate.return_value = JevSystemOneResponse( + model="jev-test", + answers={ + "tier": JevChoiceAnswer(type="choice", choice="SIMPLE", probabilities={"SIMPLE": 1.0}, confidence=1.0) + }, + ) + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1) + monkeypatch.setattr(proxy_server, "llm_router", _router()) + monkeypatch.setattr(auto_router_endpoints, "ComplexityRouter", partial(ComplexityRouter, jev_client=client)) + actor: Final = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-jev-throttle-test", + user_id="admin", + models=["cheap-model", "typesafe/jev-test"], + max_budget=max_budget, + spend=spend, + rpm_limit=100, + metadata={"throttle_on_budget_exceeded": True}, + ) + request: Final = _request( + "what is 2+2", + classifier_type="jev", + jev_classifier_config={"model": "jev-test"}, + ) + if denied: + with pytest.raises(ProxyException) as exc_info: + await preview_auto_router_routing(http_request=ROUTING_HTTP_REQUEST, data=request, user_api_key_dict=actor) + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + assert exc_info.value.code == "400" + client.evaluate.assert_not_called() + return + + response: Final = await preview_auto_router_routing( + http_request=ROUTING_HTTP_REQUEST, data=request, user_api_key_dict=actor + ) + assert response.routing_decision["cause"] == "jev_classifier" + client.evaluate.assert_awaited_once() + + @pytest.mark.parametrize("max_budget, spend", ((0.0, 0.0), (1.0, 2.0))) @pytest.mark.asyncio async def test_a_heuristic_config_does_not_need_a_budget( From 85a6a8e2063ef0932dc40b19c97eb0120a0c4235 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 18:22:58 +0000 Subject: [PATCH 10/47] test(e2e): cover fal Seedance video create, poll and download Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_nonconversational.yaml | 1 + tests/e2e/coverage_registry/schema.py | 2 + .../LLM_TRANSLATION_COVERAGE_MATRIX.md | 2 + tests/e2e/llm_translation/endpoints_client.py | 36 +++++++++- .../test_video_generation_e2e.py | 69 +++++++++++++++++++ 5 files changed, 109 insertions(+), 1 deletion(-) create mode 100644 tests/e2e/llm_translation/test_video_generation_e2e.py diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 50f9b9808b2..6970567b6f0 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -80,6 +80,7 @@ - {id: llm.images_generations.vertex.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "vertex_ai/image_generation/image_generation_handler.py", rationale: "Vertex Imagen"} - {id: llm.images_generations.bedrock.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "bedrock/image_generation/image_handler.py", rationale: "Bedrock Titan Image"} - {id: llm.images_generations.black_forest_labs.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "black_forest_labs/image_generation/handler.py", rationale: "BFL Flux via OpenAI-compat"} +- {id: llm.videos.fal_ai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: videos, route: fal_ai, capability: basic, streaming: nonstream, assertions: [works], source: "test_video_generation_e2e.py", rationale: "fal queue video create, poll, content download"} - {id: llm.audio_speech.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_audio_speech_e2e.py:22", rationale: "OpenAI TTS binary audio"} - {id: llm.audio_speech.openai.basic.stream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: basic, streaming: stream, assertions: [works], source: "proxy_server.py:9043", rationale: "TTS streaming chunk generator"} - {id: llm.audio_speech.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.6 / LIT-4778", rationale: "TTS missing input/model, invalid voice, empty input rejected"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index fa6dad90126..3ae17432863 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -44,6 +44,7 @@ LlmEndpoint = Literal[ "vector_stores", "ocr", "bedrock_native", + "videos", ] LlmRoute = Literal[ @@ -53,6 +54,7 @@ LlmRoute = Literal[ "bedrock_converse", "bedrock_invoke", "cohere", + "fal_ai", "gemini", "hosted_vllm", "openai", diff --git a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md index 44d6e79122e..178af054f2b 100644 --- a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md +++ b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md @@ -48,6 +48,7 @@ most likely to silently break and the one a mock can't prove works. |----------|---------------|-----------|------------|-------------|--------| | Chat | live (spend suite) | live (spend suite) | gap | live | partial | | Embeddings | live (spend suite) | n/a | n/a | live | covered | +| Video | live (fal.ai Seedance) | n/a | n/a | - | partial | | Responses / image / audio / rerank / realtime | - | - | - | - | gap | ## This suite's files @@ -61,6 +62,7 @@ most likely to silently break and the one a mock can't prove works. | `test_anthropic_passthrough_streaming_logs_cost` | anthropic native, stream, cost | | `test_anthropic_passthrough_tool_call_logs_cost` | anthropic native, tool call, cost | | `test_vertex_passthrough_via_managed_model_logs_cost` | vertex_ai native, non-stream, cost | +| `test_fal_seedance_video_completes_and_downloads` | fal.ai Seedance video create, poll, and content download | Vertex keeps the credential on the proxy like gemini/anthropic, but the deployment is added at runtime instead of declared in the gateway config: the test POSTs `/model/new` diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py index 4d2c73e7078..165a83e76c0 100644 --- a/tests/e2e/llm_translation/endpoints_client.py +++ b/tests/e2e/llm_translation/endpoints_client.py @@ -13,7 +13,7 @@ from dataclasses import dataclass from typing import Literal from e2e_config import SLOW_PROVIDER_TIMEOUT_SECONDS -from e2e_http import BinaryStream, Result, StreamingResponse +from e2e_http import BinaryStream, NoBody, Result, StreamingResponse from models import CacheControl, ChatMessage, LiteLLMParamsBody, RichMessage, TextBlock from proxy_client import ProxyClient from pydantic import BaseModel @@ -26,6 +26,8 @@ __all__ = [ "TextBlock", "TranscriptionForm", "TranscriptionResult", + "VideoObject", + "VideoRequest", ] @@ -127,6 +129,13 @@ class ImageRequest(BaseModel): size: str = "1024x1024" +class VideoRequest(BaseModel): + model: str + prompt: str + seconds: str = "4" + size: str = "1280x720" + + class ImageEditForm(BaseModel): model: str prompt: str @@ -267,6 +276,12 @@ class ImagesResult(BaseModel): data: list[ImageItem] = [] +class VideoObject(BaseModel): + id: str + status: str + model: str | None = None + + class TranscriptionResult(BaseModel): text: str = "" @@ -440,6 +455,25 @@ class EndpointsClient: "/v1/images/generations", key, ImageRequest(model=model, prompt=prompt) ) + def videos(self, key: str, model: str, prompt: str) -> StreamingResponse: + return self._send( + "/v1/videos", key, VideoRequest(model=model, prompt=prompt) + ) + + def video_status(self, key: str, video_id: str) -> Result[VideoObject]: + return self.proxy.transport.get( + f"/v1/videos/{video_id}", + headers=self.proxy.transport.bearer(key), + params=NoBody(), + response_type=VideoObject, + ) + + def video_content(self, key: str, video_id: str) -> StreamingResponse: + return self.proxy.transport.download( + f"/v1/videos/{video_id}/content", + headers=self.proxy.transport.bearer(key), + ) + def image_edit( self, key: str, model: str, prompt: str, image: bytes, *, filename: str = "image.png" ) -> Result[ImagesResult]: diff --git a/tests/e2e/llm_translation/test_video_generation_e2e.py b/tests/e2e/llm_translation/test_video_generation_e2e.py new file mode 100644 index 00000000000..b65529aa260 --- /dev/null +++ b/tests/e2e/llm_translation/test_video_generation_e2e.py @@ -0,0 +1,69 @@ +"""Live e2e: POST /v1/videos creates a video and serves its content. + +Registers a fal.ai Seedance deployment at runtime, polls the queued video until it +completes, and asserts the generated content is returned as binary data. +""" + +from __future__ import annotations + +import time +from typing import Final + +import pytest +from e2e_config import unique_marker +from e2e_http import require_successful_call, unwrap +from endpoints_client import EndpointsClient, VideoObject +from lifecycle import ResourceManager +from models import LiteLLMParamsBody + +pytestmark = pytest.mark.e2e + +_POLL_INTERVAL_SECONDS: Final[float] = 5.0 +_POLL_TIMEOUT_SECONDS: Final[float] = 600.0 + + +def _wait_for_completion( + endpoints_client: EndpointsClient, key: str, created: VideoObject +) -> VideoObject: + deadline = time.monotonic() + _POLL_TIMEOUT_SECONDS + while time.monotonic() < deadline: + status = unwrap(endpoints_client.video_status(key, created.id)) + assert status.id == created.id + if status.status == "completed": + return status + if status.status == "failed": + pytest.fail(f"fal.ai video generation failed: {status}") + time.sleep(_POLL_INTERVAL_SECONDS) + pytest.fail(f"fal.ai video {created.id!r} did not complete within {_POLL_TIMEOUT_SECONDS}s") + + +class TestVideoGeneration: + @pytest.mark.covers("llm.videos.fal_ai.basic.nonstream.works") + def test_fal_seedance_video_completes_and_downloads( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-fal-video-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="fal_ai/bytedance/seedance-2.5/text-to-video", + api_key="os.environ/FAL_AI_API_KEY", + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.videos( + key, model, "a red fox running through snow at dawn" + ) + require_successful_call(result) + created = VideoObject.model_validate_json(result.body) + assert created.id + assert created.model + + _wait_for_completion(endpoints_client, key, created) + + content = endpoints_client.video_content(key, created.id) + require_successful_call(content) + assert len(content.body) > 0 + assert not (content.content_type or "").startswith("application/json") From 9a63e06c63e60b3b30126aa7461115a2ccc041fe Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 23:28:51 +0000 Subject: [PATCH 11/47] test(integration): cover fal Seedance video queue wire contract Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_nonconversational.yaml | 1 - tests/e2e/coverage_registry/schema.py | 2 - .../LLM_TRANSLATION_COVERAGE_MATRIX.md | 2 - tests/e2e/llm_translation/endpoints_client.py | 36 +--------- .../test_video_generation_e2e.py | 69 ------------------ tests/integration/contracts.json | 3 + .../providers/test_fal_ai_video_wire.py | 72 +++++++++++++++++++ 7 files changed, 76 insertions(+), 109 deletions(-) delete mode 100644 tests/e2e/llm_translation/test_video_generation_e2e.py create mode 100644 tests/integration/providers/test_fal_ai_video_wire.py diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 6970567b6f0..50f9b9808b2 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -80,7 +80,6 @@ - {id: llm.images_generations.vertex.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "vertex_ai/image_generation/image_generation_handler.py", rationale: "Vertex Imagen"} - {id: llm.images_generations.bedrock.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "bedrock/image_generation/image_handler.py", rationale: "Bedrock Titan Image"} - {id: llm.images_generations.black_forest_labs.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: images_generations, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "black_forest_labs/image_generation/handler.py", rationale: "BFL Flux via OpenAI-compat"} -- {id: llm.videos.fal_ai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: videos, route: fal_ai, capability: basic, streaming: nonstream, assertions: [works], source: "test_video_generation_e2e.py", rationale: "fal queue video create, poll, content download"} - {id: llm.audio_speech.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_audio_speech_e2e.py:22", rationale: "OpenAI TTS binary audio"} - {id: llm.audio_speech.openai.basic.stream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: basic, streaming: stream, assertions: [works], source: "proxy_server.py:9043", rationale: "TTS streaming chunk generator"} - {id: llm.audio_speech.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: audio_speech, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.6 / LIT-4778", rationale: "TTS missing input/model, invalid voice, empty input rejected"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index 3ae17432863..fa6dad90126 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -44,7 +44,6 @@ LlmEndpoint = Literal[ "vector_stores", "ocr", "bedrock_native", - "videos", ] LlmRoute = Literal[ @@ -54,7 +53,6 @@ LlmRoute = Literal[ "bedrock_converse", "bedrock_invoke", "cohere", - "fal_ai", "gemini", "hosted_vllm", "openai", diff --git a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md index 178af054f2b..44d6e79122e 100644 --- a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md +++ b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md @@ -48,7 +48,6 @@ most likely to silently break and the one a mock can't prove works. |----------|---------------|-----------|------------|-------------|--------| | Chat | live (spend suite) | live (spend suite) | gap | live | partial | | Embeddings | live (spend suite) | n/a | n/a | live | covered | -| Video | live (fal.ai Seedance) | n/a | n/a | - | partial | | Responses / image / audio / rerank / realtime | - | - | - | - | gap | ## This suite's files @@ -62,7 +61,6 @@ most likely to silently break and the one a mock can't prove works. | `test_anthropic_passthrough_streaming_logs_cost` | anthropic native, stream, cost | | `test_anthropic_passthrough_tool_call_logs_cost` | anthropic native, tool call, cost | | `test_vertex_passthrough_via_managed_model_logs_cost` | vertex_ai native, non-stream, cost | -| `test_fal_seedance_video_completes_and_downloads` | fal.ai Seedance video create, poll, and content download | Vertex keeps the credential on the proxy like gemini/anthropic, but the deployment is added at runtime instead of declared in the gateway config: the test POSTs `/model/new` diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py index 165a83e76c0..4d2c73e7078 100644 --- a/tests/e2e/llm_translation/endpoints_client.py +++ b/tests/e2e/llm_translation/endpoints_client.py @@ -13,7 +13,7 @@ from dataclasses import dataclass from typing import Literal from e2e_config import SLOW_PROVIDER_TIMEOUT_SECONDS -from e2e_http import BinaryStream, NoBody, Result, StreamingResponse +from e2e_http import BinaryStream, Result, StreamingResponse from models import CacheControl, ChatMessage, LiteLLMParamsBody, RichMessage, TextBlock from proxy_client import ProxyClient from pydantic import BaseModel @@ -26,8 +26,6 @@ __all__ = [ "TextBlock", "TranscriptionForm", "TranscriptionResult", - "VideoObject", - "VideoRequest", ] @@ -129,13 +127,6 @@ class ImageRequest(BaseModel): size: str = "1024x1024" -class VideoRequest(BaseModel): - model: str - prompt: str - seconds: str = "4" - size: str = "1280x720" - - class ImageEditForm(BaseModel): model: str prompt: str @@ -276,12 +267,6 @@ class ImagesResult(BaseModel): data: list[ImageItem] = [] -class VideoObject(BaseModel): - id: str - status: str - model: str | None = None - - class TranscriptionResult(BaseModel): text: str = "" @@ -455,25 +440,6 @@ class EndpointsClient: "/v1/images/generations", key, ImageRequest(model=model, prompt=prompt) ) - def videos(self, key: str, model: str, prompt: str) -> StreamingResponse: - return self._send( - "/v1/videos", key, VideoRequest(model=model, prompt=prompt) - ) - - def video_status(self, key: str, video_id: str) -> Result[VideoObject]: - return self.proxy.transport.get( - f"/v1/videos/{video_id}", - headers=self.proxy.transport.bearer(key), - params=NoBody(), - response_type=VideoObject, - ) - - def video_content(self, key: str, video_id: str) -> StreamingResponse: - return self.proxy.transport.download( - f"/v1/videos/{video_id}/content", - headers=self.proxy.transport.bearer(key), - ) - def image_edit( self, key: str, model: str, prompt: str, image: bytes, *, filename: str = "image.png" ) -> Result[ImagesResult]: diff --git a/tests/e2e/llm_translation/test_video_generation_e2e.py b/tests/e2e/llm_translation/test_video_generation_e2e.py deleted file mode 100644 index b65529aa260..00000000000 --- a/tests/e2e/llm_translation/test_video_generation_e2e.py +++ /dev/null @@ -1,69 +0,0 @@ -"""Live e2e: POST /v1/videos creates a video and serves its content. - -Registers a fal.ai Seedance deployment at runtime, polls the queued video until it -completes, and asserts the generated content is returned as binary data. -""" - -from __future__ import annotations - -import time -from typing import Final - -import pytest -from e2e_config import unique_marker -from e2e_http import require_successful_call, unwrap -from endpoints_client import EndpointsClient, VideoObject -from lifecycle import ResourceManager -from models import LiteLLMParamsBody - -pytestmark = pytest.mark.e2e - -_POLL_INTERVAL_SECONDS: Final[float] = 5.0 -_POLL_TIMEOUT_SECONDS: Final[float] = 600.0 - - -def _wait_for_completion( - endpoints_client: EndpointsClient, key: str, created: VideoObject -) -> VideoObject: - deadline = time.monotonic() + _POLL_TIMEOUT_SECONDS - while time.monotonic() < deadline: - status = unwrap(endpoints_client.video_status(key, created.id)) - assert status.id == created.id - if status.status == "completed": - return status - if status.status == "failed": - pytest.fail(f"fal.ai video generation failed: {status}") - time.sleep(_POLL_INTERVAL_SECONDS) - pytest.fail(f"fal.ai video {created.id!r} did not complete within {_POLL_TIMEOUT_SECONDS}s") - - -class TestVideoGeneration: - @pytest.mark.covers("llm.videos.fal_ai.basic.nonstream.works") - def test_fal_seedance_video_completes_and_downloads( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> None: - model = f"e2e-fal-video-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model="fal_ai/bytedance/seedance-2.5/text-to-video", - api_key="os.environ/FAL_AI_API_KEY", - ), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - result = endpoints_client.videos( - key, model, "a red fox running through snow at dawn" - ) - require_successful_call(result) - created = VideoObject.model_validate_json(result.body) - assert created.id - assert created.model - - _wait_for_completion(endpoints_client, key, created) - - content = endpoints_client.video_content(key, created.id) - require_successful_call(content) - assert len(content.body) > 0 - assert not (content.content_type or "").startswith("application/json") diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index 6958ade50f7..01f7af6e8fe 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -160,6 +160,9 @@ "other.provider_wire.anthropic.tool_history_system_cache_and_internal_fields", "quota_management.spend_tracking.cache_tokens.disjoint_classes_use_explicit_rates" ], + "tests/integration/providers/test_fal_ai_video_wire.py::test_fal_video_create_status_and_content_follow_queue_wire_contract": [ + "other.provider_wire.fal_ai.video_queue_create_status_and_content_download" + ], "tests/integration/mcp/test_mcp_lifecycle.py::test_saved_headers_reach_real_mcp_tool_and_survive_unrelated_edit": [ "mcp.call_tool.saved_headers.reach_actual_transport" ], diff --git a/tests/integration/providers/test_fal_ai_video_wire.py b/tests/integration/providers/test_fal_ai_video_wire.py new file mode 100644 index 00000000000..c1a2655f0aa --- /dev/null +++ b/tests/integration/providers/test_fal_ai_video_wire.py @@ -0,0 +1,72 @@ +import json +import sys +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server + +_MODEL: Final = "bytedance/seedance-2.5/text-to-video" +_MP4: Final = b"\x00\x00\x00\x18ftypmp42" + uuid.uuid4().bytes * 4 + + +@pytest.mark.covers("other.provider_wire.fal_ai.video_queue_create_status_and_content_download") +def test_fal_video_create_status_and_content_follow_queue_wire_contract(gateway: Gateway) -> None: + request_id: Final = "fal-req-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.target == f"/files/{request_id}.mp4": + assert request.method == "GET" + return Reply(body=_MP4, content_type="video/mp4") + assert request.headers["authorization"] == "Key synthetic-fal-key" + if request.method == "POST": + assert request.target == f"/{_MODEL}" + assert json.loads(request.body) == { + "prompt": "a cat playing volleyball on a beach", + "duration": "4", + "resolution": "720p", + "aspect_ratio": "16:9", + } + return Reply( + body=json.dumps({"status": "IN_QUEUE", "request_id": request_id, "queue_position": 0}).encode() + ) + assert request.method == "GET" + if request.target == f"/bytedance/seedance-2.5/requests/{request_id}/status": + return Reply(body=json.dumps({"status": "COMPLETED", "request_id": request_id}).encode()) + assert request.target == f"/bytedance/seedance-2.5/requests/{request_id}" + return Reply(body=json.dumps({"video": {"url": f"{wire_url}/files/{request_id}.mp4"}}).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + wire_url: Final = wire.url + model: Final = scenario.model( + model=f"fal_ai/{_MODEL}", + api_base=wire.url, + api_key="synthetic-fal-key", + ) + created: Final = gateway.post( + "/v1/videos", + { + "model": model, + "prompt": "a cat playing volleyball on a beach", + "seconds": "4", + "size": "1280x720", + }, + ) + assert created["status"] == "queued" + video_id: Final = created["id"] + assert isinstance(video_id, str) and video_id + status: Final = gateway.get(f"/v1/videos/{video_id}") + assert status["status"] == "completed" + status_id_matches_created_id: Final = status["id"] == video_id + sys.stdout.write(f"status_id_matches_created_id={status_id_matches_created_id}\n") + content: Final = gateway.request("GET", f"/v1/videos/{video_id}/content") + assert content.status_code == 200, content.text + assert content.headers["content-type"].startswith("video/mp4") + assert content.content == _MP4 + assert [(request.method, request.target) for request in wire.drain()] == [ + ("POST", f"/{_MODEL}"), + ("GET", f"/bytedance/seedance-2.5/requests/{request_id}/status"), + ("GET", f"/bytedance/seedance-2.5/requests/{request_id}"), + ("GET", f"/files/{request_id}.mp4"), + ] From 21a2ed62448ebda3ab9de1245b2550aa0bf164e0 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 23:29:25 +0000 Subject: [PATCH 12/47] test(integration): drop id diagnostic from fal video wire test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/providers/test_fal_ai_video_wire.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/integration/providers/test_fal_ai_video_wire.py b/tests/integration/providers/test_fal_ai_video_wire.py index c1a2655f0aa..8c72810ffb6 100644 --- a/tests/integration/providers/test_fal_ai_video_wire.py +++ b/tests/integration/providers/test_fal_ai_video_wire.py @@ -1,5 +1,4 @@ import json -import sys import uuid from typing import Final @@ -58,8 +57,6 @@ def test_fal_video_create_status_and_content_follow_queue_wire_contract(gateway: assert isinstance(video_id, str) and video_id status: Final = gateway.get(f"/v1/videos/{video_id}") assert status["status"] == "completed" - status_id_matches_created_id: Final = status["id"] == video_id - sys.stdout.write(f"status_id_matches_created_id={status_id_matches_created_id}\n") content: Final = gateway.request("GET", f"/v1/videos/{video_id}/content") assert content.status_code == 200, content.text assert content.headers["content-type"].startswith("video/mp4") From b450baa402527c5be47a532508d7708fd1ca9a41 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 07:47:33 +0000 Subject: [PATCH 13/47] test(llms): migrate phase 5 provider unit tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../image/test_bedrock_image_bearer_token.py | 158 ------------------ .../test_amazon_nova_canvas_transformation.py | 0 .../test_amazon_stability3_transformation.py | 0 .../image/test_bedrock_image_bearer_token.py | 21 +++ .../test_bedrock_image_prepare_request.py | 0 .../test_amazon_nova_canvas_image_edit.py | 0 .../test_bedrock_agent_transformation.py | 0 .../guardrail_translation/test_handler.py | 0 ...test_bedrock_passthrough_transformation.py | 2 - .../realtime/test_bedrock_realtime_handler.py | 0 .../test_bedrock_realtime_transformation.py | 0 .../test_bedrock_rerank_header_forwarding.py | 0 ...est_bedrock_vector_store_transformation.py | 0 ...drock_mantle_passthrough_transformation.py | 0 .../test_bfl_image_edit_transformation.py | 0 ...est_bfl_image_generation_transformation.py | 0 .../test_bfl_common_utils.py | 0 .../chat/test_bytez_chat_transformation.py | 0 .../test_cerebras_chat_transformation.py | 0 .../llms/chat/test_converse_handler.py | 0 .../chatgpt/test_chatgpt_authenticator.py | 0 21 files changed, 21 insertions(+), 160 deletions(-) delete mode 100644 tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py rename tests/{test_litellm => unit}/llms/bedrock/image/test_amazon_nova_canvas_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/image/test_amazon_stability3_transformation.py (100%) create mode 100644 tests/unit/llms/bedrock/image/test_bedrock_image_bearer_token.py rename tests/{test_litellm => unit}/llms/bedrock/image/test_bedrock_image_prepare_request.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/passthrough/guardrail_translation/test_handler.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py (99%) rename tests/{test_litellm => unit}/llms/bedrock/realtime/test_bedrock_realtime_handler.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/realtime/test_bedrock_realtime_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py (100%) rename tests/{test_litellm => unit}/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py (100%) rename tests/{test_litellm => unit}/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py (100%) rename tests/{test_litellm => unit}/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py (100%) rename tests/{test_litellm => unit}/llms/black_forest_labs/test_bfl_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/bytez/chat/test_bytez_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/cerebras/test_cerebras_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/chat/test_converse_handler.py (100%) rename tests/{test_litellm => unit}/llms/chatgpt/test_chatgpt_authenticator.py (100%) diff --git a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py b/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py deleted file mode 100644 index 0b11a66c100..00000000000 --- a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py +++ /dev/null @@ -1,158 +0,0 @@ -import json -import os -from unittest.mock import Mock, patch -import pytest - - -import litellm -from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler - -# Mock response for Bedrock image generation -mock_image_response = {"images": ["base64_encoded_image_data"], "error": None} - - -class TestBedrockImageGeneration: - def test_image_generation_with_api_key_bearer_token(self): - """Test image generation with bearer token authentication""" - test_api_key = "test-bearer-token-12345" - model = "bedrock/stability.sd3-large-v1:0" - prompt = "A cute baby sea otter" - - with patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.image_generation" - ) as mock_bedrock_image_gen: - # Setup mock response - mock_image_response_obj = litellm.ImageResponse() - mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}] - mock_bedrock_image_gen.return_value = mock_image_response_obj - - response = litellm.image_generation( - model=model, - prompt=prompt, - aws_region_name="us-west-2", - api_key=test_api_key, - ) - - assert response is not None - assert len(response.data) > 0 - - mock_bedrock_image_gen.assert_called_once() - for call in mock_bedrock_image_gen.call_args_list: - if "headers" in call.kwargs: - headers = call.kwargs["headers"] - if ( - "Authorization" in headers - and headers["Authorization"] == f"Bearer {test_api_key}" - ): - break - - def test_image_generation_with_env_variable_bearer_token(self, monkeypatch): - """Test image generation with bearer token from environment variable""" - test_api_key = "env-bearer-token-12345" - model = "bedrock/stability.sd3-large-v1:0" - prompt = "A cute baby sea otter" - - # Mock the environment variable - with ( - patch.dict(os.environ, {"AWS_BEARER_TOKEN_BEDROCK": test_api_key}), - patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.image_generation" - ) as mock_bedrock_image_gen, - ): - - mock_image_response_obj = litellm.ImageResponse() - mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}] - mock_bedrock_image_gen.return_value = mock_image_response_obj - - response = litellm.image_generation( - model=model, prompt=prompt, aws_region_name="us-west-2" - ) - - assert response is not None - assert len(response.data) > 0 - - mock_bedrock_image_gen.assert_called_once() - for call in mock_bedrock_image_gen.call_args_list: - if "headers" in call.kwargs: - headers = call.kwargs["headers"] - if ( - "Authorization" in headers - and headers["Authorization"] == f"Bearer {test_api_key}" - ): - break - - @pytest.mark.asyncio - async def test_async_image_generation_with_bearer_token(self): - """Test async image generation with bearer token authentication""" - test_api_key = "async-bearer-token-12345" - model = "bedrock/stability.sd3-large-v1:0" - prompt = "A cute baby sea otter" - - with patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.async_image_generation" - ) as mock_async_bedrock_image_gen: - mock_image_response_obj = litellm.ImageResponse() - mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}] - mock_async_bedrock_image_gen.return_value = mock_image_response_obj - - # Call async image generation with api_key parameter - response = await litellm.aimage_generation( - model=model, - prompt=prompt, - aws_region_name="us-west-2", - api_key=test_api_key, - ) - - assert response is not None - assert len(response.data) > 0 - - mock_async_bedrock_image_gen.assert_called_once() - for call in mock_async_bedrock_image_gen.call_args_list: - if "headers" in call.kwargs: - headers = call.kwargs["headers"] - if ( - "Authorization" in headers - and headers["Authorization"] == f"Bearer {test_api_key}" - ): - break - - def test_image_generation_with_sigv4(self): - """Test image generation falls back to SigV4 auth when no bearer token is provided""" - model = "bedrock/stability.sd3-large-v1:0" - prompt = "A cute baby sea otter" - - with patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.image_generation" - ) as mock_bedrock_image_gen: - mock_image_response_obj = litellm.ImageResponse() - mock_image_response_obj.data = [{"url": "https://example.com/image.jpg"}] - mock_bedrock_image_gen.return_value = mock_image_response_obj - - response = litellm.image_generation( - model=model, prompt=prompt, aws_region_name="us-west-2" - ) - - assert response is not None - assert len(response.data) > 0 - mock_bedrock_image_gen.assert_called_once() - - -def test_image_generation_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): - """The deployment's AWS profile does not exist, so resolving SigV4 credentials - raises; a bearer-token deployment must still sign the request with the - bearer token alone.""" - from litellm.llms.bedrock.image_generation.image_handler import BedrockImageGeneration - - monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") - - request = BedrockImageGeneration()._prepare_request( - model="amazon.nova-canvas-v1:0", - prompt="A cute baby sea otter", - optional_params={"aws_region_name": "us-west-2", "aws_profile_name": "litellm-no-such-aws-profile"}, - api_base=None, - extra_headers=None, - api_key=None, - logging_obj=Mock(), - ) - - assert request.prepped.headers["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/llms/bedrock/image/test_amazon_nova_canvas_transformation.py b/tests/unit/llms/bedrock/image/test_amazon_nova_canvas_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/image/test_amazon_nova_canvas_transformation.py rename to tests/unit/llms/bedrock/image/test_amazon_nova_canvas_transformation.py diff --git a/tests/test_litellm/llms/bedrock/image/test_amazon_stability3_transformation.py b/tests/unit/llms/bedrock/image/test_amazon_stability3_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/image/test_amazon_stability3_transformation.py rename to tests/unit/llms/bedrock/image/test_amazon_stability3_transformation.py diff --git a/tests/unit/llms/bedrock/image/test_bedrock_image_bearer_token.py b/tests/unit/llms/bedrock/image/test_bedrock_image_bearer_token.py new file mode 100644 index 00000000000..599507da03d --- /dev/null +++ b/tests/unit/llms/bedrock/image/test_bedrock_image_bearer_token.py @@ -0,0 +1,21 @@ +from unittest.mock import Mock + +def test_image_generation_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): + """The deployment's AWS profile does not exist, so resolving SigV4 credentials + raises; a bearer-token deployment must still sign the request with the + bearer token alone.""" + from litellm.llms.bedrock.image_generation.image_handler import BedrockImageGeneration + + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345") + + request = BedrockImageGeneration()._prepare_request( + model="amazon.nova-canvas-v1:0", + prompt="A cute baby sea otter", + optional_params={"aws_region_name": "us-west-2", "aws_profile_name": "litellm-no-such-aws-profile"}, + api_base=None, + extra_headers=None, + api_key=None, + logging_obj=Mock(), + ) + + assert request.prepped.headers["Authorization"] == "Bearer env-bearer-token-12345" diff --git a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_prepare_request.py b/tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py similarity index 100% rename from tests/test_litellm/llms/bedrock/image/test_bedrock_image_prepare_request.py rename to tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py diff --git a/tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py b/tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py similarity index 100% rename from tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py rename to tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py diff --git a/tests/test_litellm/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py b/tests/unit/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py rename to tests/unit/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py diff --git a/tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py similarity index 100% rename from tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py rename to tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py diff --git a/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py similarity index 99% rename from tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py rename to tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index f2a9af11af7..854ef92fa4b 100644 --- a/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -367,8 +367,6 @@ def test_bedrock_passthrough_region_extraction_from_inference_profile_arn(): assert ( "us-west-2" in api_base ), f"Expected region 'us-west-2' from ARN in base URL, but got: {api_base}" - - def test_bedrock_passthrough_model_id_arn_encoding(): """ Test that model_id ARNs are properly URL-encoded when used in endpoints. diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py similarity index 100% rename from tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py rename to tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py rename to tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py similarity index 100% rename from tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py rename to tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py rename to tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py diff --git a/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py b/tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py rename to tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py diff --git a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py b/tests/unit/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py rename to tests/unit/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py diff --git a/tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py b/tests/unit/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py rename to tests/unit/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py diff --git a/tests/test_litellm/llms/black_forest_labs/test_bfl_common_utils.py b/tests/unit/llms/black_forest_labs/test_bfl_common_utils.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/test_bfl_common_utils.py rename to tests/unit/llms/black_forest_labs/test_bfl_common_utils.py diff --git a/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py b/tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py rename to tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py diff --git a/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py b/tests/unit/llms/cerebras/test_cerebras_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py rename to tests/unit/llms/cerebras/test_cerebras_chat_transformation.py diff --git a/tests/test_litellm/llms/chat/test_converse_handler.py b/tests/unit/llms/chat/test_converse_handler.py similarity index 100% rename from tests/test_litellm/llms/chat/test_converse_handler.py rename to tests/unit/llms/chat/test_converse_handler.py diff --git a/tests/test_litellm/llms/chatgpt/test_chatgpt_authenticator.py b/tests/unit/llms/chatgpt/test_chatgpt_authenticator.py similarity index 100% rename from tests/test_litellm/llms/chatgpt/test_chatgpt_authenticator.py rename to tests/unit/llms/chatgpt/test_chatgpt_authenticator.py From 4defed7f2e7eaacdf8e130857eddc6018c7bc3f7 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:03:54 +0000 Subject: [PATCH 14/47] test: migrate wave 1 phase 8 legacy llm tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat/test_hosted_vllm_ssl_verify.py | 147 ------------------ .../test_hosted_vllm_embedding_ssl_verify.py | 135 ---------------- ..._github_copilot_messages_transformation.py | 7 - ...github_copilot_responses_transformation.py | 116 +++++--------- .../test_gradient_ai_chat_transformation.py | 0 .../chat/test_groq_chat_transformation.py | 2 - .../llms/groq/test_groq_cost_calculator.py | 0 .../test_hosted_vllm_chat_transformation.py | 71 +-------- ...st_hosted_vllm_embedding_transformation.py | 8 +- ...t_hosted_vllm_image_edit_transformation.py | 0 .../responses/test_hosted_vllm_responses.py | 9 +- .../test_hosted_vllm_rerank_transformation.py | 0 .../test_hosted_vllm_video_transformation.py | 0 .../test_huggingface_rerank_transformation.py | 40 +---- .../test_inception_chat_transformation.py | 18 +-- ...est_inception_completion_transformation.py | 18 +-- .../test_jina_embedding_transformation.py | 0 .../chat/test_langflow_chat_transformation.py | 31 +--- .../litellm_proxy/test_sandbox_executor.py | 25 +-- .../litellm_proxy/test_skills_ownership.py | 73 ++------- 20 files changed, 84 insertions(+), 616 deletions(-) delete mode 100644 tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_ssl_verify.py delete mode 100644 tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_ssl_verify.py rename tests/{test_litellm => unit}/llms/github_copilot/messages/test_github_copilot_messages_transformation.py (98%) rename tests/{test_litellm => unit}/llms/github_copilot/responses/test_github_copilot_responses_transformation.py (89%) rename tests/{test_litellm => unit}/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/groq/chat/test_groq_chat_transformation.py (99%) rename tests/{test_litellm => unit}/llms/groq/test_groq_cost_calculator.py (100%) rename tests/{test_litellm => unit}/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py (82%) rename tests/{test_litellm => unit}/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py (97%) rename tests/{test_litellm => unit}/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/hosted_vllm/responses/test_hosted_vllm_responses.py (96%) rename tests/{test_litellm => unit}/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py (100%) rename tests/{test_litellm => unit}/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py (100%) rename tests/{test_litellm => unit}/llms/huggingface/rerank/test_huggingface_rerank_transformation.py (91%) rename tests/{test_litellm => unit}/llms/inception/test_inception_chat_transformation.py (96%) rename tests/{test_litellm => unit}/llms/inception/test_inception_completion_transformation.py (95%) rename tests/{test_litellm => unit}/llms/jina_ai/embedding/test_jina_embedding_transformation.py (100%) rename tests/{test_litellm => unit}/llms/langflow/chat/test_langflow_chat_transformation.py (93%) rename tests/{test_litellm => unit}/llms/litellm_proxy/test_sandbox_executor.py (84%) rename tests/{test_litellm => unit}/llms/litellm_proxy/test_skills_ownership.py (88%) diff --git a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_ssl_verify.py b/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_ssl_verify.py deleted file mode 100644 index 2364468efe1..00000000000 --- a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_ssl_verify.py +++ /dev/null @@ -1,147 +0,0 @@ -""" -Test SSL verification for hosted_vllm provider. - -This test ensures that the ssl_verify parameter is properly passed through -to the HTTP client when using the hosted_vllm provider. - -Issue: ssl_verify parameter was being ignored because hosted_vllm fell through -to the OpenAI catch-all path in main.py, which doesn't pass ssl_verify to the HTTP client. -""" - -from unittest.mock import MagicMock, patch - -import pytest - - -import litellm - - -class TestHostedVLLMSSLVerify: - """Test suite for SSL verification in hosted_vllm provider.""" - - @patch("litellm.llms.custom_httpx.llm_http_handler._get_httpx_client") - def test_hosted_vllm_ssl_verify_false_sync(self, mock_get_httpx_client): - """Test that ssl_verify=False is passed to the HTTP client for sync calls.""" - # Setup mock client - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = { - "id": "chatcmpl-test", - "object": "chat.completion", - "created": 1234567890, - "model": "test-model", - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": "Test response", - }, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 5, - "total_tokens": 15, - }, - } - mock_response.text = '{"id": "chatcmpl-test", "object": "chat.completion", "created": 1234567890, "model": "test-model", "choices": [{"index": 0, "message": {"role": "assistant", "content": "Test response"}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}' - mock_client.post.return_value = mock_response - mock_get_httpx_client.return_value = mock_client - - try: - litellm.completion( - model="hosted_vllm/test-model", - messages=[{"role": "user", "content": "Hello"}], - api_base="https://test-vllm.example.com/v1", - ssl_verify=False, - ) - except Exception: - # Even if the response parsing fails, we just need to verify - # that the mock was called with the correct ssl_verify parameter - pass - - # Verify _get_httpx_client was called with ssl_verify=False - mock_get_httpx_client.assert_called() - call_args = mock_get_httpx_client.call_args - - # Check that params contains ssl_verify=False - if call_args[0]: - # Positional argument - params = call_args[0][0] - else: - # Keyword argument - params = call_args[1].get("params", {}) - - assert ( - params.get("ssl_verify") is False - ), f"Expected ssl_verify=False in params, got {params}" - - @patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client") - @pytest.mark.asyncio - async def test_hosted_vllm_ssl_verify_false_async( - self, mock_get_async_httpx_client - ): - """Test that ssl_verify=False is passed to the HTTP client for async calls.""" - # Setup mock async client - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = { - "id": "chatcmpl-test", - "object": "chat.completion", - "created": 1234567890, - "model": "test-model", - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": "Test response", - }, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 5, - "total_tokens": 15, - }, - } - mock_response.text = '{"id": "chatcmpl-test", "object": "chat.completion", "created": 1234567890, "model": "test-model", "choices": [{"index": 0, "message": {"role": "assistant", "content": "Test response"}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}' - - async def mock_post(*args, **kwargs): - return mock_response - - mock_client.post = mock_post - mock_get_async_httpx_client.return_value = mock_client - - try: - await litellm.acompletion( - model="hosted_vllm/test-model", - messages=[{"role": "user", "content": "Hello"}], - api_base="https://test-vllm.example.com/v1", - ssl_verify=False, - ) - except Exception: - # Even if the response parsing fails, we just need to verify - # that the mock was called with the correct ssl_verify parameter - pass - - # Verify get_async_httpx_client was called with ssl_verify=False - mock_get_async_httpx_client.assert_called() - call_kwargs = mock_get_async_httpx_client.call_args[1] - - # Check that params contains ssl_verify=False - params = call_kwargs.get("params", {}) - assert ( - params.get("ssl_verify") is False - ), f"Expected ssl_verify=False in params, got {params}" - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "-s"]) diff --git a/tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_ssl_verify.py b/tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_ssl_verify.py deleted file mode 100644 index de94da49384..00000000000 --- a/tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_ssl_verify.py +++ /dev/null @@ -1,135 +0,0 @@ -""" -Test SSL verification for hosted_vllm provider embeddings. - -This test ensures that the ssl_verify parameter is properly passed through -to the HTTP client when using the hosted_vllm provider for embeddings. - -Issue: ssl_verify parameter was being ignored because hosted_vllm fell through -to the openai_like catch-all path in main.py, which doesn't pass ssl_verify to the HTTP client. -""" - -from unittest.mock import MagicMock, patch - -import pytest - - -import litellm - - -class TestHostedVLLMEmbeddingSSLVerify: - """Test suite for SSL verification in hosted_vllm provider embeddings.""" - - @patch("litellm.llms.custom_httpx.llm_http_handler._get_httpx_client") - def test_hosted_vllm_embedding_ssl_verify_false_sync(self, mock_get_httpx_client): - """Test that ssl_verify=False is passed to the HTTP client for sync embedding calls.""" - # Setup mock client - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = { - "object": "list", - "data": [ - { - "object": "embedding", - "index": 0, - "embedding": [0.1, 0.2, 0.3, 0.4, 0.5], - } - ], - "model": "text-embedding-model", - "usage": { - "prompt_tokens": 5, - "total_tokens": 5, - }, - } - mock_response.text = '{"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3, 0.4, 0.5]}], "model": "text-embedding-model", "usage": {"prompt_tokens": 5, "total_tokens": 5}}' - mock_client.post.return_value = mock_response - mock_get_httpx_client.return_value = mock_client - - try: - litellm.embedding( - model="hosted_vllm/text-embedding-model", - input=["hello world"], - api_base="https://test-vllm.example.com/v1", - ssl_verify=False, - ) - except Exception: - # Even if the response parsing fails, we just need to verify - # that the mock was called with the correct ssl_verify parameter - pass - - # Verify _get_httpx_client was called with ssl_verify=False - mock_get_httpx_client.assert_called() - call_args = mock_get_httpx_client.call_args - - # Check that params contains ssl_verify=False - if call_args[0]: - # Positional argument - params = call_args[0][0] - else: - # Keyword argument - params = call_args[1].get("params", {}) - - assert ( - params.get("ssl_verify") is False - ), f"Expected ssl_verify=False in params, got {params}" - - @patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client") - @pytest.mark.asyncio - async def test_hosted_vllm_embedding_ssl_verify_false_async( - self, mock_get_async_httpx_client - ): - """Test that ssl_verify=False is passed to the HTTP client for async embedding calls.""" - # Setup mock async client - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = { - "object": "list", - "data": [ - { - "object": "embedding", - "index": 0, - "embedding": [0.1, 0.2, 0.3, 0.4, 0.5], - } - ], - "model": "text-embedding-model", - "usage": { - "prompt_tokens": 5, - "total_tokens": 5, - }, - } - mock_response.text = '{"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3, 0.4, 0.5]}], "model": "text-embedding-model", "usage": {"prompt_tokens": 5, "total_tokens": 5}}' - - async def mock_post(*args, **kwargs): - return mock_response - - mock_client.post = mock_post - mock_get_async_httpx_client.return_value = mock_client - - try: - await litellm.aembedding( - model="hosted_vllm/text-embedding-model", - input=["hello world"], - api_base="https://test-vllm.example.com/v1", - ssl_verify=False, - ) - except Exception: - # Even if the response parsing fails, we just need to verify - # that the mock was called with the correct ssl_verify parameter - pass - - # Verify get_async_httpx_client was called with ssl_verify=False - mock_get_async_httpx_client.assert_called() - call_kwargs = mock_get_async_httpx_client.call_args[1] - - # Check that params contains ssl_verify=False - params = call_kwargs.get("params", {}) - assert ( - params.get("ssl_verify") is False - ), f"Expected ssl_verify=False in params, got {params}" - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "-s"]) diff --git a/tests/test_litellm/llms/github_copilot/messages/test_github_copilot_messages_transformation.py b/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py similarity index 98% rename from tests/test_litellm/llms/github_copilot/messages/test_github_copilot_messages_transformation.py rename to tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py index 8039e744f46..9e9760650cf 100644 --- a/tests/test_litellm/llms/github_copilot/messages/test_github_copilot_messages_transformation.py +++ b/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py @@ -10,13 +10,6 @@ from litellm.llms.github_copilot.messages.transformation import ( ) -def test_github_copilot_anthropic_messages_config_init(): - """Test GithubCopilotAnthropicMessagesConfig initialization.""" - config = GithubCopilotAnthropicMessagesConfig() - assert config is not None - assert hasattr(config, "authenticator") - - def test_github_copilot_anthropic_messages_get_complete_url(): """get_complete_url builds the /v1/messages URL from the base it is handed. diff --git a/tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py b/tests/unit/llms/github_copilot/responses/test_github_copilot_responses_transformation.py similarity index 89% rename from tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py rename to tests/unit/llms/github_copilot/responses/test_github_copilot_responses_transformation.py index 0174465b0cc..b8380b7adb4 100644 --- a/tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py +++ b/tests/unit/llms/github_copilot/responses/test_github_copilot_responses_transformation.py @@ -26,9 +26,7 @@ def use_local_model_cost_map(monkeypatch: pytest.MonkeyPatch): """Pin litellm.model_cost to the bundled local backup so tests don't depend on remote catalog fetches (and don't change behavior across remote refreshes).""" monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr( - litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url) - ) + monkeypatch.setattr(litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url)) litellm.add_known_models(model_cost_map=litellm.model_cost) @@ -44,49 +42,35 @@ class TestGithubCopilotResponsesAPITransformation: provider=LlmProviders.GITHUB_COPILOT, ) - assert ( - config is not None - ), "Config should not be None for GitHub Copilot provider" - assert isinstance( - config, GithubCopilotResponsesAPIConfig - ), f"Expected GithubCopilotResponsesAPIConfig, got {type(config)}" - assert ( - config.custom_llm_provider == LlmProviders.GITHUB_COPILOT - ), "custom_llm_provider should be GITHUB_COPILOT" + assert config is not None, "Config should not be None for GitHub Copilot provider" + assert isinstance(config, GithubCopilotResponsesAPIConfig), ( + f"Expected GithubCopilotResponsesAPIConfig, got {type(config)}" + ) + assert config.custom_llm_provider == LlmProviders.GITHUB_COPILOT, "custom_llm_provider should be GITHUB_COPILOT" @patch("litellm.llms.github_copilot.responses.transformation.Authenticator") def test_github_copilot_responses_endpoint_url(self, mock_authenticator_class): """Test that get_complete_url returns correct GitHub Copilot endpoint""" # Mock authenticator to return default base mock_auth_instance = MagicMock() - mock_auth_instance.get_api_base.return_value = ( - "https://api.individual.githubcopilot.com" - ) + mock_auth_instance.get_api_base.return_value = "https://api.individual.githubcopilot.com" mock_authenticator_class.return_value = mock_auth_instance config = GithubCopilotResponsesAPIConfig() # Test with default GitHub Copilot API base (from authenticator) url = config.get_complete_url(api_base=None, litellm_params={}) - assert ( - url == "https://api.individual.githubcopilot.com/responses" - ), f"Expected GitHub Copilot responses endpoint, got {url}" + assert url == "https://api.individual.githubcopilot.com/responses", ( + f"Expected GitHub Copilot responses endpoint, got {url}" + ) # Test with custom api_base (overrides authenticator) - custom_url = config.get_complete_url( - api_base="https://custom.githubcopilot.com", litellm_params={} - ) - assert ( - custom_url == "https://custom.githubcopilot.com/responses" - ), f"Expected custom endpoint, got {custom_url}" + custom_url = config.get_complete_url(api_base="https://custom.githubcopilot.com", litellm_params={}) + assert custom_url == "https://custom.githubcopilot.com/responses", f"Expected custom endpoint, got {custom_url}" # Test with trailing slash - url_with_slash = config.get_complete_url( - api_base="https://api.githubcopilot.com/", litellm_params={} - ) - assert ( - url_with_slash == "https://api.githubcopilot.com/responses" - ), "Should handle trailing slash" + url_with_slash = config.get_complete_url(api_base="https://api.githubcopilot.com/", litellm_params={}) + assert url_with_slash == "https://api.githubcopilot.com/responses", "Should handle trailing slash" @patch("litellm.llms.github_copilot.responses.transformation.Authenticator") def test_validate_environment_default_headers(self, mock_authenticator_class): @@ -98,9 +82,7 @@ class TestGithubCopilotResponsesAPITransformation: config = GithubCopilotResponsesAPIConfig() - headers = config.validate_environment( - headers={}, model="gpt-5.1-codex", litellm_params={} - ) + headers = config.validate_environment(headers={}, model="gpt-5.1-codex", litellm_params={}) # Check required headers assert headers["Authorization"] == "Bearer test-api-key-123" @@ -127,9 +109,7 @@ class TestGithubCopilotResponsesAPITransformation: "custom-header": "custom-value", } - headers = config.validate_environment( - headers=custom_headers, model="gpt-5.1-codex", litellm_params={} - ) + headers = config.validate_environment(headers=custom_headers, model="gpt-5.1-codex", litellm_params={}) # User header should override default assert headers["editor-version"] == "custom/2.0.0" @@ -182,9 +162,7 @@ class TestGithubCopilotResponsesAPITransformation: """Test _has_vision_input detects input_image type""" config = GithubCopilotResponsesAPIConfig() - input_with_vision = [ - {"role": "user", "content": [{"type": "input_image", "data": "base64..."}]} - ] + input_with_vision = [{"role": "user", "content": [{"type": "input_image", "data": "base64..."}]}] has_vision = config._has_vision_input(input_with_vision) assert has_vision is True, "Should detect input_image type" @@ -246,13 +224,11 @@ class TestGithubCopilotResponsesAPITransformation: } ] - headers = config.validate_environment( - headers={}, model="gpt-5.1-codex", litellm_params=mock_litellm_params - ) + headers = config.validate_environment(headers={}, model="gpt-5.1-codex", litellm_params=mock_litellm_params) - assert ( - headers.get("copilot-vision-request") == "true" - ), "Should add copilot-vision-request header for vision input" + assert headers.get("copilot-vision-request") == "true", ( + "Should add copilot-vision-request header for vision input" + ) @patch("litellm.llms.github_copilot.responses.transformation.Authenticator") def test_validate_environment_with_x_initiator(self, mock_authenticator_class): @@ -270,21 +246,15 @@ class TestGithubCopilotResponsesAPITransformation: {"role": "assistant", "content": "Hi"}, ] - headers = config.validate_environment( - headers={}, model="gpt-5.1-codex", litellm_params=mock_litellm_params - ) + headers = config.validate_environment(headers={}, model="gpt-5.1-codex", litellm_params=mock_litellm_params) - assert ( - headers.get("X-Initiator") == "agent" - ), "Should set X-Initiator to 'agent' for assistant role" + assert headers.get("X-Initiator") == "agent", "Should set X-Initiator to 'agent' for assistant role" def test_map_openai_params_no_transformation(self): """Test that map_openai_params passes through parameters unchanged""" config = GithubCopilotResponsesAPIConfig() - params = ResponsesAPIOptionalRequestParams( - temperature=0.7, max_output_tokens=1000, stream=False - ) + params = ResponsesAPIOptionalRequestParams(temperature=0.7, max_output_tokens=1000, stream=False) result = config.map_openai_params( response_api_optional_params=params, @@ -338,9 +308,9 @@ class TestGithubCopilotResponsesAPITransformation: result = config._handle_reasoning_item(reasoning_item) # encrypted_content should be preserved - assert ( - result.get("encrypted_content") == "encrypted-blob-abc123" - ), "encrypted_content must be preserved for GitHub Copilot multi-turn conversations" + assert result.get("encrypted_content") == "encrypted-blob-abc123", ( + "encrypted_content must be preserved for GitHub Copilot multi-turn conversations" + ) # status=None should be filtered out assert "status" not in result, "status=None should be filtered out" # content=None should be filtered out @@ -393,9 +363,7 @@ class TestGithubCopilotResponsesAPIRouting: in the (already-merged) model info; otherwise returns None so the dispatcher routes through the chat-completions translation bridge.""" - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_returns_config_when_mode_is_responses(self, mock_get_info): """``mode=responses`` returns native config.""" mock_get_info.return_value = {"mode": "responses"} @@ -405,9 +373,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert isinstance(config, GithubCopilotResponsesAPIConfig) - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_returns_none_when_mode_is_chat(self, mock_get_info): """``mode=chat`` returns None so dispatcher uses bridge.""" mock_get_info.return_value = {"mode": "chat"} @@ -417,9 +383,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert config is None - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_returns_none_when_mode_is_unset_and_no_endpoints(self, mock_get_info): """Entry without ``mode`` and without ``supported_endpoints`` returns None (conservative default).""" @@ -499,9 +463,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert isinstance(config, GithubCopilotResponsesAPIConfig) - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_returns_none_when_get_model_info_raises(self, mock_get_info): """Catalog lookup failure (model not registered) returns None (conservative default; bridge handles unknown models safely).""" @@ -512,9 +474,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert config is None - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_user_override_via_register_model(self, mock_get_info): """User-supplied per-deployment ``model_info`` flows through ``litellm.register_model`` (called by the router) into the merged @@ -528,9 +488,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert isinstance(config, GithubCopilotResponsesAPIConfig) - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_realistic_chat_only_entry_returns_none(self, mock_get_info): """Realistic ``model_prices_and_context_window.json`` shape for a chat-only Copilot model (e.g. github_copilot/gemini-3.1-pro-preview) @@ -554,9 +512,7 @@ class TestGithubCopilotResponsesAPIRouting: ) assert config is None - @patch( - "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" - ) + @patch("litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper") def test_realistic_responses_only_entry_returns_config(self, mock_get_info): """Realistic catalog entry for a Responses-only Copilot model (e.g. github_copilot/gpt-5.5) returns the native config.""" @@ -592,9 +548,7 @@ class TestGithubCopilotReasoningStreamItemIdNormalization: output_index group to the id from its output_item.added.""" def _config(self): - with patch( - "litellm.llms.github_copilot.responses.transformation.Authenticator" - ): + with patch("litellm.llms.github_copilot.responses.transformation.Authenticator"): return GithubCopilotResponsesAPIConfig() def _transform(self, config, chunk): diff --git a/tests/test_litellm/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py b/tests/unit/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py rename to tests/unit/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py diff --git a/tests/test_litellm/llms/groq/chat/test_groq_chat_transformation.py b/tests/unit/llms/groq/chat/test_groq_chat_transformation.py similarity index 99% rename from tests/test_litellm/llms/groq/chat/test_groq_chat_transformation.py rename to tests/unit/llms/groq/chat/test_groq_chat_transformation.py index f605958b979..f5a7a920124 100644 --- a/tests/test_litellm/llms/groq/chat/test_groq_chat_transformation.py +++ b/tests/unit/llms/groq/chat/test_groq_chat_transformation.py @@ -202,5 +202,3 @@ class TestGroqWebSearchUsageSignal: model_response = litellm.ModelResponse() GroqChatConfig()._add_web_search_usage(model_response=model_response) assert getattr(model_response, "usage", None) is None - - diff --git a/tests/test_litellm/llms/groq/test_groq_cost_calculator.py b/tests/unit/llms/groq/test_groq_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/groq/test_groq_cost_calculator.py rename to tests/unit/llms/groq/test_groq_cost_calculator.py diff --git a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py similarity index 82% rename from tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py rename to tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py index 82b05601a85..1cc6a1457fc 100644 --- a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py +++ b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py @@ -41,74 +41,9 @@ def test_hosted_vllm_chat_transformation_file_url(): ] -def test_hosted_vllm_chat_transformation_with_audio_url(): - from litellm import completion - - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = { - "id": "chatcmpl-test", - "object": "chat.completion", - "created": 1234567890, - "model": "llama-3.1-70b-instruct", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "Test response"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, - } - mock_response.text = json.dumps(mock_response.json.return_value) - mock_client.post.return_value = mock_response - - with patch( - "litellm.llms.custom_httpx.llm_http_handler._get_httpx_client", - return_value=mock_client, - ): - try: - completion( - model="hosted_vllm/llama-3.1-70b-instruct", - messages=[ - { - "role": "user", - "content": [ - { - "type": "audio_url", - "audio_url": {"url": "https://example.com/audio.mp3"}, - }, - ], - }, - ], - api_base="https://test-vllm.example.com/v1", - ) - except Exception: - pass - - mock_client.post.assert_called_once() - call_kwargs = mock_client.post.call_args[1] - request_data = json.loads(call_kwargs["data"]) - assert request_data["messages"] == [ - { - "role": "user", - "content": [ - { - "type": "audio_url", - "audio_url": {"url": "https://example.com/audio.mp3"}, - } - ], - } - ] - - def test_hosted_vllm_supports_reasoning_effort(): config = HostedVLLMChatConfig() - supported_params = config.get_supported_openai_params( - model="hosted_vllm/gpt-oss-120b" - ) + supported_params = config.get_supported_openai_params(model="hosted_vllm/gpt-oss-120b") assert "reasoning_effort" in supported_params optional_params = config.map_openai_params( non_default_params={"reasoning_effort": "high"}, @@ -129,9 +64,7 @@ def test_hosted_vllm_supports_thinking(): Related issue: https://github.com/BerriAI/litellm/issues/19761 """ config = HostedVLLMChatConfig() - supported_params = config.get_supported_openai_params( - model="hosted_vllm/GLM-4.6-FP8" - ) + supported_params = config.get_supported_openai_params(model="hosted_vllm/GLM-4.6-FP8") assert "thinking" in supported_params # Test thinking below the low threshold -> "minimal" diff --git a/tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py b/tests/unit/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py similarity index 97% rename from tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py rename to tests/unit/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py index 34be3e12abd..5854b1596b4 100644 --- a/tests/test_litellm/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py +++ b/tests/unit/llms/hosted_vllm/embedding/test_hosted_vllm_embedding_transformation.py @@ -87,9 +87,7 @@ class TestHostedVLLMEmbeddingTransformation: headers={}, ) - assert ( - "encoding_format" not in result - ), "encoding_format should not be in request when not provided" + assert "encoding_format" not in result, "encoding_format should not be in request when not provided" def test_encoding_format_not_included_when_none(self): """ @@ -278,9 +276,7 @@ class TestHostedVLLMEmbeddingTransformation: sent_data = json.loads(call_kwargs["data"]) # Assert that encoding_format is NOT in the sent data - assert ( - "encoding_format" not in sent_data - ), "encoding_format should not be in request when not provided" + assert "encoding_format" not in sent_data, "encoding_format should not be in request when not provided" assert sent_data["model"] == "BAAI/bge-small-en-v1.5" assert sent_data["input"] == ["Hello world"] diff --git a/tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py b/tests/unit/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py rename to tests/unit/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py diff --git a/tests/test_litellm/llms/hosted_vllm/responses/test_hosted_vllm_responses.py b/tests/unit/llms/hosted_vllm/responses/test_hosted_vllm_responses.py similarity index 96% rename from tests/test_litellm/llms/hosted_vllm/responses/test_hosted_vllm_responses.py rename to tests/unit/llms/hosted_vllm/responses/test_hosted_vllm_responses.py index e81bf0c4f1f..55d0ce1e68e 100644 --- a/tests/test_litellm/llms/hosted_vllm/responses/test_hosted_vllm_responses.py +++ b/tests/unit/llms/hosted_vllm/responses/test_hosted_vllm_responses.py @@ -68,9 +68,7 @@ def test_hosted_vllm_responses_create_with_string_input(): Test that hosted_vllm routes directly to the native /v1/responses endpoint when the Responses API config is registered, and correctly parses the response. """ - mock_client = _make_mock_http_client( - _make_mock_responses_api_response("I'm doing well, thanks!") - ) + mock_client = _make_mock_http_client(_make_mock_responses_api_response("I'm doing well, thanks!")) with patch( "litellm.llms.custom_httpx.llm_http_handler._get_httpx_client", @@ -109,10 +107,7 @@ def test_hosted_vllm_responses_create_with_explicit_none_extra_body(): ) # extra_body=None should be normalized to an empty dict (or absent) - assert ( - optional_params.get("extra_body") is not None - or "extra_body" not in optional_params - ) + assert optional_params.get("extra_body") is not None or "extra_body" not in optional_params def test_hosted_vllm_provider_config_registration(): diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py rename to tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py diff --git a/tests/test_litellm/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py b/tests/unit/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py rename to tests/unit/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py diff --git a/tests/test_litellm/llms/huggingface/rerank/test_huggingface_rerank_transformation.py b/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py similarity index 91% rename from tests/test_litellm/llms/huggingface/rerank/test_huggingface_rerank_transformation.py rename to tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py index 9d6b7290eb6..6fd2b006fef 100644 --- a/tests/test_litellm/llms/huggingface/rerank/test_huggingface_rerank_transformation.py +++ b/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py @@ -219,29 +219,6 @@ def test_huggingface_rerank_return_documents(mock_post): assert "text" in result["document"] -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_huggingface_rerank_error_handling(mock_post): - """Test HuggingFace rerank error handling.""" - - def return_val(): - return {"error": "Unauthorized"} - - mock_response = MagicMock() - mock_response.status_code = 401 - mock_response.json = return_val - mock_response.text = "Unauthorized" - mock_post.return_value = mock_response - - with pytest.raises(litellm.APIConnectionError): - litellm.rerank( - model="huggingface/BAAI/bge-reranker-base", - query="hello", - documents=["hello", "world"], - top_n=2, - api_key="invalid_key", - ) - - def test_huggingface_rerank_config(): """Test HuggingFaceRerankConfig class functionality.""" from litellm.llms.huggingface.rerank.transformation import HuggingFaceRerankConfig @@ -249,10 +226,7 @@ def test_huggingface_rerank_config(): config = HuggingFaceRerankConfig() # Test complete URL generation - assert ( - config.get_complete_url(None, "test") - == "https://api-inference.huggingface.co/rerank" - ) + assert config.get_complete_url(None, "test") == "https://api-inference.huggingface.co/rerank" # Test custom API base custom_url = config.get_complete_url("https://custom.huggingface.co", "test") @@ -292,13 +266,9 @@ def test_request_transformation(): config = HuggingFaceRerankConfig() - optional_params = OptionalRerankParams( - query="hello", texts=["hello", "world"], top_n=2, return_text=True - ) + optional_params = OptionalRerankParams(query="hello", texts=["hello", "world"], top_n=2, return_text=True) - request_body = config.transform_rerank_request( - model="test", optional_rerank_params=optional_params, headers={} - ) + request_body = config.transform_rerank_request(model="test", optional_rerank_params=optional_params, headers={}) assert request_body["query"] == "hello" assert request_body["texts"] == ["hello", "world"] @@ -368,9 +338,7 @@ def test_validate_environment(): # Test headers override custom_headers = {"custom": "header"} - headers = config.validate_environment( - headers=custom_headers, model="test", api_key="test_key" - ) + headers = config.validate_environment(headers=custom_headers, model="test", api_key="test_key") assert "custom" in headers assert headers["custom"] == "header" diff --git a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py b/tests/unit/llms/inception/test_inception_chat_transformation.py similarity index 96% rename from tests/test_litellm/llms/inception/test_inception_chat_transformation.py rename to tests/unit/llms/inception/test_inception_chat_transformation.py index 1d12be2adee..c4c023077fc 100644 --- a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py +++ b/tests/unit/llms/inception/test_inception_chat_transformation.py @@ -188,21 +188,15 @@ def test_inception_does_not_leak_key_to_caller_api_base(): caller also supplies their own key. """ config = InceptionChatConfig() - with mock.patch.dict( - os.environ, {"INCEPTION_API_KEY": "server-secret"}, clear=True - ): + with mock.patch.dict(os.environ, {"INCEPTION_API_KEY": "server-secret"}, clear=True): with mock.patch.object(litellm, "inception_key", "module-secret"): # caller overrides api_base without a key -> server key withheld - api_base, api_key = config._get_openai_compatible_provider_info( - "https://attacker.example/v1", None - ) + api_base, api_key = config._get_openai_compatible_provider_info("https://attacker.example/v1", None) assert api_base == "https://attacker.example/v1" assert api_key is None # caller overrides api_base AND supplies their own key -> used as-is - _, api_key = config._get_openai_compatible_provider_info( - "https://attacker.example/v1", "caller-key" - ) + _, api_key = config._get_openai_compatible_provider_info("https://attacker.example/v1", "caller-key") assert api_key == "caller-key" # default/server base -> server-managed key resolved @@ -217,9 +211,7 @@ def test_get_llm_provider_inception(): assert model == "mercury-2" assert provider == "inception" - model, provider, _, api_base = get_llm_provider( - "mercury-2", api_base="https://api.inceptionlabs.ai/v1" - ) + model, provider, _, api_base = get_llm_provider("mercury-2", api_base="https://api.inceptionlabs.ai/v1") assert model == "mercury-2" assert provider == "inception" assert api_base == "https://api.inceptionlabs.ai/v1" @@ -293,5 +285,3 @@ def test_inception_completion_targets_inception_endpoint(): assert captured["body"]["model"] == "mercury-2" assert captured["body"]["tool_choice"] == "auto" assert response.choices[0].message.content == "hi" - - diff --git a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py b/tests/unit/llms/inception/test_inception_completion_transformation.py similarity index 95% rename from tests/test_litellm/llms/inception/test_inception_completion_transformation.py rename to tests/unit/llms/inception/test_inception_completion_transformation.py index ed3f34fc744..84923229e20 100644 --- a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py +++ b/tests/unit/llms/inception/test_inception_completion_transformation.py @@ -22,9 +22,7 @@ def _fim_response_bytes(): "object": "text_completion", "created": 1, "model": "mercury-edit-2", - "choices": [ - {"text": "a + b", "index": 0, "finish_reason": "stop", "logprobs": None} - ], + "choices": [{"text": "a + b", "index": 0, "finish_reason": "stop", "logprobs": None}], "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, } ).encode() @@ -47,9 +45,7 @@ def test_inception_fim_supports_suffix_param(): def test_inception_fim_supported_params_match_schema(): """FIM exposes the OpenAI subset of Inception's FIMCompletionRequest only""" - params = InceptionTextCompletionConfig().get_supported_openai_params( - "mercury-edit-2" - ) + params = InceptionTextCompletionConfig().get_supported_openai_params("mercury-edit-2") for p in ("suffix", "top_p", "frequency_penalty", "presence_penalty", "stop"): assert p in params # Chat-only sampling controls are not part of Inception's FIM schema @@ -75,11 +71,7 @@ def test_inception_get_supported_openai_params_dispatch(): @pytest.mark.parametrize("provider", ["inception", "text-completion-inception"]) def test_inception_validate_environment(provider): - model = ( - "inception/mercury-2" - if provider == "inception" - else "text-completion-inception/mercury-edit-2" - ) + model = "inception/mercury-2" if provider == "inception" else "text-completion-inception/mercury-edit-2" with mock.patch.dict(os.environ, {}, clear=True): result = litellm.validate_environment(model) @@ -217,9 +209,7 @@ def test_inception_fim_does_not_leak_global_api_key(): content=_fim_response_bytes(), ) - with mock.patch.dict( - os.environ, {"INCEPTION_API_KEY": "sk-inception-correct"}, clear=True - ): + with mock.patch.dict(os.environ, {"INCEPTION_API_KEY": "sk-inception-correct"}, clear=True): with mock.patch.object(litellm, "inception_key", None): with mock.patch.object(litellm, "api_key", "sk-global-should-not-leak"): with mock.patch("httpx.Client.send", new=fake_send): diff --git a/tests/test_litellm/llms/jina_ai/embedding/test_jina_embedding_transformation.py b/tests/unit/llms/jina_ai/embedding/test_jina_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/jina_ai/embedding/test_jina_embedding_transformation.py rename to tests/unit/llms/jina_ai/embedding/test_jina_embedding_transformation.py diff --git a/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py b/tests/unit/llms/langflow/chat/test_langflow_chat_transformation.py similarity index 93% rename from tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py rename to tests/unit/llms/langflow/chat/test_langflow_chat_transformation.py index 383a7afbe93..179a6cad4aa 100644 --- a/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py +++ b/tests/unit/llms/langflow/chat/test_langflow_chat_transformation.py @@ -46,7 +46,7 @@ def test_langflow_config_get_complete_url(): def test_langflow_config_get_complete_url_requires_api_base(): config = LangFlowConfig() - with pytest.raises(ValueError, match='api_base is required for LangFlow\\. Set it via'): + with pytest.raises(ValueError, match="api_base is required for LangFlow\\. Set it via"): config.get_complete_url( api_base=None, api_key=None, @@ -225,9 +225,7 @@ def test_langflow_extra_body_cannot_inject_tweaks_into_run_payload(): posted_bodies.append(json.loads(body) if isinstance(body, str) else body) resp = MagicMock(spec=httpx.Response) resp.status_code = 200 - resp.json.return_value = { - "outputs": [{"outputs": [{"results": {"message": {"text": "hi"}}}]}] - } + resp.json.return_value = {"outputs": [{"outputs": [{"results": {"message": {"text": "hi"}}}]}]} resp.headers = {} resp.text = "{}" return resp @@ -275,9 +273,7 @@ def test_langflow_config_extract_response_from_outputs_dict(): "outputs": [ { "results": {}, - "outputs": { - "message": {"message": {"text": "via outputs dict"}} - }, + "outputs": {"message": {"message": {"text": "via outputs dict"}}}, } ] } @@ -292,14 +288,9 @@ def test_langflow_extract_response_returns_none_when_no_message(): assert config._extract_content_from_response({"outputs": []}) is None assert config._extract_content_from_response({"detail": "flow failed"}) is None assert config._extract_content_from_response({"outputs": ["not-a-dict"]}) is None + assert config._extract_content_from_response({"outputs": [{"outputs": ["bad"]}]}) is None assert ( - config._extract_content_from_response({"outputs": [{"outputs": ["bad"]}]}) - is None - ) - assert ( - config._extract_content_from_response( - {"outputs": [{"outputs": [{"results": {"message": {"text": ""}}}]}]} - ) + config._extract_content_from_response({"outputs": [{"outputs": [{"results": {"message": {"text": ""}}}]}]}) is None ) @@ -310,9 +301,7 @@ def test_langflow_transform_response_builds_model_response_with_usage(): status_code=200, json={ "session_id": "sess-abc", - "outputs": [ - {"outputs": [{"results": {"message": {"text": "Hello from LangFlow"}}}]} - ], + "outputs": [{"outputs": [{"results": {"message": {"text": "Hello from LangFlow"}}}]}], }, ) @@ -332,9 +321,7 @@ def test_langflow_transform_response_builds_model_response_with_usage(): assert result.choices[0].finish_reason == "stop" assert result.model == "langflow/my-flow-id" assert result.usage.completion_tokens > 0 - assert result.usage.total_tokens == ( - result.usage.prompt_tokens + result.usage.completion_tokens - ) + assert result.usage.total_tokens == (result.usage.prompt_tokens + result.usage.completion_tokens) def test_langflow_transform_response_raises_on_unparseable_body(): @@ -357,9 +344,7 @@ def test_langflow_transform_response_raises_on_unparseable_body(): def test_langflow_transform_response_raises_on_non_json_body(): config = LangFlowConfig() - raw_response = httpx.Response( - status_code=200, content=b"not json", headers={"content-type": "text/plain"} - ) + raw_response = httpx.Response(status_code=200, content=b"not json", headers={"content-type": "text/plain"}) with pytest.raises(LangFlowError): config.transform_response( diff --git a/tests/test_litellm/llms/litellm_proxy/test_sandbox_executor.py b/tests/unit/llms/litellm_proxy/test_sandbox_executor.py similarity index 84% rename from tests/test_litellm/llms/litellm_proxy/test_sandbox_executor.py rename to tests/unit/llms/litellm_proxy/test_sandbox_executor.py index 422e7a3cf4d..e7a03b9231a 100644 --- a/tests/test_litellm/llms/litellm_proxy/test_sandbox_executor.py +++ b/tests/unit/llms/litellm_proxy/test_sandbox_executor.py @@ -55,9 +55,7 @@ def _install_fake_sandbox(monkeypatch, session_cls=_FakeSandboxSession): def test_execute_installs_inline_requirements_file(monkeypatch): _install_fake_sandbox(monkeypatch) executor = SkillsSandboxExecutor() - monkeypatch.setattr( - executor, "_collect_generated_files", lambda *args, **kwargs: [] - ) + monkeypatch.setattr(executor, "_collect_generated_files", lambda *args, **kwargs: []) requirements = "git+https://example.com/repo.git#egg=foo\n-r extra.txt\n-e ./pkg\n" result = executor.execute( @@ -69,22 +67,15 @@ def test_execute_installs_inline_requirements_file(monkeypatch): assert result["success"] is True created_session = _FakeSandboxSession.last_instance - assert created_session.copied_contents[ - "/sandbox/.litellm_requirements.txt" - ] == requirements.encode("utf-8") - assert ( - "pip', 'install', '-r', '.litellm_requirements.txt'" - in created_session.run_calls[0] - ) + assert created_session.copied_contents["/sandbox/.litellm_requirements.txt"] == requirements.encode("utf-8") + assert "pip', 'install', '-r', '.litellm_requirements.txt'" in created_session.run_calls[0] assert "os.chdir('/sandbox')" in created_session.run_calls[1] def test_execute_uses_skill_requirements_txt(monkeypatch): _install_fake_sandbox(monkeypatch) executor = SkillsSandboxExecutor() - monkeypatch.setattr( - executor, "_collect_generated_files", lambda *args, **kwargs: [] - ) + monkeypatch.setattr(executor, "_collect_generated_files", lambda *args, **kwargs: []) result = executor.execute( code="print('hello')", @@ -97,9 +88,7 @@ def test_execute_uses_skill_requirements_txt(monkeypatch): assert result["success"] is True created_session = _FakeSandboxSession.last_instance - copied_paths = { - sandbox_path for _, sandbox_path in created_session.copy_to_runtime_calls - } + copied_paths = {sandbox_path for _, sandbox_path in created_session.copy_to_runtime_calls} assert "/sandbox/requirements.txt" in copied_paths assert "/sandbox/.litellm_requirements.txt" not in copied_paths assert "pip', 'install', '-r', 'requirements.txt'" in created_session.run_calls[0] @@ -118,9 +107,7 @@ def test_execute_returns_install_failure(monkeypatch): _install_fake_sandbox(monkeypatch, session_cls=_FailingSandboxSession) executor = SkillsSandboxExecutor() - monkeypatch.setattr( - executor, "_collect_generated_files", lambda *args, **kwargs: [] - ) + monkeypatch.setattr(executor, "_collect_generated_files", lambda *args, **kwargs: []) result = executor.execute( code="print('hello')", diff --git a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py b/tests/unit/llms/litellm_proxy/test_skills_ownership.py similarity index 88% rename from tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py rename to tests/unit/llms/litellm_proxy/test_skills_ownership.py index e538c50cde8..6caa2da3169 100644 --- a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py +++ b/tests/unit/llms/litellm_proxy/test_skills_ownership.py @@ -37,12 +37,7 @@ def _skill(skill_id: str, created_by: str | None) -> LiteLLM_SkillsTable: def test_should_extract_skill_auth_from_supported_metadata_fields(): auth = UserAPIKeyAuth(user_id="user-1") - assert ( - skills_main._get_user_api_key_auth_from_kwargs( - {"metadata": {"user_api_key_auth": auth}} - ) - is auth - ) + assert skills_main._get_user_api_key_auth_from_kwargs({"metadata": {"user_api_key_auth": auth}}) is auth assert ( skills_main._get_user_api_key_auth_from_kwargs( {"metadata": {}, "litellm_metadata": {"user_api_key_auth": auth}} @@ -122,9 +117,7 @@ def test_should_forward_skill_auth_through_sdk_entrypoints(monkeypatch): == "deleted" ) - assert handler.create_skill_handler.call_args.kwargs["metadata"] == { - "source": "request" - } + assert handler.create_skill_handler.call_args.kwargs["metadata"] == {"source": "request"} assert handler.create_skill_handler.call_args.kwargs["user_api_key_dict"] is auth assert handler.list_skills_handler.call_args.kwargs["user_api_key_dict"] is auth assert handler.get_skill_handler.call_args.kwargs["user_api_key_dict"] is auth @@ -149,9 +142,7 @@ def test_should_build_resource_owner_scopes_for_auth_context(): ] assert resource_ownership.get_primary_resource_owner_scope(auth) == "user-1" assert resource_ownership.user_can_access_resource_owner("team:team-1", auth) - assert resource_ownership.get_resource_owner_scopes( - UserAPIKeyAuth(token="token-hash") - ) == ["key:token-hash"] + assert resource_ownership.get_resource_owner_scopes(UserAPIKeyAuth(token="token-hash")) == ["key:token-hash"] # Identity-less callers get an empty scope set — sharing a sentinel # would collapse every identity-less caller into the same logical # owner, which is a cross-tenant data-access primitive. @@ -165,9 +156,7 @@ def test_should_allow_admin_and_anonymous_resource_owner_paths(): assert resource_ownership.is_proxy_admin(admin) assert resource_ownership.user_can_access_resource_owner(None, admin) assert resource_ownership.user_can_access_resource_owner(None, None) - assert not resource_ownership.user_can_access_resource_owner( - None, UserAPIKeyAuth(user_id="user-1") - ) + assert not resource_ownership.user_can_access_resource_owner(None, UserAPIKeyAuth(user_id="user-1")) @pytest.mark.asyncio @@ -218,9 +207,7 @@ async def test_should_forward_skill_auth_through_transformation_handler(monkeypa async def test_should_store_team_owner_for_keys_without_user_id(monkeypatch): table = AsyncMock() table.create.side_effect = lambda data: _skill(data["skill_id"], data["created_by"]) - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -242,9 +229,7 @@ async def test_should_store_team_owner_for_keys_without_user_id(monkeypatch): async def test_should_store_token_owner_for_keys_without_user_team_or_org(monkeypatch): table = AsyncMock() table.create.side_effect = lambda data: _skill(data["skill_id"], data["created_by"]) - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -268,9 +253,7 @@ async def test_should_reject_skill_create_for_identityless_proxy_auth(monkeypatc sentinel as ``created_by`` would let any two such callers see each other's skills via the resulting shared owner scope.""" table = AsyncMock() - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -291,9 +274,7 @@ async def test_should_reject_skill_create_for_identityless_proxy_auth(monkeypatc async def test_should_filter_list_skills_to_authenticated_owner_scopes(monkeypatch): table = AsyncMock() table.find_many.return_value = [_skill("litellm_skill_owner", "user-1")] - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -318,9 +299,7 @@ async def test_should_filter_list_skills_to_authenticated_owner_scopes(monkeypat async def test_should_hide_skill_from_different_owner(monkeypatch): table = AsyncMock() table.find_unique.return_value = _skill("litellm_skill_other", "user-2") - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -340,9 +319,7 @@ async def test_should_hide_skill_from_different_owner(monkeypatch): async def test_should_hide_unowned_skill_by_default(monkeypatch): table = AsyncMock() table.find_unique.return_value = _skill("litellm_skill_unowned", None) - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -364,9 +341,7 @@ async def test_list_skills_excludes_unowned_for_non_admin(monkeypatch): with ``created_by IS NULL`` are excluded — admin-only.""" table = AsyncMock() table.find_many.return_value = [] - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( LiteLLMSkillsHandler, "_get_prisma_client", @@ -422,9 +397,7 @@ async def test_load_skill_uses_cache_after_first_db_hit(monkeypatch): fake_skill = Mock(created_by="user-1", skill_id="litellm_skill_a") table = AsyncMock() table.find_unique = AsyncMock(return_value=fake_skill) - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( skills_handler.LiteLLMSkillsHandler, "_get_prisma_client", @@ -432,10 +405,7 @@ async def test_load_skill_uses_cache_after_first_db_hit(monkeypatch): ) for _ in range(3): - assert ( - await skills_handler.LiteLLMSkillsHandler._load_skill("litellm_skill_a") - is fake_skill - ) + assert await skills_handler.LiteLLMSkillsHandler._load_skill("litellm_skill_a") is fake_skill assert table.find_unique.await_count == 1 @@ -445,9 +415,7 @@ async def test_load_skill_caches_negative_lookups(monkeypatch): the DB and the caller still sees ``None``.""" table = AsyncMock() table.find_unique = AsyncMock(return_value=None) - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( skills_handler.LiteLLMSkillsHandler, "_get_prisma_client", @@ -466,9 +434,7 @@ async def test_delete_skill_invalidates_cache(monkeypatch): table = AsyncMock() table.find_unique = AsyncMock(return_value=fake_skill) table.delete = AsyncMock() - prisma_client = type( - "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} - )() + prisma_client = type("Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()})() monkeypatch.setattr( skills_handler.LiteLLMSkillsHandler, "_get_prisma_client", @@ -480,12 +446,7 @@ async def test_delete_skill_invalidates_cache(monkeypatch): assert skills_handler._SKILL_CACHE.get_cache("litellm_skill_a") is fake_skill auth = UserAPIKeyAuth(user_id="user-1") - await skills_handler.LiteLLMSkillsHandler.delete_skill( - "litellm_skill_a", user_api_key_dict=auth - ) + await skills_handler.LiteLLMSkillsHandler.delete_skill("litellm_skill_a", user_api_key_dict=auth) # Post-delete, the cache holds the negative sentinel — not the stale row. - assert ( - skills_handler._SKILL_CACHE.get_cache("litellm_skill_a") - == skills_handler._NEGATIVE_SKILL_SENTINEL - ) + assert skills_handler._SKILL_CACHE.get_cache("litellm_skill_a") == skills_handler._NEGATIVE_SKILL_SENTINEL From d47008129c175a068b23c9e0afe306463e39974f Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:08:12 +0000 Subject: [PATCH 15/47] test(llms): migrate phase 7 provider unit tests to tests/unit Move the wave 1 phase 7 batch (fireworks_ai, gemini, gigachat, github_copilot; 20 files) from tests/test_litellm to tests/unit after judging every test function under a behaviour mutation. Seven wiring or mock-echo tests that stayed green are deleted. The fireworks cost calculator tests get a local model_cost save/restore fixture since the tests/unit tree has no shared conftest for it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_fireworks_ai_chat_transformation.py | 0 ...test_fireworks_ai_rerank_transformation.py | 0 ...t_fireworks_ai_responses_transformation.py | 15 -------- .../test_fireworks_ai_cache_pricing.py | 0 .../test_fireworks_ai_common_utils.py | 0 .../test_fireworks_ai_cost_calculator.py | 22 ++++++++--- ...mini_audio_transcription_transformation.py | 0 .../files/test_gemini_files_transformation.py | 0 .../test_google_genai_guardrail_handler.py | 0 .../test_gemini_image_edit_transformation.py | 0 .../test_gemini_realtime_transformation.py | 0 .../test_gemini_video_transformation.py | 0 .../chat/test_gigachat_chat_streaming.py | 0 .../chat/test_gigachat_chat_transformation.py | 28 -------------- .../test_gigachat_embedding_transformation.py | 30 --------------- ...est_gigachat_passthrough_transformation.py | 0 .../llms/gigachat/test_authenticator.py | 0 .../llms/gigachat/test_file_handler.py | 38 ------------------- .../llms/gigachat/test_utils.py | 0 ...github_copilot_embedding_transformation.py | 0 20 files changed, 16 insertions(+), 117 deletions(-) rename tests/{test_litellm => unit}/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py (100%) rename tests/{test_litellm => unit}/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py (97%) rename tests/{test_litellm => unit}/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py (100%) rename tests/{test_litellm => unit}/llms/fireworks_ai/test_fireworks_ai_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py (90%) rename tests/{test_litellm => unit}/llms/gemini/audio_transcription/test_gemini_audio_transcription_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gemini/files/test_gemini_files_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/gemini/image_edit/test_gemini_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gemini/realtime/test_gemini_realtime_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gemini/videos/test_gemini_video_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gigachat/chat/test_gigachat_chat_streaming.py (100%) rename tests/{test_litellm => unit}/llms/gigachat/chat/test_gigachat_chat_transformation.py (95%) rename tests/{test_litellm => unit}/llms/gigachat/embedding/test_gigachat_embedding_transformation.py (91%) rename tests/{test_litellm => unit}/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py (100%) rename tests/{test_litellm => unit}/llms/gigachat/test_authenticator.py (100%) rename tests/{test_litellm => unit}/llms/gigachat/test_file_handler.py (91%) rename tests/{test_litellm => unit}/llms/gigachat/test_utils.py (100%) rename tests/{test_litellm => unit}/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py (100%) diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py rename to tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py diff --git a/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py b/tests/unit/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py rename to tests/unit/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py diff --git a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py similarity index 97% rename from tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py rename to tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py index d0697ca9b0e..05e3812152e 100644 --- a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py +++ b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py @@ -406,21 +406,6 @@ def test_responses_call_sends_session_affinity_for_caller_session_id() -> None: assert headers["x-session-affinity"] == "sess-42" -def test_responses_call_keeps_caller_supplied_session_affinity_header() -> None: - client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) - pinned: Final[Mapping[str, str]] = MappingProxyType({"x-session-affinity": "explicit-node"}) - with patch(HTTPX_CLIENT_FACTORY, return_value=client): - litellm.responses( - model="fireworks_ai/kimi-k3", - input="hi", - api_key="fw-test-key", - litellm_session_id="sess-42", - extra_headers=pinned, - ) - _, headers, _ = _sent_request(client) - assert headers["x-session-affinity"] == "explicit-node" - - def test_responses_call_maps_provider_errors_to_fireworks_ai() -> None: client: Final = MagicMock() request: Final = httpx.Request("POST", FIREWORKS_RESPONSES_URL) diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py b/tests/unit/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py similarity index 100% rename from tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py rename to tests/unit/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py b/tests/unit/llms/fireworks_ai/test_fireworks_ai_common_utils.py similarity index 100% rename from tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py rename to tests/unit/llms/fireworks_ai/test_fireworks_ai_common_utils.py diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py b/tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py similarity index 90% rename from tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py rename to tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py index 52222f22a51..c6096ba2745 100644 --- a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py +++ b/tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py @@ -1,4 +1,5 @@ import math +from collections.abc import Generator from datetime import datetime, timezone from typing import Final @@ -24,6 +25,15 @@ CACHE_READ_COST = litellm.get_model_info(model=MODEL, custom_llm_provider="firew OUTPUT_COST = 4.4e-06 +@pytest.fixture(autouse=True) +def restore_model_cost() -> Generator[None, None, None]: + original: Final = litellm.model_cost + litellm.get_model_info.cache_clear() + yield + litellm.model_cost = original + litellm.get_model_info.cache_clear() + + def _usage(prompt_tokens: int, cached_tokens: int, completion_tokens: int) -> Usage: return Usage( prompt_tokens=prompt_tokens, @@ -57,7 +67,7 @@ def _register_off_peak_model( cache_read_cost: float | None = STANDARD_CACHE_READ_COST, model: str = OFF_PEAK_MODEL, ) -> None: - litellm.model_cost = { # test-quality-ok: conftest restores litellm.model_cost after each test + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test **litellm.model_cost, f"fireworks_ai/{model}": { "litellm_provider": "fireworks_ai", @@ -151,7 +161,7 @@ def test_an_entry_without_a_cache_read_rate_bills_cached_tokens_at_the_documente """Fireworks documents a default 50% cached-token discount for serverless models: https://docs.fireworks.ai/guides/prompt-caching, accessed 2026-09-19.""" model = "accounts/fireworks/models/default-cache-read-test" - litellm.model_cost = { # test-quality-ok: conftest restores litellm.model_cost after each test + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test **litellm.model_cost, f"fireworks_ai/{model}": { "litellm_provider": "fireworks_ai", @@ -171,7 +181,7 @@ def test_an_entry_without_a_cache_read_rate_bills_cached_tokens_at_the_documente def test_fireworks_cache_read_rates_match_breakdown_and_caching_savings(): model = "accounts/fireworks/models/breakdown-cache-read-test" - litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing **litellm.model_cost, f"fireworks_ai/{model}": { "litellm_provider": "fireworks_ai", @@ -204,7 +214,7 @@ def test_fireworks_cache_read_rates_match_breakdown_and_caching_savings(): def test_generic_cost_per_token_applies_fireworks_cache_read_default_with_or_without_model_info(): model = "accounts/fireworks/models/generic-cache-read-test" - litellm.model_cost = { # test-quality-ok: conftest restores litellm.model_cost after each test + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test **litellm.model_cost, f"fireworks_ai/{model}": { "litellm_provider": "fireworks_ai", @@ -257,7 +267,7 @@ COMPONENT_AUDIO_OUT_COST = 6e-06 def test_cache_write_reasoning_and_audio_tokens_are_billed_at_their_component_rates(): - litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing **litellm.model_cost, f"fireworks_ai/{COMPONENT_MODEL}": { "litellm_provider": "fireworks_ai", @@ -302,7 +312,7 @@ def test_cache_write_reasoning_and_audio_tokens_are_billed_at_their_component_ra def test_an_entry_without_an_input_rate_gets_no_cache_read_fallback(): - litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing + litellm.model_cost = { # test-quality-ok: the restore_model_cost fixture returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing **litellm.model_cost, # pyright: ignore[reportUnknownMemberType] # the SDK types model_cost as dict[Unknown, Unknown] "fireworks_ai/accounts/fireworks/models/no-input-rate-test": { "litellm_provider": "fireworks_ai", diff --git a/tests/test_litellm/llms/gemini/audio_transcription/test_gemini_audio_transcription_transformation.py b/tests/unit/llms/gemini/audio_transcription/test_gemini_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/audio_transcription/test_gemini_audio_transcription_transformation.py rename to tests/unit/llms/gemini/audio_transcription/test_gemini_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py b/tests/unit/llms/gemini/files/test_gemini_files_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py rename to tests/unit/llms/gemini/files/test_gemini_files_transformation.py diff --git a/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py b/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py rename to tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py diff --git a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py b/tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py rename to tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py diff --git a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py rename to tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py diff --git a/tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py b/tests/unit/llms/gemini/videos/test_gemini_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py rename to tests/unit/llms/gemini/videos/test_gemini_video_transformation.py diff --git a/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_streaming.py b/tests/unit/llms/gigachat/chat/test_gigachat_chat_streaming.py similarity index 100% rename from tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_streaming.py rename to tests/unit/llms/gigachat/chat/test_gigachat_chat_streaming.py diff --git a/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py b/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py similarity index 95% rename from tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py rename to tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py index 2f9511e642c..8e84072e549 100644 --- a/tests/test_litellm/llms/gigachat/chat/test_gigachat_chat_transformation.py +++ b/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py @@ -141,22 +141,6 @@ class TestValidateEnvironment: assert self.config._current_credentials == "my-creds" assert self.config._current_api_base == "https://my-api.example.com" - @patch(f"{TRANSFORM_MODULE}.get_access_token", return_value="token") - @patch(f"{TRANSFORM_MODULE}.get_secret_str") - def test_falls_back_to_env_for_credentials( # test-quality-ok: mock-echo of internal wiring - self, mock_get_secret, mock_get_token - ): - mock_get_secret.return_value = "env-creds" - self.config.validate_environment( - headers={}, - model="GigaChat", - messages=[], - optional_params={}, - litellm_params={}, - api_key=None, - api_base=None, - ) - mock_get_secret.assert_any_call("GIGACHAT_CREDENTIALS") # test-quality-ok: mock-echo of internal wiring class TestGetSupportedOpenAiParams: @@ -865,18 +849,6 @@ class TestUploadImage: def setup_method(self): self.config = GigaChatConfig() - @patch(f"{TRANSFORM_MODULE}.upload_file_sync", return_value="file-uploaded") - def test_upload_image_success(self, mock_upload): - self.config._current_credentials = "creds" - self.config._current_api_base = "https://api.example.com" - result = self.config._upload_image("https://example.com/img.jpg") - assert result == "file-uploaded" - mock_upload.assert_called_once_with( - image_url="https://example.com/img.jpg", - credentials="creds", - api_base="https://api.example.com", - ) - @patch(f"{TRANSFORM_MODULE}.upload_file_sync", side_effect=Exception("fail")) def test_upload_image_failure_returns_none(self, mock_upload): result = self.config._upload_image("https://example.com/img.jpg") diff --git a/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py b/tests/unit/llms/gigachat/embedding/test_gigachat_embedding_transformation.py similarity index 91% rename from tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py rename to tests/unit/llms/gigachat/embedding/test_gigachat_embedding_transformation.py index 8537793ea72..01fe66ca4c7 100644 --- a/tests/test_litellm/llms/gigachat/embedding/test_gigachat_embedding_transformation.py +++ b/tests/unit/llms/gigachat/embedding/test_gigachat_embedding_transformation.py @@ -37,17 +37,6 @@ def _make_httpx_response(body: dict, status_code: int = 200) -> httpx.Response: # --------------------------------------------------------------------------- -class TestGetConfig: - def setup_method(self): - self.config = GigaChatEmbeddingConfig() - - def test_contains_only_abc_impl(self): - """get_config returns ABC internal data due to inheritance.""" - result = self.config.get_config() - # The only key should be _abc_impl from ABC base class - assert set(result.keys()) == {"_abc_impl"} - - class TestGetSupportedOpenAiParams: def setup_method(self): self.config = GigaChatEmbeddingConfig() @@ -287,25 +276,6 @@ class TestTransformEmbeddingResponse: ) assert result.model == "Embeddings" - def test_calls_logging_post_call(self): - raw = self._make_gigachat_response([ - {"object": "embedding", "embedding": [0.1], "index": 0}, - ]) - model_response = EmbeddingResponse() - self.config.transform_embedding_response( - model="gigachat/Embeddings", - raw_response=raw, - model_response=model_response, - logging_obj=self.logging_obj, - api_key="test-api-key", - request_data={"input": ["hello"]}, - optional_params={}, - litellm_params={}, - ) - self.logging_obj.post_call.assert_called_once() - args = self.logging_obj.post_call.call_args.kwargs - assert args["api_key"] == "test-api-key" - assert args["input"] == ["hello"] class TestValidateEnvironment: diff --git a/tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py b/tests/unit/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py rename to tests/unit/llms/gigachat/passthrough/test_gigachat_passthrough_transformation.py diff --git a/tests/test_litellm/llms/gigachat/test_authenticator.py b/tests/unit/llms/gigachat/test_authenticator.py similarity index 100% rename from tests/test_litellm/llms/gigachat/test_authenticator.py rename to tests/unit/llms/gigachat/test_authenticator.py diff --git a/tests/test_litellm/llms/gigachat/test_file_handler.py b/tests/unit/llms/gigachat/test_file_handler.py similarity index 91% rename from tests/test_litellm/llms/gigachat/test_file_handler.py rename to tests/unit/llms/gigachat/test_file_handler.py index ce9505f11f2..de83b2ddf5f 100644 --- a/tests/test_litellm/llms/gigachat/test_file_handler.py +++ b/tests/unit/llms/gigachat/test_file_handler.py @@ -344,25 +344,6 @@ class TestUploadFileSync: assert result is None - @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") - @patch(f"{FILE_MODULE}.get_access_token", return_value="test-token") - @patch(f"{FILE_MODULE}._get_httpx_client") - def test_uploads_without_optional_args( - self, mock_http_handler_cls, mock_get_token, mock_get_api_base - ): - """Verify that credentials, api_base, and litellm_params are optional.""" - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.json.return_value = {"id": "file-no-args"} - mock_response.raise_for_status = MagicMock() - mock_client.post.return_value = mock_response - mock_http_handler_cls.return_value = mock_client - - result = upload_file_sync(image_url=_RED_PNG_DATA_URL) - - assert result == "file-no-args" - # Should still have called get_access_token without args - mock_get_token.assert_called_once_with(credentials=None, litellm_params=None) # --------------------------------------------------------------------------- @@ -483,22 +464,3 @@ class TestUploadFileAsync: ) assert result is None - - @pytest.mark.asyncio - @patch(f"{FILE_MODULE}.get_api_base", return_value="https://api.example.com") - @patch(f"{FILE_MODULE}.get_access_token_async", return_value="test-token-async") - @patch(f"{FILE_MODULE}.get_async_httpx_client") - async def test_uploads_without_optional_args( - self, mock_get_client, mock_get_token, mock_get_api_base - ): - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.json = MagicMock(return_value={"id": "async-no-args"}) - mock_response.raise_for_status = MagicMock() - mock_client.post = AsyncMock(return_value=mock_response) - mock_get_client.return_value = mock_client - - result = await upload_file_async(image_url=_RED_PNG_DATA_URL) - - assert result == "async-no-args" - mock_get_token.assert_called_once_with(credentials=None, litellm_params=None) \ No newline at end of file diff --git a/tests/test_litellm/llms/gigachat/test_utils.py b/tests/unit/llms/gigachat/test_utils.py similarity index 100% rename from tests/test_litellm/llms/gigachat/test_utils.py rename to tests/unit/llms/gigachat/test_utils.py diff --git a/tests/test_litellm/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py b/tests/unit/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py rename to tests/unit/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py From feca00248cf78fda0da5a1fa5b7f1a164fbe5608 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:11:56 +0000 Subject: [PATCH 16/47] test(bedrock): isolate host AWS config in realtime and rerank unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock/realtime/test_bedrock_realtime_handler.py | 9 +++++++++ .../rerank/test_bedrock_rerank_header_forwarding.py | 9 +++++++++ 2 files changed, 18 insertions(+) diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index a7f0f64ef68..73a78a94e9f 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -18,6 +18,15 @@ from litellm.llms.bedrock.realtime.handler import BedrockRealtime from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig +@pytest.fixture(autouse=True) +def _isolate_host_aws_config(monkeypatch, tmp_path): + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "credentials")) + monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config")) + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + for env_var in ("AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION_NAME", "AWS_DEFAULT_REGION"): + monkeypatch.delenv(env_var, raising=False) + + class FakePayloadPart: def __init__(self, bytes_): self.bytes_ = bytes_ diff --git a/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index 2ea61b5e978..aa93ddb21b8 100644 --- a/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -15,6 +15,15 @@ from litellm.llms.bedrock.base_aws_llm import Boto3CredentialsInfo from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + +@pytest.fixture(autouse=True) +def _isolate_host_aws_config(monkeypatch, tmp_path): + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "credentials")) + monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config")) + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + for env_var in ("AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION_NAME", "AWS_DEFAULT_REGION"): + monkeypatch.delenv(env_var, raising=False) + # Mock response for Bedrock rerank # Format based on Bedrock rerank API response structure bedrock_rerank_response = { From fcabb626acdcf240c43ba78dfb92d7822e239cfe Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:19:05 +0000 Subject: [PATCH 17/47] test(unit): migrate wave 1 phase 3 anthropic, apiserpent, azure and azure_ai legacy tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 11 +++ .../test_reasoning_effort_fields.py | 0 .../test_anthropic_files_transformation.py | 0 .../messages/test_advisor_orchestration.py | 0 .../llms/apiserpent/test_apiserpent_search.py | 0 .../test_azure_image_edit_transformation.py | 0 .../test_azure_image_generation_init.py | 82 ------------------- .../test_azure_passthrough_transformation.py | 0 .../realtime/test_azure_realtime_handler.py | 35 -------- .../response/test_azure_transformation.py | 0 .../foundry_responses_web_search_fixture.json | 0 ...st_bing_grounding_search_transformation.py | 0 .../test_azure_tts_transformation.py | 0 ...test_azure_vector_stores_transformation.py | 0 .../chat/test_azure_ai_transformation.py | 15 ---- .../embed/test_azure_ai_embed_handler.py | 0 ...test_azure_ai_image_edit_transformation.py | 0 .../test_mai_image_edit_transformation.py | 0 ...st_azure_ai_cohere_parse_transformation.py | 0 ...est_azure_ai_passthrough_transformation.py | 0 .../test_azure_ai_rerank_transformation.py | 0 .../test_azure_ai_responses_transformation.py | 0 22 files changed, 11 insertions(+), 132 deletions(-) rename tests/{test_litellm => unit}/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/files/test_anthropic_files_transformation.py (100%) rename tests/{test_litellm => unit}/llms/anthropic/messages/test_advisor_orchestration.py (100%) rename tests/{test_litellm => unit}/llms/apiserpent/test_apiserpent_search.py (100%) rename tests/{test_litellm => unit}/llms/azure/image_edit/test_azure_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/image_generation/test_azure_image_generation_init.py (91%) rename tests/{test_litellm => unit}/llms/azure/passthrough/test_azure_passthrough_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/realtime/test_azure_realtime_handler.py (94%) rename tests/{test_litellm => unit}/llms/azure/response/test_azure_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/search/foundry_responses_web_search_fixture.json (100%) rename tests/{test_litellm => unit}/llms/azure/search/test_bing_grounding_search_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/text_to_speech/test_azure_tts_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure/vector_stores/test_azure_vector_stores_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/chat/test_azure_ai_transformation.py (97%) rename tests/{test_litellm => unit}/llms/azure_ai/embed/test_azure_ai_embed_handler.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py (100%) rename tests/{test_litellm => unit}/llms/azure_ai/responses/test_azure_ai_responses_transformation.py (100%) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 3bdab1d231a..d002452d14c 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -2,6 +2,8 @@ from collections.abc import Iterator from typing import Final import pytest + +import litellm from pytest_socket import enable_socket, socket_allow_hosts LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] @@ -21,3 +23,12 @@ def block_external_sockets() -> Iterator[None]: @pytest.hookimpl(trylast=True) def pytest_runtest_setup() -> None: _allow_loopback_only() + + +@pytest.fixture +def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py b/tests/unit/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py rename to tests/unit/llms/anthropic/experimental_pass_through/test_reasoning_effort_fields.py diff --git a/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py b/tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py similarity index 100% rename from tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py rename to tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py diff --git a/tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py b/tests/unit/llms/anthropic/messages/test_advisor_orchestration.py similarity index 100% rename from tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py rename to tests/unit/llms/anthropic/messages/test_advisor_orchestration.py diff --git a/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py b/tests/unit/llms/apiserpent/test_apiserpent_search.py similarity index 100% rename from tests/test_litellm/llms/apiserpent/test_apiserpent_search.py rename to tests/unit/llms/apiserpent/test_apiserpent_search.py diff --git a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py b/tests/unit/llms/azure/image_edit/test_azure_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py rename to tests/unit/llms/azure/image_edit/test_azure_image_edit_transformation.py diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py similarity index 91% rename from tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py rename to tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py index cfde1760389..eabd5c8427d 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py @@ -133,88 +133,6 @@ def test_azure_image_generation_flattens_extra_body(): assert data["size"] == "1024x1024" -def test_azure_image_generation_creates_token_provider_from_credentials(): - """ - Test that azure_ad_token_provider is created from tenant_id, client_id, client_secret. - - This test verifies the fix in images/main.py where we now create the - azure_ad_token_provider from credentials in litellm_params if it's not already provided. - """ - # Simulate the fix in images/main.py - litellm_params_dict = { - "tenant_id": "test-tenant-id", - "client_id": "test-client-id", - "client_secret": "test-client-secret", - "azure_scope": None, - } - - azure_ad_token_provider = None - - # This is the logic we added in images/main.py - if azure_ad_token_provider is None: - tenant_id = litellm_params_dict.get("tenant_id") - client_id = litellm_params_dict.get("client_id") - client_secret = litellm_params_dict.get("client_secret") - azure_scope = ( - litellm_params_dict.get("azure_scope") - or "https://cognitiveservices.azure.com/.default" - ) - - # Verify the credentials are extracted correctly - assert tenant_id == "test-tenant-id" - assert client_id == "test-client-id" - assert client_secret == "test-client-secret" - assert azure_scope == "https://cognitiveservices.azure.com/.default" - - # Verify the condition to create token provider is met - assert ( - tenant_id and client_id and client_secret - ), "Credentials should be present to create token provider" - - -def test_azure_image_generation_headers_without_api_key(): - """ - Test that when api_key is None, the api-key header is not added to headers. - - This prevents the httpx TypeError: "Header value must be str or bytes, not " - that was occurring when api_key was None and being set in headers. - - This is a unit test for the fix in images/main.py where we now check: - if api_key is not None: - default_headers["api-key"] = api_key - """ - from litellm.images.main import image_generation - - # Test the header building logic directly - api_key = None - - default_headers = { - "Content-Type": "application/json", - } - - # This is the fix: only add api-key if it's not None - if api_key is not None: - default_headers["api-key"] = api_key - - # Verify api-key is not in headers when api_key is None - assert "api-key" not in default_headers - - # Verify Content-Type is still there - assert default_headers["Content-Type"] == "application/json" - - # Test with a valid api_key - api_key = "valid-key-123" - default_headers_with_key = { - "Content-Type": "application/json", - } - if api_key is not None: - default_headers_with_key["api-key"] = api_key - - # Verify api-key is added when api_key is valid - assert "api-key" in default_headers_with_key - assert default_headers_with_key["api-key"] == "valid-key-123" - - def test_azure_image_generation_drop_params_response_format(): """ Test that unsupported params like response_format are dropped when drop_params=True. diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/unit/llms/azure/passthrough/test_azure_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py rename to tests/unit/llms/azure/passthrough/test_azure_passthrough_transformation.py diff --git a/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py similarity index 94% rename from tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py rename to tests/unit/llms/azure/realtime/test_azure_realtime_handler.py index 7d24e604569..c1ba286f8c0 100644 --- a/tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py +++ b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py @@ -426,41 +426,6 @@ async def test_async_realtime_beta_without_api_version_raises(): ) -@pytest.mark.asyncio -async def test_realtime_protocol_env_var_fallback(): - """ - Test that LITELLM_AZURE_REALTIME_PROTOCOL env var is used as fallback. - Fixes #22127: no way to set realtime_protocol from config. - """ - from litellm.realtime_api.main import _arealtime - from litellm.types.router import GenericLiteLLMParams - - with patch.dict(os.environ, {"LITELLM_AZURE_REALTIME_PROTOCOL": "v1"}): - # Create a GenericLiteLLMParams without realtime_protocol - litellm_params = GenericLiteLLMParams() - # The env var should be picked up as fallback - realtime_protocol = ( - {}.get("realtime_protocol") - or litellm_params.get("realtime_protocol") - or os.environ.get("LITELLM_AZURE_REALTIME_PROTOCOL") - or "beta" - ) - assert realtime_protocol == "v1" - - -@pytest.mark.asyncio -async def test_realtime_protocol_from_litellm_params(): - """ - Test that realtime_protocol is read from litellm_params (config.yaml extra field). - Fixes #22127: realtime_protocol in litellm_params was not used. - """ - from litellm.types.router import GenericLiteLLMParams - - # Simulate config.yaml with realtime_protocol as an extra field - litellm_params = GenericLiteLLMParams(realtime_protocol="GA") - assert litellm_params.get("realtime_protocol") == "GA" - - @pytest.mark.asyncio async def test_arealtime_transcription_intent_defaults_to_ga(monkeypatch): """ diff --git a/tests/test_litellm/llms/azure/response/test_azure_transformation.py b/tests/unit/llms/azure/response/test_azure_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/response/test_azure_transformation.py rename to tests/unit/llms/azure/response/test_azure_transformation.py diff --git a/tests/test_litellm/llms/azure/search/foundry_responses_web_search_fixture.json b/tests/unit/llms/azure/search/foundry_responses_web_search_fixture.json similarity index 100% rename from tests/test_litellm/llms/azure/search/foundry_responses_web_search_fixture.json rename to tests/unit/llms/azure/search/foundry_responses_web_search_fixture.json diff --git a/tests/test_litellm/llms/azure/search/test_bing_grounding_search_transformation.py b/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/search/test_bing_grounding_search_transformation.py rename to tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py diff --git a/tests/test_litellm/llms/azure/text_to_speech/test_azure_tts_transformation.py b/tests/unit/llms/azure/text_to_speech/test_azure_tts_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/text_to_speech/test_azure_tts_transformation.py rename to tests/unit/llms/azure/text_to_speech/test_azure_tts_transformation.py diff --git a/tests/test_litellm/llms/azure/vector_stores/test_azure_vector_stores_transformation.py b/tests/unit/llms/azure/vector_stores/test_azure_vector_stores_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/vector_stores/test_azure_vector_stores_transformation.py rename to tests/unit/llms/azure/vector_stores/test_azure_vector_stores_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py similarity index 97% rename from tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py rename to tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py index f8cc0b5071e..e4a33d5772c 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -352,21 +352,6 @@ def test_azure_model_router_stamps_selected_model_on_hidden_params(): ) -def test_azure_model_router_stamp_does_not_leak_across_responses(): - """ - ModelResponse declares _hidden_params as a class-level dict, so the stamp has to be written - as a fresh dict. Mutating in place would bleed the selected model into unrelated responses. - """ - from litellm.llms.azure_ai.common_utils import ( - AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, - ) - from litellm.types.utils import ModelResponse - - untouched = ModelResponse() - - assert AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY not in (untouched._hidden_params or {}) - - def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name(): """ Regression test: Azure AI returns 400 when tools contain copilot_mcp_server_name. diff --git a/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py b/tests/unit/llms/azure_ai/embed/test_azure_ai_embed_handler.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py rename to tests/unit/llms/azure_ai/embed/test_azure_ai_embed_handler.py diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py rename to tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py rename to tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py b/tests/unit/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py rename to tests/unit/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py rename to tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py b/tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py rename to tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py b/tests/unit/llms/azure_ai/responses/test_azure_ai_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py rename to tests/unit/llms/azure_ai/responses/test_azure_ai_responses_transformation.py From cf2a9b372cffe0e00e44ee252168384c8f5c058d Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:21:51 +0000 Subject: [PATCH 18/47] test(gigachat): cover env credential fallback by its resulting auth header Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat/test_gigachat_chat_transformation.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py b/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py index 8e84072e549..b1307f56336 100644 --- a/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py +++ b/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py @@ -141,6 +141,25 @@ class TestValidateEnvironment: assert self.config._current_credentials == "my-creds" assert self.config._current_api_base == "https://my-api.example.com" + @patch( + f"{TRANSFORM_MODULE}.get_access_token", + side_effect=lambda credentials, litellm_params: f"token-for-{credentials}", + ) + def test_falls_back_to_env_credentials_when_api_key_missing( + self, mock_get_token, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setenv("GIGACHAT_CREDENTIALS", "env-creds") + result = self.config.validate_environment( + headers={}, + model="GigaChat", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + assert result["Authorization"] == "Bearer token-for-env-creds" + assert self.config._current_credentials == "env-creds" class TestGetSupportedOpenAiParams: From decb28f5b575378cff5e8ce941c999da24df3875 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 08:34:55 +0000 Subject: [PATCH 19/47] test(unit): restore live router and runtime model cost state between unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index d002452d14c..d47f71b21a8 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -2,9 +2,11 @@ from collections.abc import Iterator from typing import Final import pytest +from pytest_socket import enable_socket, socket_allow_hosts import litellm -from pytest_socket import enable_socket, socket_allow_hosts +import litellm.router as litellm_router_module +import litellm.utils as litellm_utils_module LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] @@ -25,6 +27,24 @@ def pytest_runtest_setup() -> None: _allow_loopback_only() +@pytest.fixture(autouse=True) +def isolate_router_model_cost_state() -> Iterator[None]: + original_live_routers: Final = frozenset(litellm_router_module._live_routers) + original_runtime_registered_model_cost: Final = { + model_key: dict(model_value) + for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items() + } + yield + for router in tuple(litellm_router_module._live_routers): + litellm_router_module._live_routers.discard(router) + for router in original_live_routers: + litellm_router_module._live_routers.add(router) + litellm_utils_module._runtime_registered_model_cost.clear() + litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost) + litellm_utils_module._invalidate_model_cost_lowercase_map() + litellm.get_model_info.cache_clear() + + @pytest.fixture def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") From 03a650c8a95802633c792d290a95bb8cbbb1159c Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 09:17:57 +0000 Subject: [PATCH 20/47] test: migrate wave 1 phase 1 legacy tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...t_responses_bridge_provider_propagation.py | 153 ------------------ .../test_pydantic_ai_agent_headers.py | 0 .../test_pydantic_ai_agent_transformation.py | 0 ...test_watsonx_orchestrate_transformation.py | 18 --- .../test_exception_mapping_utils.py | 0 .../batches/test_batch_utils.py | 0 .../batches/test_main.py | 0 .../batches/test_responses_batch_cost.py | 25 +-- .../chat_completions/test_dispatch.py | 13 +- ...responses_transformation_transformation.py | 0 .../compression/test_compress.py | 0 .../test_transformation.py | 0 .../test_callback_controls.py | 0 .../enterprise_callbacks/test_llm_guard.py | 0 .../test_secret_detection.py | 0 .../test_compression_interception_handler.py | 0 .../gcs_bucket/test_gcs_bucket_base.py | 0 .../integrations/gcs_pubsub/test_pub_sub.py | 0 .../helicone/test_helicone_gemini.py | 34 ---- 19 files changed, 14 insertions(+), 229 deletions(-) delete mode 100644 tests/test_litellm/completion_extras/test_responses_bridge_provider_propagation.py rename tests/{test_litellm => unit}/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py (100%) rename tests/{test_litellm => unit}/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py (100%) rename tests/{test_litellm => unit}/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py (95%) rename tests/{test_litellm => unit}/anthropic_interface/exceptions/test_exception_mapping_utils.py (100%) rename tests/{test_litellm => unit}/batches/test_batch_utils.py (100%) rename tests/{test_litellm => unit}/batches/test_main.py (100%) rename tests/{test_litellm => unit}/batches/test_responses_batch_cost.py (87%) rename tests/{test_litellm => unit}/chat_completions/test_dispatch.py (93%) rename tests/{test_litellm => unit}/completion_extras/test_litellm_responses_transformation_transformation.py (100%) rename tests/{test_litellm => unit}/compression/test_compress.py (100%) rename tests/{test_litellm => unit}/endpoints/speech/speech_to_completion_bridge/test_transformation.py (100%) rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/test_callback_controls.py (100%) rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/test_llm_guard.py (100%) rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/test_secret_detection.py (100%) rename tests/{test_litellm => unit}/integrations/compression_interception/test_compression_interception_handler.py (100%) rename tests/{test_litellm => unit}/integrations/gcs_bucket/test_gcs_bucket_base.py (100%) rename tests/{test_litellm => unit}/integrations/gcs_pubsub/test_pub_sub.py (100%) rename tests/{test_litellm => unit}/integrations/helicone/test_helicone_gemini.py (73%) diff --git a/tests/test_litellm/completion_extras/test_responses_bridge_provider_propagation.py b/tests/test_litellm/completion_extras/test_responses_bridge_provider_propagation.py deleted file mode 100644 index 8036c72679e..00000000000 --- a/tests/test_litellm/completion_extras/test_responses_bridge_provider_propagation.py +++ /dev/null @@ -1,153 +0,0 @@ -""" -Regression test for https://github.com/BerriAI/litellm/issues/28505 - -the Responses API bridge double-strips the provider prefix from the -model name when a Chat Completions request has both `tools` and -`reasoning_effort`. - -Root cause: the bridge handler called `litellm.responses()` / -`litellm.aresponses()` without passing the already-resolved -`custom_llm_provider`. The downstream call then re-invoked -`get_llm_provider()` with `custom_llm_provider=None`, which stripped -a second provider prefix from a `provider/provider/model` deployment -string. - -This test pins both the sync and async bridge handler call sites: -the resolved `custom_llm_provider` must be forwarded to the underlying -`responses` / `aresponses` call so the provider isn't re-detected. -""" - -from unittest.mock import MagicMock, patch - -import pytest - -from litellm.completion_extras.litellm_responses_transformation.handler import ( - ResponsesToCompletionBridgeHandler, -) - - -def _validated_kwargs(): - return { - "model": "openai/openai/openai/gpt-5.5", - "messages": [{"role": "user", "content": "hi"}], - "optional_params": {}, - "litellm_params": {}, - "headers": {}, - "model_response": MagicMock(), - "logging_obj": MagicMock(), - "custom_llm_provider": "openai", - } - - -def test_sync_completion_forwards_custom_llm_provider(): - handler = ResponsesToCompletionBridgeHandler() - handler.transformation_handler = MagicMock() - handler.transformation_handler.transform_request.return_value = { - "model": "openai/openai/openai/gpt-5.5", - "input": [], - # `_build_sanitized_litellm_params` spreads `custom_llm_provider` from - # `litellm_params` into request_data on the real bridge path. Seed - # it here so the test exercises the overwrite (not an explicit kwarg - # that would TypeError against an already-present key). - "custom_llm_provider": "should-be-overwritten", - } - handler.transformation_handler.transform_response.return_value = ( - _validated_kwargs()["model_response"] - ) - with ( - patch.object( - handler, "validate_input_kwargs", return_value=_validated_kwargs() - ), - patch( - "litellm.responses", - return_value=MagicMock(spec=[]), - ) as mock_responses, - ): - # The handler routes ResponsesAPIResponse through transform_response. - # We just want to verify the kwargs going INTO responses(). - try: - handler.completion(acompletion=False) - except Exception: - # Downstream handling (transform_response, type checks) is not - # the subject of this test. - pass - assert mock_responses.called - kwargs = mock_responses.call_args.kwargs - assert kwargs.get("custom_llm_provider") == "openai", ( - "sync bridge must forward custom_llm_provider to litellm.responses() " - "so the downstream get_llm_provider() call does not re-strip the " - "provider prefix on a provider/provider/model deployment string" - ) - - -@pytest.mark.asyncio -async def test_async_completion_forwards_custom_llm_provider(): - handler = ResponsesToCompletionBridgeHandler() - handler.transformation_handler = MagicMock() - handler.transformation_handler.transform_request.return_value = { - "model": "openai/openai/openai/gpt-5.5", - "input": [], - # `_build_sanitized_litellm_params` spreads `custom_llm_provider` from - # `litellm_params` into request_data on the real bridge path. Seed - # it here so the test exercises the overwrite (not an explicit kwarg - # that would TypeError against an already-present key). - "custom_llm_provider": "should-be-overwritten", - } - - async def _fake_aresponses(**kwargs): - _fake_aresponses.kwargs = kwargs - return MagicMock(spec=[]) - - _fake_aresponses.kwargs = {} - - with ( - patch.object( - handler, "validate_input_kwargs", return_value=_validated_kwargs() - ), - patch("litellm.aresponses", _fake_aresponses), - ): - try: - await handler.acompletion() - except Exception: - pass - assert _fake_aresponses.kwargs.get("custom_llm_provider") == "openai", ( - "async bridge must forward custom_llm_provider to litellm.aresponses() " - "so the downstream get_llm_provider() call does not re-strip the " - "provider prefix on a provider/provider/model deployment string" - ) - - -@pytest.mark.asyncio -async def test_async_completion_forwards_aws_region_name(): - handler = ResponsesToCompletionBridgeHandler() - handler.transformation_handler = MagicMock() - handler.transformation_handler.transform_request.return_value = { - "model": "openai.gpt-5.5", - "input": [], - "aws_region_name": "us-east-2", - "api_base": "https://bedrock-mantle.us-east-1.api.aws/v1", - "custom_llm_provider": "bedrock_mantle", - } - - async def _fake_aresponses(**kwargs): - _fake_aresponses.kwargs = kwargs - return MagicMock(spec=[]) - - _fake_aresponses.kwargs = {} - - validated = _validated_kwargs() - validated["custom_llm_provider"] = "bedrock_mantle" - validated["litellm_params"] = { - "aws_region_name": "us-east-2", - "api_base": "https://bedrock-mantle.us-east-1.api.aws/v1", - "custom_llm_provider": "bedrock_mantle", - } - - with ( - patch.object(handler, "validate_input_kwargs", return_value=validated), - patch("litellm.aresponses", _fake_aresponses), - ): - try: - await handler.acompletion() - except Exception: - pass - assert _fake_aresponses.kwargs.get("aws_region_name") == "us-east-2" diff --git a/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py b/tests/unit/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py similarity index 100% rename from tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py rename to tests/unit/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py diff --git a/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py b/tests/unit/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py similarity index 100% rename from tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py rename to tests/unit/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py diff --git a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py similarity index 95% rename from tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py rename to tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py index 43dfdaba02d..a300560ae9d 100644 --- a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py +++ b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py @@ -1,7 +1,6 @@ import asyncio import json import time -from pathlib import Path import httpx import pytest @@ -571,20 +570,3 @@ def test_config_manager_returns_wxo_provider(): ) assert config is not None assert config.__class__.__name__ == "WatsonxOrchestrateA2AConfig" - - -def test_wxo_dashboard_auth_fields(): - fields_path = ( - Path(__file__).resolve().parents[5] - / "litellm/proxy/public_endpoints/agent_create_fields.json" - ) - agent_fields = json.loads(fields_path.read_text()) - wxo_agent = next( - agent for agent in agent_fields if agent["agent_type"] == "watsonx_orchestrate" - ) - fields_by_key = {field["key"]: field for field in wxo_agent["credential_fields"]} - - assert fields_by_key["auth_mode"]["default_value"] == "cp4d" - # Username is CP4D-only; UI does not require it so ibm_cloud users are not blocked. - assert fields_by_key["username"]["required"] is False - assert "cp4d" in fields_by_key["username"]["tooltip"].lower() diff --git a/tests/test_litellm/anthropic_interface/exceptions/test_exception_mapping_utils.py b/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py similarity index 100% rename from tests/test_litellm/anthropic_interface/exceptions/test_exception_mapping_utils.py rename to tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py similarity index 100% rename from tests/test_litellm/batches/test_batch_utils.py rename to tests/unit/batches/test_batch_utils.py diff --git a/tests/test_litellm/batches/test_main.py b/tests/unit/batches/test_main.py similarity index 100% rename from tests/test_litellm/batches/test_main.py rename to tests/unit/batches/test_main.py diff --git a/tests/test_litellm/batches/test_responses_batch_cost.py b/tests/unit/batches/test_responses_batch_cost.py similarity index 87% rename from tests/test_litellm/batches/test_responses_batch_cost.py rename to tests/unit/batches/test_responses_batch_cost.py index b634f5f73db..63b28fb3b42 100644 --- a/tests/test_litellm/batches/test_responses_batch_cost.py +++ b/tests/unit/batches/test_responses_batch_cost.py @@ -12,17 +12,28 @@ Line shape decides the parse, not the batch's declared endpoint, so an output file mixing Responses-shaped and chat-shaped lines sums across both. """ -from typing import Literal, get_args, get_type_hints import pytest import litellm import litellm.batches.batch_utils as bu -from litellm.types.llms.openai import CreateBatchRequest MODEL = "gpt-5.6" +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + def _responses_line(input_tokens: int, output_tokens: int) -> dict: return { "response": { @@ -107,13 +118,3 @@ async def test_mixed_shape_batch_output_sums_across_both_line_shapes(local_model assert result.cost == pytest.approx( 133 * model_info["input_cost_per_token_batches"] + 107 * model_info["output_cost_per_token_batches"] ) - - -def test_create_batch_endpoint_accepts_v1_responses(): - """A type-checked caller can pass endpoint="/v1/responses", which the runtime - already forwarded correctly.""" - endpoint_annotation = get_type_hints(CreateBatchRequest)["endpoint"] - assert "/v1/responses" in get_args(endpoint_annotation) - - for create_fn in (litellm.create_batch, litellm.acreate_batch): - assert "/v1/responses" in get_args(get_type_hints(create_fn)["endpoint"]) diff --git a/tests/test_litellm/chat_completions/test_dispatch.py b/tests/unit/chat_completions/test_dispatch.py similarity index 93% rename from tests/test_litellm/chat_completions/test_dispatch.py rename to tests/unit/chat_completions/test_dispatch.py index d4bfeaf8d70..63821c74208 100644 --- a/tests/test_litellm/chat_completions/test_dispatch.py +++ b/tests/unit/chat_completions/test_dispatch.py @@ -1,11 +1,9 @@ -import inspect from collections.abc import Awaitable, Callable, Mapping -from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect +from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures import pytest import litellm -from litellm import main as python_chat from litellm.chat_completions.dispatch import ( _ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch _DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch @@ -40,15 +38,6 @@ def acompletion_binding(native: NativeAcompletion | None) -> NativeBinding[Nativ return binding -def test_public_signature_is_the_legacy_signature() -> None: - public_completion: Final = cast(Callable[..., object], litellm.completion) - legacy_completion: Final = cast(Callable[..., object], python_chat.completion) - public_acompletion: Final = cast(Callable[..., object], litellm.acompletion) - legacy_acompletion: Final = cast(Callable[..., object], python_chat.acompletion) - assert inspect.signature(public_completion) == inspect.signature(legacy_completion) - assert inspect.signature(public_acompletion) == inspect.signature(legacy_acompletion) - - def test_python_route_forwards_original_call_shape() -> None: metadata: Final = {"user_id": "u"} args: Final[tuple[object, ...]] = ("gpt-4o", MESSAGES) diff --git a/tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/test_litellm_responses_transformation_transformation.py similarity index 100% rename from tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py rename to tests/unit/completion_extras/test_litellm_responses_transformation_transformation.py diff --git a/tests/test_litellm/compression/test_compress.py b/tests/unit/compression/test_compress.py similarity index 100% rename from tests/test_litellm/compression/test_compress.py rename to tests/unit/compression/test_compress.py diff --git a/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py b/tests/unit/endpoints/speech/speech_to_completion_bridge/test_transformation.py similarity index 100% rename from tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py rename to tests/unit/endpoints/speech/speech_to_completion_bridge/test_transformation.py diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_callback_controls.py b/tests/unit/enterprise/enterprise_callbacks/test_callback_controls.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/test_callback_controls.py rename to tests/unit/enterprise/enterprise_callbacks/test_callback_controls.py diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py b/tests/unit/enterprise/enterprise_callbacks/test_llm_guard.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/test_llm_guard.py rename to tests/unit/enterprise/enterprise_callbacks/test_llm_guard.py diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py rename to tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py diff --git a/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py b/tests/unit/integrations/compression_interception/test_compression_interception_handler.py similarity index 100% rename from tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py rename to tests/unit/integrations/compression_interception/test_compression_interception_handler.py diff --git a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py b/tests/unit/integrations/gcs_bucket/test_gcs_bucket_base.py similarity index 100% rename from tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py rename to tests/unit/integrations/gcs_bucket/test_gcs_bucket_base.py diff --git a/tests/test_litellm/integrations/gcs_pubsub/test_pub_sub.py b/tests/unit/integrations/gcs_pubsub/test_pub_sub.py similarity index 100% rename from tests/test_litellm/integrations/gcs_pubsub/test_pub_sub.py rename to tests/unit/integrations/gcs_pubsub/test_pub_sub.py diff --git a/tests/test_litellm/integrations/helicone/test_helicone_gemini.py b/tests/unit/integrations/helicone/test_helicone_gemini.py similarity index 73% rename from tests/test_litellm/integrations/helicone/test_helicone_gemini.py rename to tests/unit/integrations/helicone/test_helicone_gemini.py index 8ce02784345..667b16a48a1 100644 --- a/tests/test_litellm/integrations/helicone/test_helicone_gemini.py +++ b/tests/unit/integrations/helicone/test_helicone_gemini.py @@ -3,7 +3,6 @@ Test HeliconeLogger Gemini/Vertex AI support. Fixes: https://github.com/BerriAI/litellm/issues/19093 """ -import pytest def test_helicone_gemini_model_in_list(): @@ -36,39 +35,6 @@ def test_helicone_gemini_models_recognized(): assert is_recognized, f"{model} should be recognized by helicone_model_list" -def test_helicone_vertex_ai_models_recognized(): - """ - Test that Vertex AI models (GLM, DeepSeek, etc.) are recognized via custom_llm_provider. - """ - # Test models that don't contain "gemini" but are vertex_ai - test_models = [ - "vertex_ai/zai-org/glm-4.7-maas", - "vertex_ai/deepseek-ai/deepseek-v3", - "vertex_ai/meta/llama-3.1-405b", - ] - for model in test_models: - is_vertex_ai = model.startswith("vertex_ai/") - assert is_vertex_ai, f"{model} should be recognized as vertex_ai model" - - -def test_helicone_vertex_ai_via_custom_llm_provider(): - """ - Test that vertex_ai models are recognized when custom_llm_provider is set. - """ - # Models without vertex_ai/ prefix but with custom_llm_provider="vertex_ai" - test_cases = [ - ("zai-org/glm-4.7-maas", "vertex_ai"), - ("deepseek-ai/deepseek-v3", "vertex_ai"), - ] - for model, custom_llm_provider in test_cases: - is_vertex_ai = custom_llm_provider == "vertex_ai" or model.startswith( - "vertex_ai/" - ) - assert ( - is_vertex_ai - ), f"{model} with custom_llm_provider={custom_llm_provider} should be recognized as vertex_ai" - - def test_helicone_vertex_gemini_gets_vertex_provider_url(): """ Test that vertex_ai/gemini-* models route to aiplatform.googleapis.com, From ba629f2537a73f74f5d466e8b7d2703902b74014 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 09:37:17 +0000 Subject: [PATCH 21/47] test(bedrock): wrap long lines flagged by review in migrated unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../image/test_bedrock_image_prepare_request.py | 9 ++++++--- .../test_bedrock_passthrough_transformation.py | 13 ++++++++++--- .../realtime/test_bedrock_realtime_handler.py | 8 +++++++- .../rerank/test_bedrock_rerank_header_forwarding.py | 11 +++++++++-- .../test_bedrock_vector_store_transformation.py | 3 ++- 5 files changed, 34 insertions(+), 10 deletions(-) diff --git a/tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py b/tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py index 1575ccb5739..b010db3a840 100644 --- a/tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py +++ b/tests/unit/llms/bedrock/image/test_bedrock_image_prepare_request.py @@ -11,7 +11,8 @@ def test_bedrock_image_prepare_request_with_arn() -> None: with ( patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration._get_boto_credentials_from_optional_params" + "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration." + "_get_boto_credentials_from_optional_params" ), patch( "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.get_request_headers" @@ -31,7 +32,8 @@ def test_bedrock_image_prepare_request_with_arn() -> None: assert ( request.endpoint_url - == "https://bedrock-runtime.test.com/model/arn%3Aaws%3Abedrock%3Aus-east-1%3A123456789012%3Aapplication-inference-profile%2Fabcdefghi123/invoke" + == "https://bedrock-runtime.test.com/model/arn%3Aaws%3Abedrock%3Aus-east-1%3A123456789012" + "%3Aapplication-inference-profile%2Fabcdefghi123/invoke" ) @@ -41,7 +43,8 @@ def test_bedrock_image_prepare_request_without_arn() -> None: with ( patch( - "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration._get_boto_credentials_from_optional_params" + "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration." + "_get_boto_credentials_from_optional_params" ), patch( "litellm.llms.bedrock.image_generation.image_handler.BedrockImageGeneration.get_request_headers" diff --git a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index 854ef92fa4b..d1d636a15f7 100644 --- a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -419,7 +419,9 @@ def test_bedrock_passthrough_model_id_arn_encoding(): ), f"ARN slash should be encoded, but found unencoded version in: {url_str}" # Verify the complete expected URL structure - expected_encoded_model_id = "arn:aws:bedrock:us-east-1:590183661440:application-inference-profile%2Fb943q2qbl3m7" + expected_encoded_model_id = ( + "arn:aws:bedrock:us-east-1:590183661440:application-inference-profile%2Fb943q2qbl3m7" + ) expected_url = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{expected_encoded_model_id}/converse" assert url_str == expected_url, f"Expected {expected_url}, but got: {url_str}" @@ -515,7 +517,10 @@ def test_bedrock_passthrough_model_id_without_arn(): def _event_frame(event_type: str, payload: dict) -> bytes: def header(name: str, value: str) -> bytes: name_b, value_b = name.encode(), value.encode() - return struct.pack("!B", len(name_b)) + name_b + struct.pack("!B", 7) + struct.pack("!H", len(value_b)) + value_b + return ( + struct.pack("!B", len(name_b)) + name_b + + struct.pack("!B", 7) + struct.pack("!H", len(value_b)) + value_b + ) payload_b = json.dumps(payload, separators=(",", ":")).encode() headers_b = ( @@ -589,7 +594,9 @@ def _feed(collector: PassthroughStreamCollector, stream: bytes, chunk_size: int def test_converse_stream_collector_keeps_usage_without_retaining_the_stream(): texts = [f"tok{i} " for i in range(4000)] - stream = _event_frame("messageStart", {"role": "assistant"}) + _text_block(0, texts) + _stream_tail("end_turn", 4000) + stream = ( + _event_frame("messageStart", {"role": "assistant"}) + _text_block(0, texts) + _stream_tail("end_turn", 4000) + ) _feed(_converse_stream_collector(), stream) tracemalloc.start() diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 73a78a94e9f..3aa827beb80 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -23,7 +23,13 @@ def _isolate_host_aws_config(monkeypatch, tmp_path): monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "credentials")) monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config")) monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") - for env_var in ("AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION_NAME", "AWS_DEFAULT_REGION"): + for env_var in ( + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_BEARER_TOKEN_BEDROCK", + "AWS_REGION_NAME", + "AWS_DEFAULT_REGION", + ): monkeypatch.delenv(env_var, raising=False) diff --git a/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index aa93ddb21b8..c40830b238f 100644 --- a/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/unit/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -21,7 +21,13 @@ def _isolate_host_aws_config(monkeypatch, tmp_path): monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "credentials")) monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config")) monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") - for env_var in ("AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION_NAME", "AWS_DEFAULT_REGION"): + for env_var in ( + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_BEARER_TOKEN_BEDROCK", + "AWS_REGION_NAME", + "AWS_DEFAULT_REGION", + ): monkeypatch.delenv(env_var, raising=False) # Mock response for Bedrock rerank @@ -39,7 +45,8 @@ bedrock_rerank_response = { test_query = "What is the capital of the United States?" test_documents = [ "Carson City is the capital city of the American state of Nevada.", - "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.", + "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. " + "Its capital is Saipan.", "Washington, D.C. is the capital of the United States.", ] diff --git a/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index ab5a2531461..b45e70e31d3 100644 --- a/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/unit/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -46,7 +46,8 @@ def test_transform_search_request_encodes_vector_store_id(): assert ( url - == "https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases/..%2F..%2Fknowledgebases%2Fother%3Fx%3D1%23frag/retrieve" + == "https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases/..%2F..%2Fknowledgebases%2Fother" + "%3Fx%3D1%23frag/retrieve" ) assert body["retrievalQuery"].get("text") == "hello" From f61abe00a6bd3eff086e0fc5fc503e4418850d3e Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 09:58:23 +0000 Subject: [PATCH 22/47] test(llms): wrap remaining lines over 120 chars in migrated unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_translation/test_handler.py | 12 ++++-- ...drock_mantle_passthrough_transformation.py | 8 +++- .../chat/test_bytez_chat_transformation.py | 42 +++++++++++++++---- tests/unit/llms/chat/test_converse_handler.py | 5 ++- 4 files changed, 53 insertions(+), 14 deletions(-) diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index dee8366ce2d..da7ed635dcb 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -1072,7 +1072,8 @@ class TestDeAnonymizeConverseStream: @pytest.mark.asyncio async def test_reasoning_text_delta_de_anonymized(self): - """Reasoning deltas carry model output; their text must be guardrailed while the reasoning signature is left untouched.""" + """Reasoning deltas carry model output; their text must be guardrailed while the + reasoning signature is left untouched.""" stream_bytes = ( _build_event_stream_frame("messageStart", {"role": "assistant"}) + _build_event_stream_frame( @@ -1105,7 +1106,8 @@ class TestDeAnonymizeConverseStream: @pytest.mark.asyncio async def test_tool_use_input_delta_de_anonymized(self): - """toolUse.input deltas carry model-generated tool arguments and must be guardrailed instead of being forwarded raw.""" + """toolUse.input deltas carry model-generated tool arguments and must be + guardrailed instead of being forwarded raw.""" stream_bytes = _build_event_stream_frame( "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"toolUse": {"input": '{"q":""}'}}}, @@ -1154,7 +1156,8 @@ class TestDeAnonymizeConverseStream: @pytest.mark.asyncio async def test_text_and_reasoning_deltas_de_anonymized_independently(self): - """Distinct delta kinds must each be guardrailed and written back into their own field without bleeding the de-anonymized text across kinds.""" + """Distinct delta kinds must each be guardrailed and written back into their own + field without bleeding the de-anonymized text across kinds.""" captured = {} async def mock_hook(data, user_api_key_dict, response): @@ -1192,7 +1195,8 @@ class TestDeAnonymizeConverseStream: @pytest.mark.asyncio async def test_reasoning_signature_only_frame_left_unmodified(self): - """A reasoning delta carrying only a signature has no guardrailable text; it must be forwarded untouched and the guardrail must not run.""" + """A reasoning delta carrying only a signature has no guardrailable text; it must + be forwarded untouched and the guardrail must not run.""" stream_bytes = _build_event_stream_frame( "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"reasoningContent": {"signature": "sig"}}}, diff --git a/tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py b/tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py index 090de0a9d3e..27f6c9a9140 100644 --- a/tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py +++ b/tests/unit/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py @@ -109,7 +109,13 @@ def test_region_falls_back_to_the_mantle_default_without_any_hint(no_ambient_aws ({}, {"AWS_BEARER_TOKEN_BEDROCK": "aws-env-key"}, "aws-env-key"), ], ) -def test_sign_request_uses_the_deployment_bearer_token(no_ambient_aws, monkeypatch, litellm_params, env, expected_bearer): +def test_sign_request_uses_the_deployment_bearer_token( + no_ambient_aws, + monkeypatch, + litellm_params, + env, + expected_bearer, +): for name, value in env.items(): monkeypatch.setenv(name, value) headers, body = BedrockMantlePassthroughConfig().sign_request( diff --git a/tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py b/tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py index 440304aeac1..157e2e51175 100644 --- a/tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py +++ b/tests/unit/llms/bytez/chat/test_bytez_chat_transformation.py @@ -8,6 +8,32 @@ from litellm.llms.bytez.chat.transformation import BytezChatConfig, API_BASE, ve TEST_API_KEY = "MOCK_BYTEZ_API_KEY" TEST_MODEL_NAME = "google/gemma-3-4b-it" TEST_MODEL = f"bytez/{TEST_MODEL_NAME}" +CAT_IMAGE_URL = ( + "https://images.squarespace-cdn.com/content/v1/5452d441e4b0c188b51fef1a/1615326541809-TW01PVTOJ4PXQUX" + "VRLHI/male-orange-tabby-cat.jpg" +) +KAGGLE_AUDIO_URL = ( + "https://storage.googleapis.com/kagglesdsdata/datasets/1736753/2838478/dataset/dataset/B_ANI01_MC_FN_" + "SIM01_101.wav?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-1616" + "07.iam.gserviceaccount.com%2F20250711%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T192905Z&" + "X-Goog-Expires=345600&X-Goog-SignedHeaders=host&X-Goog-Signature=812b4bd6fcf9296f8e34f67664d900a81cf" + "81a4c8a4f439ce12befc89b4bef07c2645cab20ce5ba8f6b311dffa85aa05b70b4efbe53bced50a43a5e7622ea1ee0d8cc39" + "0679cdc6a6aae2c27f75debc1ce2361c595b3c9e1b8c88e2756ffc6b4f290af7f3dfa7232dc69ccc9a2181be756e0d538250" + "f9761a8b05ba1ac6c6b5d946f97a16aa14a5609ae62a2c4713c2077fcd34d129dbcdac6bb543ae547507b1a424e4fd09f817" + "000943c11507e0a74c514ec212b17427b7fc9e2ce87a250db1258645e4862a4261e3790fd99c9186148ad0653acd2b6a9468" + "adbeb94f17b5a685551037fd2cc9fe72fa405a006c0bd42d03be1e4c0dc4023ed3a77171edff3" +) +KAGGLE_VIDEO_URL = ( + "https://storage.googleapis.com/kagglesdsdata/datasets/3957252/6888743/dog1.mp4?X-Goog-Algorithm=GOOG" + "4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-161607.iam.gserviceaccount.com%2F202507" + "11%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T193025Z&X-Goog-Expires=345600&X-Goog-Signed" + "Headers=host&X-Goog-Signature=810961d9abcbc2437954fdf19ef216deb65d3977eb354ec10af0d4644627cc6b143a5f" + "c6450996bae1787c09d26334de7cd6ff887510a5ac2a6eed3cfcc6673a47686c84c1f2b0bf543009388d83f2cd9551ad5f72" + "084513c6a7acd2c718849a4ebe951ccc5631bed014b0d115225c048b9f5de68673a37db24a98ad39cf3d0ba16fb764bf38eb" + "90c78c295c21a4ddac08c3c661b65efd511ccb86bacb87a2e2a97a06f53ea1c64d5dcf274001a61bc20867802549601301d9" + "99f5a5b2e49fd444b7db860c68e1c67df6e8edd5ad97171eaafb4fa1462453924ea4d78733be411cb6b5c910d4f829cd7189" + "c28dc1b22c8ae2a4da844a0d202e9e64bc7fb17947" +) TEST_MESSAGES = [{"role": "user", "content": "Hello"}] @@ -148,7 +174,7 @@ class TestBytezChatConfig: "What color is this cat?", { "type": "image_url", - "url": "https://images.squarespace-cdn.com/content/v1/5452d441e4b0c188b51fef1a/1615326541809-TW01PVTOJ4PXQUXVRLHI/male-orange-tabby-cat.jpg", + "url": CAT_IMAGE_URL, }, ], } @@ -160,7 +186,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What color is this cat?"}, { "type": "image", - "url": "https://images.squarespace-cdn.com/content/v1/5452d441e4b0c188b51fef1a/1615326541809-TW01PVTOJ4PXQUXVRLHI/male-orange-tabby-cat.jpg", + "url": CAT_IMAGE_URL, }, ], } @@ -174,7 +200,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What color is this cat?"}, { "type": "image_url", - "url": "https://images.squarespace-cdn.com/content/v1/5452d441e4b0c188b51fef1a/1615326541809-TW01PVTOJ4PXQUXVRLHI/male-orange-tabby-cat.jpg", + "url": CAT_IMAGE_URL, }, ], } @@ -186,7 +212,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What color is this cat?"}, { "type": "image", - "url": "https://images.squarespace-cdn.com/content/v1/5452d441e4b0c188b51fef1a/1615326541809-TW01PVTOJ4PXQUXVRLHI/male-orange-tabby-cat.jpg", + "url": CAT_IMAGE_URL, }, ], } @@ -200,7 +226,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What kind of cat meow is this?"}, { "type": "input_audio", - "url": "https://storage.googleapis.com/kagglesdsdata/datasets/1736753/2838478/dataset/dataset/B_ANI01_MC_FN_SIM01_101.wav?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-161607.iam.gserviceaccount.com%2F20250711%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T192905Z&X-Goog-Expires=345600&X-Goog-SignedHeaders=host&X-Goog-Signature=812b4bd6fcf9296f8e34f67664d900a81cf81a4c8a4f439ce12befc89b4bef07c2645cab20ce5ba8f6b311dffa85aa05b70b4efbe53bced50a43a5e7622ea1ee0d8cc390679cdc6a6aae2c27f75debc1ce2361c595b3c9e1b8c88e2756ffc6b4f290af7f3dfa7232dc69ccc9a2181be756e0d538250f9761a8b05ba1ac6c6b5d946f97a16aa14a5609ae62a2c4713c2077fcd34d129dbcdac6bb543ae547507b1a424e4fd09f817000943c11507e0a74c514ec212b17427b7fc9e2ce87a250db1258645e4862a4261e3790fd99c9186148ad0653acd2b6a9468adbeb94f17b5a685551037fd2cc9fe72fa405a006c0bd42d03be1e4c0dc4023ed3a77171edff3", + "url": KAGGLE_AUDIO_URL, }, ], } @@ -212,7 +238,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What kind of cat meow is this?"}, { "type": "audio", - "url": "https://storage.googleapis.com/kagglesdsdata/datasets/1736753/2838478/dataset/dataset/B_ANI01_MC_FN_SIM01_101.wav?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-161607.iam.gserviceaccount.com%2F20250711%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T192905Z&X-Goog-Expires=345600&X-Goog-SignedHeaders=host&X-Goog-Signature=812b4bd6fcf9296f8e34f67664d900a81cf81a4c8a4f439ce12befc89b4bef07c2645cab20ce5ba8f6b311dffa85aa05b70b4efbe53bced50a43a5e7622ea1ee0d8cc390679cdc6a6aae2c27f75debc1ce2361c595b3c9e1b8c88e2756ffc6b4f290af7f3dfa7232dc69ccc9a2181be756e0d538250f9761a8b05ba1ac6c6b5d946f97a16aa14a5609ae62a2c4713c2077fcd34d129dbcdac6bb543ae547507b1a424e4fd09f817000943c11507e0a74c514ec212b17427b7fc9e2ce87a250db1258645e4862a4261e3790fd99c9186148ad0653acd2b6a9468adbeb94f17b5a685551037fd2cc9fe72fa405a006c0bd42d03be1e4c0dc4023ed3a77171edff3", + "url": KAGGLE_AUDIO_URL, }, ], } @@ -226,7 +252,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What kind of dog is this?"}, { "type": "video_url", - "url": "https://storage.googleapis.com/kagglesdsdata/datasets/3957252/6888743/dog1.mp4?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-161607.iam.gserviceaccount.com%2F20250711%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T193025Z&X-Goog-Expires=345600&X-Goog-SignedHeaders=host&X-Goog-Signature=810961d9abcbc2437954fdf19ef216deb65d3977eb354ec10af0d4644627cc6b143a5fc6450996bae1787c09d26334de7cd6ff887510a5ac2a6eed3cfcc6673a47686c84c1f2b0bf543009388d83f2cd9551ad5f72084513c6a7acd2c718849a4ebe951ccc5631bed014b0d115225c048b9f5de68673a37db24a98ad39cf3d0ba16fb764bf38eb90c78c295c21a4ddac08c3c661b65efd511ccb86bacb87a2e2a97a06f53ea1c64d5dcf274001a61bc20867802549601301d999f5a5b2e49fd444b7db860c68e1c67df6e8edd5ad97171eaafb4fa1462453924ea4d78733be411cb6b5c910d4f829cd7189c28dc1b22c8ae2a4da844a0d202e9e64bc7fb17947", + "url": KAGGLE_VIDEO_URL, }, ], } @@ -238,7 +264,7 @@ class TestBytezChatConfig: {"type": "text", "text": "What kind of dog is this?"}, { "type": "video", - "url": "https://storage.googleapis.com/kagglesdsdata/datasets/3957252/6888743/dog1.mp4?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=databundle-worker-v2%40kaggle-161607.iam.gserviceaccount.com%2F20250711%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250711T193025Z&X-Goog-Expires=345600&X-Goog-SignedHeaders=host&X-Goog-Signature=810961d9abcbc2437954fdf19ef216deb65d3977eb354ec10af0d4644627cc6b143a5fc6450996bae1787c09d26334de7cd6ff887510a5ac2a6eed3cfcc6673a47686c84c1f2b0bf543009388d83f2cd9551ad5f72084513c6a7acd2c718849a4ebe951ccc5631bed014b0d115225c048b9f5de68673a37db24a98ad39cf3d0ba16fb764bf38eb90c78c295c21a4ddac08c3c661b65efd511ccb86bacb87a2e2a97a06f53ea1c64d5dcf274001a61bc20867802549601301d999f5a5b2e49fd444b7db860c68e1c67df6e8edd5ad97171eaafb4fa1462453924ea4d78733be411cb6b5c910d4f829cd7189c28dc1b22c8ae2a4da844a0d202e9e64bc7fb17947", + "url": KAGGLE_VIDEO_URL, }, ], } diff --git a/tests/unit/llms/chat/test_converse_handler.py b/tests/unit/llms/chat/test_converse_handler.py index 12b5f03aedc..05debee0602 100644 --- a/tests/unit/llms/chat/test_converse_handler.py +++ b/tests/unit/llms/chat/test_converse_handler.py @@ -106,7 +106,10 @@ class TestBedrockRegionInModelPath: ), f"modelId mismatch for {model!r}: got {model_id!r}, expected {expected_model_id!r}" assert ( optional_params.get("aws_region_name") == expected_region - ), f"region mismatch for {model!r}: got {optional_params.get('aws_region_name')!r}, expected {expected_region!r}" + ), ( + f"region mismatch for {model!r}: " + f"got {optional_params.get('aws_region_name')!r}, expected {expected_region!r}" + ) def test_explicit_aws_region_name_not_overridden(self): """ From ccd8997b0846469ff0e624b2cd837c4d0f8a3da6 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 10:16:48 +0000 Subject: [PATCH 23/47] refactor(types): replace Any with proven types in 34 files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock_agentcore/transformation.py | 2 +- .../providers/watsonx_orchestrate/config.py | 4 ++-- litellm/a2a_protocol/utils.py | 6 ++--- .../gitlab/gitlab_prompt_manager.py | 4 ++-- litellm/integrations/otel/presets/agentops.py | 7 ++++-- .../integrations/vantage/vantage_logger.py | 4 ++-- litellm/interactions/agents/http_handler.py | 22 +++++++++---------- litellm/litellm_core_utils/logging_utils.py | 4 +++- litellm/litellm_core_utils/url_utils.py | 4 ++-- litellm/llms/anthropic/chat/handler.py | 2 +- litellm/llms/azure/completion/handler.py | 9 ++++---- .../document_intelligence/transformation.py | 8 +++---- .../guardrail_translation/base_translation.py | 2 +- .../llms/bedrock/batches/transformation.py | 10 ++++----- ...mazon_twelvelabs_pegasus_transformation.py | 2 +- litellm/llms/bytez/chat/transformation.py | 6 ++--- litellm/llms/custom_httpx/aiohttp_handler.py | 4 ++-- .../llms/deprecated_providers/aleph_alpha.py | 2 +- litellm/llms/lemonade/chat/transformation.py | 2 +- .../llms/openai/image_edit/transformation.py | 2 +- .../runwayml/text_to_speech/transformation.py | 2 +- .../llms/vertex_ai/files/transformation.py | 6 ++--- .../batch_embed_content_handler.py | 4 ++-- .../llms/vertex_ai/vertex_ai_non_gemini.py | 2 +- .../audio_transcription/transformation.py | 4 ++-- litellm/proxy/a2a/agent_card.py | 16 +++++++------- .../proxy/agent_endpoints/a2a_endpoints.py | 2 +- litellm/proxy/client/credentials.py | 5 +++-- .../guardrails_ai/guardrails_ai.py | 4 ++-- .../guardrail_hooks/singulr/singulr.py | 6 ++--- .../object_permission_utils.py | 13 ++++++----- .../proxy/policy_engine/pipeline_executor.py | 12 +++++----- .../router_strategy/adaptive_router/hooks.py | 2 +- .../auto_router/auto_router.py | 4 ++-- 34 files changed, 98 insertions(+), 90 deletions(-) diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py index 9fa9db48af8..c486f1f6d95 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py @@ -85,7 +85,7 @@ def _filter_reserved_headers( def _request_scoped_runtime_session_id( - params: Mapping[str, Any], + params: Mapping[str, object], litellm_params: Mapping[str, Any], ) -> str | None: context_id: Final = get_session_id_from_a2a_params(params) diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py index ca84d3e07b4..44873edf271 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py @@ -20,7 +20,7 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): params: dict[str, Any], api_base: str | None = None, **kwargs: Any, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Handle a non-streaming A2A request via WXO runs API.""" litellm_params: Final = kwargs.get("litellm_params") if not litellm_params: @@ -40,7 +40,7 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): params: dict[str, Any], api_base: str | None = None, **kwargs: Any, - ) -> AsyncIterator[dict[str, Any]]: + ) -> AsyncIterator[dict[str, object]]: """Handle a streaming A2A request via WXO streaming runs API.""" litellm_params: Final = kwargs.get("litellm_params") if not litellm_params: diff --git a/litellm/a2a_protocol/utils.py b/litellm/a2a_protocol/utils.py index 47f561068cd..7400844bf28 100644 --- a/litellm/a2a_protocol/utils.py +++ b/litellm/a2a_protocol/utils.py @@ -17,7 +17,7 @@ class A2ARequestUtils: """Utility class for A2A request/response processing.""" @staticmethod - def extract_text_from_message(message: Any) -> str: + def extract_text_from_message(message: object) -> str: """ Extract text content from A2A message parts. @@ -142,7 +142,7 @@ class A2ARequestUtils: return prompt_tokens, completion_tokens, total_tokens -def get_session_id_from_a2a_params(params: Mapping[str, Any]) -> str | None: +def get_session_id_from_a2a_params(params: Mapping[str, object]) -> str | None: message: Final = params.get("message", {}) if isinstance(message, dict): return message.get("contextId") @@ -166,7 +166,7 @@ def scope_session_to_principal(session_id: str, principal: str | None) -> str: # Backwards compatibility aliases -def extract_text_from_a2a_message(message: Any) -> str: +def extract_text_from_a2a_message(message: object) -> str: return A2ARequestUtils.extract_text_from_message(message) diff --git a/litellm/integrations/gitlab/gitlab_prompt_manager.py b/litellm/integrations/gitlab/gitlab_prompt_manager.py index d4602176650..817d280074f 100644 --- a/litellm/integrations/gitlab/gitlab_prompt_manager.py +++ b/litellm/integrations/gitlab/gitlab_prompt_manager.py @@ -200,8 +200,8 @@ class GitLabTemplateManager: metadata=metadata, ) - def _parse_yaml_basic(self, yaml_str: str) -> dict[str, Any]: - result: Final[dict[str, Any]] = {} + def _parse_yaml_basic(self, yaml_str: str) -> dict[str, bool | int | float | str]: + result: Final[dict[str, bool | int | float | str]] = {} for line in yaml_str.split("\n"): line = line.strip() if ":" in line and not line.startswith("#"): diff --git a/litellm/integrations/otel/presets/agentops.py b/litellm/integrations/otel/presets/agentops.py index 965213f2ee4..58123656caa 100644 --- a/litellm/integrations/otel/presets/agentops.py +++ b/litellm/integrations/otel/presets/agentops.py @@ -9,9 +9,12 @@ this preset registers a custom exporter (``kind="agentops"``) that mints the JWT worker thread, off any event loop — and caches it for the process lifetime. """ +from collections.abc import Sequence from typing import Any, Final import httpx +from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult from pydantic import Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -71,7 +74,7 @@ def agentops_preset( ) -def _build_agentops_exporter(spec: ExporterSpec) -> Any: +def _build_agentops_exporter(spec: ExporterSpec) -> SpanExporter: """Factory for the ``agentops`` exporter kind: a lazy-auth OTLP/HTTP exporter.""" from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( OTLPSpanExporter, @@ -106,7 +109,7 @@ def _build_agentops_exporter(spec: ExporterSpec) -> Any: except Exception as e: verbose_logger.debug("AgentOps JWT fetch failed: %s", e) - def export(self, spans: Any) -> Any: + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: self._ensure_authenticated() return super().export(spans) diff --git a/litellm/integrations/vantage/vantage_logger.py b/litellm/integrations/vantage/vantage_logger.py index c219ba392ab..48a492fdd72 100644 --- a/litellm/integrations/vantage/vantage_logger.py +++ b/litellm/integrations/vantage/vantage_logger.py @@ -59,7 +59,7 @@ class VantageLogger(FocusLogger): raw_interval, ) - destination_config: Final[dict[str, Any]] = {} + destination_config: Final[dict[str, str]] = {} if resolved_api_key: destination_config["api_key"] = resolved_api_key if resolved_token: @@ -93,7 +93,7 @@ class VantageLogger(FocusLogger): pod_lock_manager = 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/interactions/agents/http_handler.py b/litellm/interactions/agents/http_handler.py index ec9df0fb488..2afedc34d36 100644 --- a/litellm/interactions/agents/http_handler.py +++ b/litellm/interactions/agents/http_handler.py @@ -7,7 +7,7 @@ duplicated. BaseAgentsAPIConfig stays as pure transform code. """ from collections.abc import Coroutine, Mapping -from typing import Any, Final +from typing import Final import httpx @@ -38,7 +38,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_body: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, @@ -93,7 +93,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, extra_body: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, @@ -141,7 +141,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): agents_api_config: BaseAgentsAPIConfig, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, @@ -181,7 +181,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): agents_api_config: BaseAgentsAPIConfig, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> AgentListResponse: @@ -216,7 +216,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, @@ -259,7 +259,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> AgentCreateResponse: @@ -295,7 +295,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, @@ -338,7 +338,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> AgentDeleteResult: @@ -374,7 +374,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, _is_async: bool = False, @@ -417,7 +417,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): name: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> AgentVersionsResponse: diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 5be9dd7be2f..38a501ecaae 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -67,7 +67,9 @@ def _truncate_base64_in_string(value: str) -> str: return _DATA_URI_RE.sub(_base64_data_uri_replacer, value) -def _truncate_base64_in_value(value: Any) -> Any: +def _truncate_base64_in_value( + value: str | dict[str, object] | list[object] | None, +) -> str | dict[str, object] | list[object] | None: """Iteratively truncate base64 data URIs in a JSON-like value (str/list/dict). Uses an explicit stack instead of recursion to satisfy the project's diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index fa070a648f5..b94d5a6886d 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -418,7 +418,7 @@ def _extract_redirect_url(response: httpx.Response, request_url: str) -> str: return str(httpx.URL(request_url).join(location)) -def safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response: +def safe_get(client: _UrlFetcher, url: str, **kwargs: Any) -> httpx.Response: """ Fetch a user-supplied URL with SSRF protection on every redirect hop. @@ -461,7 +461,7 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response: raise SSRFError("Too many redirects") -async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response: +async def async_safe_get(client: _AsyncUrlFetcher, url: str, **kwargs: Any) -> httpx.Response: """Async version of safe_get.""" if not getattr(litellm, "user_url_validation", True): kwargs.setdefault("follow_redirects", True) diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 359b8bb08c9..ef0f45d8f8b 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -596,7 +596,7 @@ class ModelResponseIterator: self.reasoning_content_chunks: list[str] = [] # Track server tool use inputs and results for code_interpreter_results - self._server_tool_inputs: dict[str, Any] = {} + self._server_tool_inputs: dict[str, object] = {} self.tool_results: list[dict[str, Any]] = [] self._current_server_tool_id: str | None = None self._container_id: str | None = None diff --git a/litellm/llms/azure/completion/handler.py b/litellm/llms/azure/completion/handler.py index 80934e994f6..23eef51e7ee 100644 --- a/litellm/llms/azure/completion/handler.py +++ b/litellm/llms/azure/completion/handler.py @@ -1,6 +1,7 @@ from collections.abc import Callable -from typing import Any, Final +from typing import Final +import httpx from openai import AsyncAzureOpenAI, AzureOpenAI from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -191,7 +192,7 @@ class AzureTextCompletion(BaseAzureLLM): model: str, api_base: str, data: dict, - timeout: Any, + timeout: float | httpx.Timeout | None, model_response: ModelResponse, logging_obj: LiteLLMLoggingObj, max_retries: int, @@ -253,7 +254,7 @@ class AzureTextCompletion(BaseAzureLLM): api_version: str, data: dict, model: str, - timeout: Any, + timeout: float | httpx.Timeout | None, azure_ad_token: str | None = None, client=None, litellm_params: dict = {}, @@ -306,7 +307,7 @@ class AzureTextCompletion(BaseAzureLLM): api_version: str, data: dict, model: str, - timeout: Any, + timeout: float | httpx.Timeout | None, azure_ad_token: str | None = None, client=None, litellm_params: dict = {}, diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index 3a2af8a5aba..23f532be757 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -12,7 +12,7 @@ import asyncio import re import time from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from urllib.parse import quote import httpx @@ -127,7 +127,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): def map_ocr_params( self, - non_default_params: dict, + non_default_params: Mapping[str, object], optional_params: dict, model: str, ) -> dict: @@ -164,7 +164,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): raise UnsupportedParamsError(message=f"{e}", model=model, llm_provider="azure_ai") from e @staticmethod - def _normalize_pages_param(pages: Any) -> str: + def _normalize_pages_param(pages: object) -> str: """ Convert a caller-provided `pages` value to Azure DI's query-string form. Azure expects 1-based page numbers, grammar: `^(\\d+(-\\d+)?)(,\\s*(\\d+(-\\d+)?))*$`. @@ -412,7 +412,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): raise ValueError("Document URL is required") # Build Azure DI request - data: Final[dict[str, Any]] = {} + data: Final[dict[str, str]] = {} # Check if it's a data URI (base64) if document_url.startswith("data:"): diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 89ad67f0485..ace5af8124f 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -81,7 +81,7 @@ class BaseTranslation(ABC): @staticmethod def transform_user_api_key_dict_to_metadata( - user_api_key_dict: Any | None, + user_api_key_dict: Optional["UserAPIKeyAuth"], ) -> dict[str, object]: """ Transform user_api_key_dict to a metadata dict with prefixed keys. diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index ae0f8c5935b..973388ca5bd 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -2,7 +2,7 @@ import os import re import time from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Final, Literal, cast from httpx import Headers, Response from pydantic import TypeAdapter, ValidationError @@ -170,7 +170,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): create_batch_data: CreateBatchRequest, optional_params: dict, litellm_params: dict, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform the batch creation request to Bedrock format. @@ -354,7 +354,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): ) @staticmethod - def _get_openai_compatible_batch_metadata(metadata: Any) -> dict[str, str]: + def _get_openai_compatible_batch_metadata(metadata: object) -> dict[str, str]: """ OpenAI Batch metadata only accepts string values. """ @@ -379,7 +379,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): batch_id: str, optional_params: dict, litellm_params: dict, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform batch retrieval request for Bedrock. @@ -523,7 +523,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): ) # Enrich metadata with useful Bedrock fields - enriched_metadata_raw: Final[dict[str, Any]] = { + enriched_metadata_raw: Final[dict[str, object]] = { "jobName": response_data.get("jobName"), "clientRequestToken": response_data.get("clientRequestToken"), "modelId": response_data.get("modelId"), diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py index d12c8aee48c..39cded4ed64 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py @@ -110,7 +110,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): headers: dict, ) -> dict: input_prompt: Final = self._convert_messages_to_prompt(messages=messages) - request_data: Final[dict[str, Any]] = {"inputPrompt": input_prompt} + request_data: Final[dict[str, object]] = {"inputPrompt": input_prompt} media_source: Final = self._build_media_source(optional_params) if media_source is not None: diff --git a/litellm/llms/bytez/chat/transformation.py b/litellm/llms/bytez/chat/transformation.py index d9a0c98b6db..7977db0f056 100644 --- a/litellm/llms/bytez/chat/transformation.py +++ b/litellm/llms/bytez/chat/transformation.py @@ -335,10 +335,10 @@ class BytezChatConfig(BaseConfig): class BytezCustomStreamWrapper(CustomStreamWrapper): - def chunk_creator(self, chunk: Any): + def chunk_creator(self, chunk: object): try: model_response: Final = self.model_response_creator() - response_obj: dict[str, Any] = {} + response_obj: dict[str, object] = {} response_obj = { "text": chunk, @@ -346,7 +346,7 @@ class BytezCustomStreamWrapper(CustomStreamWrapper): "finish_reason": "", } - completion_obj: Final[dict[str, Any]] = {"content": chunk} + completion_obj: Final[dict[str, object]] = {"content": chunk} return self.return_processed_chunk_logic( completion_obj=completion_obj, diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index 7035ce58ae1..0809ef5274f 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -1,5 +1,5 @@ import ssl -from collections.abc import Callable +from collections.abc import AsyncIterable, Callable, Iterable from typing import TYPE_CHECKING, Any, Final, cast import aiohttp @@ -212,7 +212,7 @@ class BaseLLMAIOHTTPHandler: litellm_params: dict, stream: bool = False, files: dict | None = None, - content: Any = None, + content: str | bytes | Iterable[bytes] | AsyncIterable[bytes] | None = None, params: dict | None = None, ) -> httpx.Response: max_retry_on_unprocessable_entity_error: Final = provider_config.max_retry_on_unprocessable_entity_error diff --git a/litellm/llms/deprecated_providers/aleph_alpha.py b/litellm/llms/deprecated_providers/aleph_alpha.py index 4a29549b6aa..2ad9ce4edc8 100644 --- a/litellm/llms/deprecated_providers/aleph_alpha.py +++ b/litellm/llms/deprecated_providers/aleph_alpha.py @@ -146,7 +146,7 @@ class AlephAlphaConfig: setattr(self.__class__, key, value) @classmethod - def get_config(cls): + def get_config(cls) -> dict[str, object]: return { k: v for k, v in cls.__dict__.items() diff --git a/litellm/llms/lemonade/chat/transformation.py b/litellm/llms/lemonade/chat/transformation.py index 553478aec16..c01ad2a0edd 100644 --- a/litellm/llms/lemonade/chat/transformation.py +++ b/litellm/llms/lemonade/chat/transformation.py @@ -170,7 +170,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig): model: str, api_base: str | None = None, api_key: str | None = None, - ) -> Any: + ) -> dict[str, object]: if model.startswith("lemonade/"): model = model.split("/", 1)[1] diff --git a/litellm/llms/openai/image_edit/transformation.py b/litellm/llms/openai/image_edit/transformation.py index f55084adbde..d54522597a0 100644 --- a/litellm/llms/openai/image_edit/transformation.py +++ b/litellm/llms/openai/image_edit/transformation.py @@ -66,7 +66,7 @@ class OpenAIImageEditConfig(BaseImageEditConfig): def _add_image_to_files( self, files_list: list[tuple[str, Any]], - image: Any, + image: object, field_name: str, ) -> None: """Add an image to the files list with appropriate content type""" diff --git a/litellm/llms/runwayml/text_to_speech/transformation.py b/litellm/llms/runwayml/text_to_speech/transformation.py index 19e6d8ff494..6769accc1d6 100644 --- a/litellm/llms/runwayml/text_to_speech/transformation.py +++ b/litellm/llms/runwayml/text_to_speech/transformation.py @@ -78,7 +78,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): aspeech: bool, api_base: str | None, api_key: str | None, - **kwargs: Any, + **kwargs: object, ) -> Union[ "HttpxBinaryResponseContent", Coroutine[object, object, "HttpxBinaryResponseContent"], diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 85ec2911464..80d32289c94 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -651,7 +651,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, Any]], + map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, object]], ) -> tuple[Mapping[str, object], ...]: """ Transforms a single OpenAI JSONL batch entry into the Vertex rows it maps to. @@ -774,7 +774,7 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): def __init__( self, openai_file_content: FileTypes, - map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], + map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, object]], ) -> None: self._openai_file_content = openai_file_content self._map_openai_to_vertex_params = map_openai_to_vertex_params @@ -948,7 +948,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _map_openai_to_vertex_params( self, openai_request_body: dict[str, Any], - ) -> dict[str, Any]: + ) -> dict[str, object]: """ wrapper to call VertexGeminiConfig.map_openai_params """ diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index f81d4ca777e..c6ac87d646b 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -3,7 +3,7 @@ Google AI Studio /batchEmbedContents Embeddings Endpoint """ import json -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Final, Literal import httpx @@ -210,7 +210,7 @@ class GoogleBatchEmbeddings(VertexLLM): ) ### TRANSFORMATION (sync path) ### - request_data: Any + request_data: VertexAIBatchEmbeddingsRequestBody | dict[str, object] if use_embed_content: resolved_files = {} if api_key: diff --git a/litellm/llms/vertex_ai/vertex_ai_non_gemini.py b/litellm/llms/vertex_ai/vertex_ai_non_gemini.py index 1c582c7c376..a7a1ea8d88d 100644 --- a/litellm/llms/vertex_ai/vertex_ai_non_gemini.py +++ b/litellm/llms/vertex_ai/vertex_ai_non_gemini.py @@ -64,7 +64,7 @@ def _get_client_from_cache(client_cache_key: str): return litellm.in_memory_llm_clients_cache.get_cache(client_cache_key) -def _set_client_in_cache(client_cache_key: str, vertex_llm_model: Any): +def _set_client_in_cache(client_cache_key: str, vertex_llm_model: object): litellm.in_memory_llm_clients_cache.set_cache( key=client_cache_key, value=vertex_llm_model, diff --git a/litellm/llms/watsonx/audio_transcription/transformation.py b/litellm/llms/watsonx/audio_transcription/transformation.py index 7d1aba63428..2169b9bf49a 100644 --- a/litellm/llms/watsonx/audio_transcription/transformation.py +++ b/litellm/llms/watsonx/audio_transcription/transformation.py @@ -4,7 +4,7 @@ Translates from OpenAI's `/v1/audio/transcriptions` to IBM WatsonX's `/ml/v1/aud WatsonX follows the OpenAI spec for audio transcription. """ -from typing import Any, Final +from typing import Final from httpx import Response @@ -124,7 +124,7 @@ class IBMWatsonXAudioTranscriptionConfig(IBMWatsonXMixin, OpenAIWhisperAudioTran } # Convert TypedDict to regular dict for AudioTranscriptionRequestData - form_data_dict: Final[dict[str, Any]] = dict(form_data) + form_data_dict: Final[dict[str, object]] = dict(form_data) return AudioTranscriptionRequestData(data=form_data_dict, files=files) diff --git a/litellm/proxy/a2a/agent_card.py b/litellm/proxy/a2a/agent_card.py index 5bec5158bcc..3718359fbb6 100644 --- a/litellm/proxy/a2a/agent_card.py +++ b/litellm/proxy/a2a/agent_card.py @@ -10,7 +10,7 @@ and uses LiteLLM auth. import re from collections.abc import Mapping from copy import deepcopy -from typing import Any, Final, Literal +from typing import Final, Literal SupportedA2AVersion = Literal["0.3", "1.0"] @@ -44,7 +44,7 @@ def normalize_protocol_version(version: object) -> SupportedA2AVersion | None: return next((supported for supported in SUPPORTED_A2A_PROTOCOL_VERSIONS if supported == major_minor), None) -def resolve_served_protocol_version(card: Mapping[str, Any] | None) -> str: +def resolve_served_protocol_version(card: Mapping[str, object] | None) -> str: """Return the validated protocol version an agent card pins, else the default.""" normalized: Final = normalize_protocol_version(card.get("protocolVersion") if card else None) return normalized if normalized is not None else LITELLM_A2A_PROTOCOL_VERSION @@ -53,7 +53,7 @@ def resolve_served_protocol_version(card: Mapping[str, Any] | None) -> str: # Security scheme exposed by the LiteLLM-fronted agent card. Always replaces # whatever upstream advertised — the client must authenticate to the proxy, # not the upstream agent. -LITELLM_SECURITY_SCHEMES: Final[dict[str, dict[str, Any]]] = { +LITELLM_SECURITY_SCHEMES: Final[dict[str, dict[str, str]]] = { "LiteLLMKey": { "type": "http", "scheme": "bearer", @@ -112,7 +112,7 @@ _ALLOWED_TOP_LEVEL_KEYS: Final = { "url", } -_DEFAULT_SKILLS: Final[list[dict[str, Any]]] = [ +_DEFAULT_SKILLS: Final[list[dict[str, str | list[str]]]] = [ { "id": "chat", "name": "Chat", @@ -129,7 +129,7 @@ _DEFAULT_MODES: Final[list[str]] = ["text"] _DEFAULT_AGENT_VERSION: Final = "1.0.0" -def _filter_capabilities(upstream_capabilities: Any) -> dict[str, Any]: +def _filter_capabilities(upstream_capabilities: object) -> dict[str, object]: """Return a capabilities dict containing only allowlisted, truthy keys.""" if not isinstance(upstream_capabilities, dict): return {} @@ -143,13 +143,13 @@ def _default_litellm_provider(proxy_base_url: str) -> dict[str, str]: def merge_agent_card( - upstream_card: Mapping[str, Any] | None, + upstream_card: Mapping[str, object] | None, *, proxy_url: str, proxy_base_url: str, name: str | None = None, description: str | None = None, -) -> dict[str, Any]: +) -> dict[str, object]: """ Build the LiteLLM-fronted agent card. @@ -169,7 +169,7 @@ def merge_agent_card( A dict suitable for serving as the proxy's agent card. Only keys in the v1.0 AgentCard schema (plus ``supportedInterfaces``) are emitted. """ - base: Final[dict[str, Any]] = deepcopy(dict(upstream_card)) if upstream_card else {} + base: Final[dict[str, object]] = deepcopy(dict(upstream_card)) if upstream_card else {} # Keep the upstream ``url`` on the stored card: the runtime A2A # invocation path reads it from ``agent_card_params`` to know where to diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 834c16ba6dc..faf3e98a3a7 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -880,7 +880,7 @@ async def invoke_agent_a2a( logging_obj._enqueue_deferred_logging = None _enqueue_fn() - response_dict: Final[dict[str, Any]] = ( + response_dict: Final[dict[str, object]] = ( response.model_dump(mode="json", exclude_none=True) if hasattr(response, "model_dump") else response diff --git a/litellm/proxy/client/credentials.py b/litellm/proxy/client/credentials.py index a9bff67b1c5..d9edecd2eb7 100644 --- a/litellm/proxy/client/credentials.py +++ b/litellm/proxy/client/credentials.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Final import requests @@ -69,8 +70,8 @@ class CredentialsManagementClient: def create( self, credential_name: str, - credential_info: dict[str, Any], - credential_values: dict[str, Any], + credential_info: Mapping[str, object], + credential_values: Mapping[str, object], return_request: bool = False, ) -> dict[str, Any] | requests.Request: """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py index 47324471650..18451df574f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py @@ -7,7 +7,7 @@ import json import os -from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict +from typing import TYPE_CHECKING, Final, Literal, TypedDict from fastapi import HTTPException @@ -181,7 +181,7 @@ class GuardrailsAI(CustomGuardrail): ): # raise exception if invalid, return a str for the user to receive - if rejected, or return a modified dictionary for passing into litellm return await self.process_input(data=data, call_type=call_type) - async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]: + async def async_logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]: if call_type == "acompletion" or call_type == "completion": kwargs = await self.process_input(data=kwargs, call_type=call_type) diff --git a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py index 06d4b39f5f6..bd5b18e368d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py @@ -24,7 +24,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk from litellm.types.proxy.guardrails.guardrail_hooks.base import ( GuardrailConfigModel, ) @@ -36,7 +36,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.singulr import ( ToolCall, ToolCallFunction, ) -from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs +from litellm.types.utils import CallTypes, ChatCompletionMessageToolCall, GenericGuardrailAPIInputs _DEFAULT_API_BASE: Final = "http://localhost:8003" _GUARD_ENDPOINT: Final = "/api/v1/ai-gateway/litellm-v2" @@ -339,7 +339,7 @@ class SingulrGuardrail(CustomGuardrail): return inputs @staticmethod - def _build_tool_call(tool_call: Mapping[str, Any]) -> "ToolCall | None": + def _build_tool_call(tool_call: ChatCompletionToolCallChunk | ChatCompletionMessageToolCall) -> "ToolCall | None": tool_call_id: Final = tool_call.get("id") fun: Final = tool_call.get("function") if not tool_call_id or not fun: diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index daab38d3662..61e432daa16 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -8,7 +8,7 @@ from collections.abc import Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Final, Optional from fastapi import HTTPException, status from pydantic import TypeAdapter @@ -156,12 +156,13 @@ async def handle_update_object_permission_common( if prisma_client is None: raise ValueError("Prisma client not found") - new_object_permission: dict | str | None = data_json.pop("object_permission", None) - if new_object_permission is None: + raw_object_permission: Final[dict | str | None] = data_json.pop("object_permission", None) + if raw_object_permission is None: return None - if isinstance(new_object_permission, str): - new_object_permission = json.loads(new_object_permission) + new_object_permission: Final[object] = ( + json.loads(raw_object_permission) if isinstance(raw_object_permission, str) else raw_object_permission + ) upsert: Final = await prepare_object_permission_upsert( new_object_permission=new_object_permission if isinstance(new_object_permission, dict) else {}, @@ -230,7 +231,7 @@ def _dedupe_preserving_order(values: list[str]) -> list[str]: return result -def _mcp_server_identifier_matches(server: Any, identifier: str) -> bool: +def _mcp_server_identifier_matches(server: object, identifier: str) -> bool: return identifier in { getattr(server, "server_id", None), getattr(server, "alias", None), diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 0b81e7af84d..e9d23436b59 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -8,7 +8,7 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding. import copy import time from collections.abc import Callable, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar +from typing import TYPE_CHECKING, Final, Literal, TypeVar from pydantic import BaseModel @@ -314,11 +314,11 @@ class PipelineExecutor: steps: list[PipelineStep], mode: str, data: dict, - user_api_key_dict: Any, + user_api_key_dict: "UserAPIKeyAuth", call_type: str, policy_name: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data - streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step + streaming_chunks: list[object] | None = None, # mutable-ok: shared buffered-stream chunks, read per step endpoint_translation: "BaseTranslation | None" = None, ) -> PipelineExecutionResult: """ @@ -490,10 +490,10 @@ class PipelineExecutor: step: PipelineStep, mode: str, data: dict, - user_api_key_dict: Any, + user_api_key_dict: "UserAPIKeyAuth", call_type: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data - streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step + streaming_chunks: list[object] | None = None, # mutable-ok: shared buffered-stream chunks, read per step endpoint_translation: "BaseTranslation | None" = None, ) -> tuple[ Literal["pass", "fail", "error"], @@ -722,7 +722,7 @@ def _extract_error_message(e: Exception) -> str: if isinstance(e, ModifyResponseException): return str(e) if HTTPException is not None and isinstance(e, HTTPException): - detail: Final = getattr(e, "detail", None) + detail: Final[object] = getattr(e, "detail", None) if detail: return str(detail) return str(e) diff --git a/litellm/router_strategy/adaptive_router/hooks.py b/litellm/router_strategy/adaptive_router/hooks.py index 709910753f2..c4a2eae1ef9 100644 --- a/litellm/router_strategy/adaptive_router/hooks.py +++ b/litellm/router_strategy/adaptive_router/hooks.py @@ -86,7 +86,7 @@ def _resolve_session_key(kwargs: dict[str, Any]) -> str | None: return hashlib.sha256(payload.encode("utf-8")).hexdigest() -def _last_user_content(messages: list[dict[str, Any]] | None) -> str | None: +def _last_user_content(messages: Sequence[Mapping[str, object]] | None) -> str | None: if not messages: return None for msg in reversed(messages): diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index d08afa8c1f6..250201a46d1 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -3,7 +3,7 @@ Auto-Routing Strategy that works with a Semantic Router Config """ import asyncio -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Optional from pydantic import BaseModel, ConfigDict @@ -158,7 +158,7 @@ class AutoRouter(CustomLogger): return await asyncio.shield(build_task) @staticmethod - def _extract_text_from_messages(messages: list[dict[str, Any]]) -> str: + def _extract_text_from_messages(messages: Sequence[Mapping[str, object]]) -> str: """ Extract text content from the last user message for routing. From dbf566873ea32a82a9915876615dad67febd2e57 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 10:28:06 +0000 Subject: [PATCH 24/47] refactor(types): keep agentops preset imports optional Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/presets/agentops.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/litellm/integrations/otel/presets/agentops.py b/litellm/integrations/otel/presets/agentops.py index 58123656caa..965213f2ee4 100644 --- a/litellm/integrations/otel/presets/agentops.py +++ b/litellm/integrations/otel/presets/agentops.py @@ -9,12 +9,9 @@ this preset registers a custom exporter (``kind="agentops"``) that mints the JWT worker thread, off any event loop — and caches it for the process lifetime. """ -from collections.abc import Sequence from typing import Any, Final import httpx -from opentelemetry.sdk.trace import ReadableSpan -from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult from pydantic import Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -74,7 +71,7 @@ def agentops_preset( ) -def _build_agentops_exporter(spec: ExporterSpec) -> SpanExporter: +def _build_agentops_exporter(spec: ExporterSpec) -> Any: """Factory for the ``agentops`` exporter kind: a lazy-auth OTLP/HTTP exporter.""" from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( OTLPSpanExporter, @@ -109,7 +106,7 @@ def _build_agentops_exporter(spec: ExporterSpec) -> SpanExporter: except Exception as e: verbose_logger.debug("AgentOps JWT fetch failed: %s", e) - def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: + def export(self, spans: Any) -> Any: self._ensure_authenticated() return super().export(spans) From 729a96e4eaff1d6fe4e68d5a1137ba83b07c0fe0 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 10:40:27 +0000 Subject: [PATCH 25/47] test: migrate openai, openai_like and openrouter legacy tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/llms/openai/evals/__init__.py | 1 - tests/test_litellm/llms/openai_like/embedding/__init__.py | 1 - tests/test_litellm/llms/openai_like/messages/__init__.py | 0 tests/test_litellm/llms/openrouter/image_edit/__init__.py | 0 .../guardrail_translation/test_embeddings_guardrail_handler.py | 0 .../llms/openai/evals/test_openai_evals_transformation.py | 0 .../llms/openai/image_generation/test_gpt_transformation.py | 0 .../image_generation/test_image_generation_guardrail_handler.py | 0 .../test_openai_image_generation_extra_headers.py | 0 .../llms/openai/speech/test_text_to_speech_guardrail_handler.py | 0 .../transcriptions/test_audio_transcription_guardrail_handler.py | 0 .../openai/transcriptions/test_transcription_duration_hidden.py | 0 .../llms/openai/transcriptions/test_whisper_transformation.py | 0 .../test_openai_vector_store_files_transformation.py | 0 .../vector_stores/test_openai_vector_stores_transformation.py | 0 .../llms/openai/videos/test_openai_video_transformation.py | 0 .../openai_like/chat/test_openai_like_chat_transformation.py | 0 .../llms/openai_like/embedding/test_openai_like_embedding.py | 0 .../test_openai_like_anthropic_messages_transformation.py | 0 .../llms/openrouter/chat/test_openrouter_chat_transformation.py | 0 .../image_edit/test_openrouter_image_edit_transformation.py | 0 .../image_generation/test_openrouter_image_gen_transformation.py | 0 .../llms/openrouter/test_openrouter_embedding_transformation.py | 0 .../llms/openrouter/test_openrouter_provider_routing.py | 0 24 files changed, 2 deletions(-) delete mode 100644 tests/test_litellm/llms/openai/evals/__init__.py delete mode 100644 tests/test_litellm/llms/openai_like/embedding/__init__.py delete mode 100644 tests/test_litellm/llms/openai_like/messages/__init__.py delete mode 100644 tests/test_litellm/llms/openrouter/image_edit/__init__.py rename tests/{test_litellm => unit}/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/evals/test_openai_evals_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/image_generation/test_gpt_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/image_generation/test_image_generation_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/image_generation/test_openai_image_generation_extra_headers.py (100%) rename tests/{test_litellm => unit}/llms/openai/speech/test_text_to_speech_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py (100%) rename tests/{test_litellm => unit}/llms/openai/transcriptions/test_transcription_duration_hidden.py (100%) rename tests/{test_litellm => unit}/llms/openai/transcriptions/test_whisper_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/vector_stores/test_openai_vector_stores_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai/videos/test_openai_video_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/chat/test_openai_like_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/embedding/test_openai_like_embedding.py (100%) rename tests/{test_litellm => unit}/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openrouter/chat/test_openrouter_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openrouter/test_openrouter_embedding_transformation.py (100%) rename tests/{test_litellm => unit}/llms/openrouter/test_openrouter_provider_routing.py (100%) diff --git a/tests/test_litellm/llms/openai/evals/__init__.py b/tests/test_litellm/llms/openai/evals/__init__.py deleted file mode 100644 index 47a8a2f0aed..00000000000 --- a/tests/test_litellm/llms/openai/evals/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""OpenAI Evals API tests""" diff --git a/tests/test_litellm/llms/openai_like/embedding/__init__.py b/tests/test_litellm/llms/openai_like/embedding/__init__.py deleted file mode 100644 index 2cb77227ed0..00000000000 --- a/tests/test_litellm/llms/openai_like/embedding/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Test module for OpenAI-like embedding handler diff --git a/tests/test_litellm/llms/openai_like/messages/__init__.py b/tests/test_litellm/llms/openai_like/messages/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/openrouter/image_edit/__init__.py b/tests/test_litellm/llms/openrouter/image_edit/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py b/tests/unit/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py rename to tests/unit/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py diff --git a/tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py b/tests/unit/llms/openai/evals/test_openai_evals_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py rename to tests/unit/llms/openai/evals/test_openai_evals_transformation.py diff --git a/tests/test_litellm/llms/openai/image_generation/test_gpt_transformation.py b/tests/unit/llms/openai/image_generation/test_gpt_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/image_generation/test_gpt_transformation.py rename to tests/unit/llms/openai/image_generation/test_gpt_transformation.py diff --git a/tests/test_litellm/llms/openai/image_generation/test_image_generation_guardrail_handler.py b/tests/unit/llms/openai/image_generation/test_image_generation_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/image_generation/test_image_generation_guardrail_handler.py rename to tests/unit/llms/openai/image_generation/test_image_generation_guardrail_handler.py diff --git a/tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py b/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py similarity index 100% rename from tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py rename to tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py diff --git a/tests/test_litellm/llms/openai/speech/test_text_to_speech_guardrail_handler.py b/tests/unit/llms/openai/speech/test_text_to_speech_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/speech/test_text_to_speech_guardrail_handler.py rename to tests/unit/llms/openai/speech/test_text_to_speech_guardrail_handler.py diff --git a/tests/test_litellm/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py b/tests/unit/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py rename to tests/unit/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py diff --git a/tests/test_litellm/llms/openai/transcriptions/test_transcription_duration_hidden.py b/tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py similarity index 100% rename from tests/test_litellm/llms/openai/transcriptions/test_transcription_duration_hidden.py rename to tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py diff --git a/tests/test_litellm/llms/openai/transcriptions/test_whisper_transformation.py b/tests/unit/llms/openai/transcriptions/test_whisper_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/transcriptions/test_whisper_transformation.py rename to tests/unit/llms/openai/transcriptions/test_whisper_transformation.py diff --git a/tests/test_litellm/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 similarity index 100% rename from tests/test_litellm/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py rename to tests/unit/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py diff --git a/tests/test_litellm/llms/openai/vector_stores/test_openai_vector_stores_transformation.py b/tests/unit/llms/openai/vector_stores/test_openai_vector_stores_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/vector_stores/test_openai_vector_stores_transformation.py rename to tests/unit/llms/openai/vector_stores/test_openai_vector_stores_transformation.py diff --git a/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py b/tests/unit/llms/openai/videos/test_openai_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py rename to tests/unit/llms/openai/videos/test_openai_video_transformation.py diff --git a/tests/test_litellm/llms/openai_like/chat/test_openai_like_chat_transformation.py b/tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai_like/chat/test_openai_like_chat_transformation.py rename to tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py diff --git a/tests/test_litellm/llms/openai_like/embedding/test_openai_like_embedding.py b/tests/unit/llms/openai_like/embedding/test_openai_like_embedding.py similarity index 100% rename from tests/test_litellm/llms/openai_like/embedding/test_openai_like_embedding.py rename to tests/unit/llms/openai_like/embedding/test_openai_like_embedding.py diff --git a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py rename to tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py diff --git a/tests/test_litellm/llms/openrouter/chat/test_openrouter_chat_transformation.py b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/openrouter/chat/test_openrouter_chat_transformation.py rename to tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py diff --git a/tests/test_litellm/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py b/tests/unit/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py rename to tests/unit/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py diff --git a/tests/test_litellm/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py b/tests/unit/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py similarity index 100% rename from tests/test_litellm/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py rename to tests/unit/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py diff --git a/tests/test_litellm/llms/openrouter/test_openrouter_embedding_transformation.py b/tests/unit/llms/openrouter/test_openrouter_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/openrouter/test_openrouter_embedding_transformation.py rename to tests/unit/llms/openrouter/test_openrouter_embedding_transformation.py diff --git a/tests/test_litellm/llms/openrouter/test_openrouter_provider_routing.py b/tests/unit/llms/openrouter/test_openrouter_provider_routing.py similarity index 100% rename from tests/test_litellm/llms/openrouter/test_openrouter_provider_routing.py rename to tests/unit/llms/openrouter/test_openrouter_provider_routing.py From 7cf9a3035ce130fb4a07a02c7e139df6915fd378 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 10:59:02 +0000 Subject: [PATCH 26/47] test: migrate phase 16 legacy tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/rust_bridge/__init__.py | 0 tests/unit/rust_bridge/ocr/__init__.py | 0 tests/unit/rust_bridge/ocr/test_route_host.py | 85 ++ tests/unit/rust_bridge/responses/__init__.py | 0 .../rust_bridge/responses/test_route_host.py | 57 + tests/unit/sandbox/test_e2b_sandbox.py | 318 +++++ .../unit/sandbox/test_opensandbox_sandbox.py | 647 ++++++++++ tests/unit/sandbox/test_sandbox_tools.py | 181 +++ tests/unit/skills/test_skills_main.py | 57 + .../test_enforce_model_rate_limits.py | 468 ++++++++ .../test_router/test_io_token_rate_limits.py | 1041 +++++++++++++++++ .../types/llms/test_types_llms_bedrock.py | 46 + .../unit/types/llms/test_types_llms_openai.py | 591 ++++++++++ .../types/proxy/policy_engine/__init__.py | 0 .../policy_engine/test_pipeline_types.py | 168 +++ .../proxy/policy_engine/test_policy_types.py | 15 + .../policy_engine/test_resolver_types.py | 115 ++ tests/unit/videos/__init__.py | 0 tests/unit/videos/test_main.py | 455 +++++++ tests/unit/videos/test_utils.py | 181 +++ 20 files changed, 4425 insertions(+) create mode 100644 tests/unit/rust_bridge/__init__.py create mode 100644 tests/unit/rust_bridge/ocr/__init__.py create mode 100644 tests/unit/rust_bridge/ocr/test_route_host.py create mode 100644 tests/unit/rust_bridge/responses/__init__.py create mode 100644 tests/unit/rust_bridge/responses/test_route_host.py create mode 100644 tests/unit/sandbox/test_e2b_sandbox.py create mode 100644 tests/unit/sandbox/test_opensandbox_sandbox.py create mode 100644 tests/unit/sandbox/test_sandbox_tools.py create mode 100644 tests/unit/skills/test_skills_main.py create mode 100644 tests/unit/test_router/test_enforce_model_rate_limits.py create mode 100644 tests/unit/test_router/test_io_token_rate_limits.py create mode 100644 tests/unit/types/llms/test_types_llms_bedrock.py create mode 100644 tests/unit/types/llms/test_types_llms_openai.py create mode 100644 tests/unit/types/proxy/policy_engine/__init__.py create mode 100644 tests/unit/types/proxy/policy_engine/test_pipeline_types.py create mode 100644 tests/unit/types/proxy/policy_engine/test_policy_types.py create mode 100644 tests/unit/types/proxy/policy_engine/test_resolver_types.py create mode 100644 tests/unit/videos/__init__.py create mode 100644 tests/unit/videos/test_main.py create mode 100644 tests/unit/videos/test_utils.py diff --git a/tests/unit/rust_bridge/__init__.py b/tests/unit/rust_bridge/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/rust_bridge/ocr/__init__.py b/tests/unit/rust_bridge/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/rust_bridge/ocr/test_route_host.py b/tests/unit/rust_bridge/ocr/test_route_host.py new file mode 100644 index 00000000000..699492e4424 --- /dev/null +++ b/tests/unit/rust_bridge/ocr/test_route_host.py @@ -0,0 +1,85 @@ +from typing import Final + +import pytest + +import litellm +from litellm.rust_bridge.ocr.route_host import UpstreamFailure, map_failure +from litellm.rust_bridge.ocr.route_host import response as build_ocr_response +from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest + +REQUEST: Final = LiteLLMOcrRequest( + model="mistral/mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/file.pdf"}, + api_key="test-key", + api_base=None, + timeout=None, + custom_llm_provider=None, + extra_headers=None, + kwargs={"req_format": "markdown"}, +) + + +class RustUpstreamError(Exception): + def __init__(self, status: int, body: str, headers: tuple[tuple[str, str], ...]) -> None: + super().__init__(status, body) + self.headers: Final = list(headers) + + +class RustFormatError(Exception): + ocr_request_format_error: Final = True + + +def test_rust_ocr_response_retains_provider_native_response(): + provider_response = {"status": "succeeded", "analyzeResult": {"content": "native"}} + response = build_ocr_response( + { + "pages": [], + "model": "prebuilt-layout", + "document_annotation": None, + "usage_info": {"pages_processed": 0}, + "object": "ocr", + "provider_native_response": provider_response, + } + ) + + assert response.get_provider_native_response() == provider_response + assert response.model_dump().get("provider_native_response") is None + + +def test_map_failure_builds_public_error_from_upstream_status_and_headers() -> None: + error: Final = RustUpstreamError(429, '{"message": "slow down"}', (("retry-after", "7"),)) + + public_error: Final = map_failure(error, REQUEST, "mistral") + + assert isinstance(public_error, litellm.RateLimitError) + assert public_error.status_code == 429 + assert public_error.response.headers["retry-after"] == "7" + assert public_error.response.text == '{"message": "slow down"}' + assert public_error.__context__ is error + assert public_error.llm_provider == "mistral" + + +def test_map_failure_maps_upstream_401_to_authentication_error() -> None: + error: Final = RustUpstreamError(401, '{"message": "Unauthorized"}', ()) + + public_error: Final = map_failure(error, REQUEST, "mistral") + + assert isinstance(public_error, litellm.AuthenticationError) + assert public_error.status_code == 401 + assert public_error.response.text == '{"message": "Unauthorized"}' + assert public_error.__context__ is error + + +def test_map_failure_leaves_non_upstream_errors_unwrapped() -> None: + error: Final = RuntimeError("bridge exploded") + + public_error: Final = map_failure(error, REQUEST, "mistral") + + assert not isinstance(public_error, UpstreamFailure) + assert isinstance(public_error, litellm.APIConnectionError) + assert "bridge exploded" in str(public_error) + + +def test_map_failure_reports_invalid_request_format_as_unsupported_params() -> None: + with pytest.raises(litellm.UnsupportedParamsError, match="Invalid `req_format`: 'markdown'"): + raise map_failure(RustFormatError(), REQUEST, "mistral") diff --git a/tests/unit/rust_bridge/responses/__init__.py b/tests/unit/rust_bridge/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/rust_bridge/responses/test_route_host.py b/tests/unit/rust_bridge/responses/test_route_host.py new file mode 100644 index 00000000000..49bf19e7d8a --- /dev/null +++ b/tests/unit/rust_bridge/responses/test_route_host.py @@ -0,0 +1,57 @@ +from types import MappingProxyType +from typing import Final + +import pytest +from pydantic import ValidationError + +from litellm.rust_bridge.responses.route_host import arguments, response +from litellm.rust_bridge.responses.entrypoints import LiteLLMResponsesRequest +from litellm.types.llms.openai import ResponsesAPIResponse + + +def test_response_validates_into_the_public_responses_model() -> None: + built: Final = response( + MappingProxyType( + { + "id": "resp_native", + "object": "response", + "created_at": 1, + "model": "gpt-4o", + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_native", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "native", "annotations": []}], + } + ], + } + ) + ) + + assert isinstance(built, ResponsesAPIResponse) + assert built.id == "resp_native" + assert built.output[0].content[0].text == "native" + + +def test_response_rejects_a_payload_missing_required_fields() -> None: + with pytest.raises(ValidationError): + response(MappingProxyType({"object": "response"})) + + +def test_arguments_are_the_public_kwargs_view() -> None: + kwargs: Final = MappingProxyType({"litellm_metadata": {"user_id": "u"}}) + request: Final = LiteLLMResponsesRequest( + model="gpt-4o", + input="hi", + stream=None, + api_key=None, + api_base=None, + custom_llm_provider="openai", + extra_headers=None, + kwargs=kwargs, + ) + + assert arguments(request) is kwargs diff --git a/tests/unit/sandbox/test_e2b_sandbox.py b/tests/unit/sandbox/test_e2b_sandbox.py new file mode 100644 index 00000000000..e01b9120416 --- /dev/null +++ b/tests/unit/sandbox/test_e2b_sandbox.py @@ -0,0 +1,318 @@ +""" +Tests for the e2b code execution sandbox primitive. + +Unit tests inject a fake async HTTP client (dependency injection, no +monkeypatching) and assert request shapes and result mapping. Real-network +integration tests live in tests/integration/sandbox/test_e2b_sandbox.py. +""" + +import json + +import httpx +import pytest + +import litellm +from litellm.llms.base_llm.sandbox.transformation import ContainerHandle +from litellm.llms.e2b.sandbox.transformation import ( + MAX_OUTPUT_BYTES, + E2BSandboxConfig, +) + + +class FakeResponse: + def __init__(self, *, json_data=None, lines=None, status_code=200): + self._json = json_data + self._lines = lines or [] + self.status_code = status_code + + def json(self): + return self._json + + async def aiter_lines(self): + for line in self._lines: + yield line + + +class FakeHTTPClient: + """Records outbound requests and returns canned responses keyed by URL.""" + + def __init__( + self, + *, + create_json=None, + execute_lines=None, + delete_status=204, + execute_raises=None, + ): + self.create_json = create_json or { + "sandboxID": "sbx_123", + "domain": "e2b.app", + "envdAccessToken": "tok_abc", + } + self.execute_lines = execute_lines or [] + self.delete_status = delete_status + self.execute_raises = execute_raises + self.calls = [] + + async def post(self, url, headers=None, json=None, stream=False, **kwargs): + self.calls.append(("POST", url, headers, json)) + if url.endswith("/sandboxes"): + return FakeResponse(json_data=self.create_json) + if url.endswith("/execute"): + if self.execute_raises is not None: + raise self.execute_raises + return FakeResponse(lines=self.execute_lines) + raise AssertionError(f"unexpected POST {url}") + + async def delete(self, url, headers=None, **kwargs): + self.calls.append(("DELETE", url, headers, None)) + if not (200 <= self.delete_status < 300): + raise httpx.HTTPStatusError( + f"status {self.delete_status}", + request=httpx.Request("DELETE", url), + response=httpx.Response(self.delete_status), + ) + return FakeResponse(status_code=self.delete_status) + + +# ---------- pure parser ---------- + + +def test_parse_lines_stdout_and_count(): + lines = [ + json.dumps({"type": "stdout", "text": "6\n", "timestamp": 1}), + json.dumps({"type": "number_of_executions", "execution_count": 1}), + ] + result = E2BSandboxConfig._parse_lines(lines) + assert result.stdout == "6\n" + assert result.execution_count == 1 + assert result.error is None + + +def test_parse_lines_error_surfaces_name_and_traceback(): + lines = [ + json.dumps( + { + "type": "error", + "name": "ZeroDivisionError", + "value": "division by zero", + "traceback": "Traceback (most recent call last): ...", + } + ) + ] + result = E2BSandboxConfig._parse_lines(lines) + assert result.error["name"] == "ZeroDivisionError" + assert "Traceback" in result.error["traceback"] + + +def test_parse_lines_result_carries_png(): + lines = [ + json.dumps({"type": "result", "png": "BASE64DATA", "is_main_result": True}) + ] + result = E2BSandboxConfig._parse_lines(lines) + assert result.results and result.results[0]["png"] == "BASE64DATA" + assert "type" not in result.results[0] + + +# ---------- request shapes ---------- + + +@pytest.mark.asyncio +async def test_template_flows_into_create_request_as_templateID(): + client = FakeHTTPClient() + cfg = E2BSandboxConfig() + handle = await cfg.acreate_sandbox( + template="my-custom-template", api_key="e2b_key", client=client + ) + + method, url, headers, body = client.calls[0] + assert method == "POST" + assert url.endswith("/sandboxes") + assert body["templateID"] == "my-custom-template" # not "template" + assert body["secure"] is True + assert headers["X-API-Key"] == "e2b_key" + assert handle.id == "sbx_123" + assert handle._hidden_params["envd_access_token"] == "tok_abc" + + +@pytest.mark.asyncio +async def test_create_defaults_template_when_omitted(): + client = FakeHTTPClient() + await E2BSandboxConfig().acreate_sandbox(api_key="e2b_key", client=client) + _, _, _, body = client.calls[0] + assert body["templateID"] == "code-interpreter-v1" + + +@pytest.mark.asyncio +async def test_run_code_targets_jupyter_host_with_access_token(): + client = FakeHTTPClient( + execute_lines=[json.dumps({"type": "stdout", "text": "42\n", "timestamp": 1})] + ) + handle = ContainerHandle(id="sbx_xyz", provider="e2b", domain="e2b.app") + handle._hidden_params = {"envd_access_token": "tok_run"} + + result = await E2BSandboxConfig().arun_code( + container=handle, code="print(6*7)", client=client + ) + + method, url, headers, body = client.calls[0] + assert url == "https://49999-sbx_xyz.e2b.app/execute" + assert headers["X-Access-Token"] == "tok_run" + assert body["code"] == "print(6*7)" + assert result.stdout.strip() == "42" + + +@pytest.mark.asyncio +async def test_delete_issues_delete_to_sandbox_id(): + client = FakeHTTPClient(delete_status=204) + handle = ContainerHandle(id="sbx_del", provider="e2b", domain="e2b.app") + handle._hidden_params = {"api_key": "e2b_key"} + + ok = await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) + + method, url, headers, _ = client.calls[0] + assert method == "DELETE" + assert url.endswith("/sandboxes/sbx_del") + assert ok is True + + +@pytest.mark.asyncio +async def test_delete_returns_false_on_404(): + client = FakeHTTPClient(delete_status=404) + handle = ContainerHandle(id="sbx_gone", provider="e2b", domain="e2b.app") + handle._hidden_params = {"api_key": "e2b_key"} + ok = await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) + assert ok is False + + +# ---------- ephemeral teardown ---------- + + +@pytest.mark.asyncio +async def test_code_interpreter_tool_deletes_even_when_run_raises(): + client = FakeHTTPClient(execute_raises=RuntimeError("boom")) + + with pytest.raises(RuntimeError, match="boom"): + await litellm.acode_interpreter_tool( + provider="e2b", code="1/0", api_key="e2b_key", client=client + ) + + methods = [c[0] for c in client.calls] + urls = [c[1] for c in client.calls] + assert methods == ["POST", "POST", "DELETE"] # create, run(raises), delete + assert urls[0].endswith("/sandboxes") + assert urls[1].endswith("/execute") + assert urls[2].endswith("/sandboxes/sbx_123") + + +# ---------- correctness guards ---------- + + +@pytest.mark.asyncio +async def test_delete_reraises_non_404_http_error(): + client = FakeHTTPClient(delete_status=500) + handle = ContainerHandle(id="sbx_err", provider="e2b", domain="e2b.app") + handle._hidden_params = {"api_key": "e2b_key"} + with pytest.raises(httpx.HTTPStatusError): + await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) + + +@pytest.mark.asyncio +async def test_create_preserves_explicit_zero_timeout(): + client = FakeHTTPClient() + await E2BSandboxConfig().acreate_sandbox( + timeout=0, api_key="e2b_key", client=client + ) + _, _, _, body = client.calls[0] + assert body["timeout"] == 0 + + +@pytest.mark.asyncio +async def test_run_code_rejects_bare_id_without_access_token(): + client = FakeHTTPClient() + with pytest.raises(ValueError, match="access token"): + await E2BSandboxConfig().arun_code( + container="sbx_no_token", code="print(1)", client=client + ) + assert client.calls == [] # never reached the network + + +def test_parse_lines_skips_non_json_lines(): + lines = [ + "not-json-heartbeat", + json.dumps({"type": "stdout", "text": "ok\n"}), + "", + "{partial", + ] + result = E2BSandboxConfig._parse_lines(lines) + assert result.stdout == "ok\n" + assert result.error is None + + +@pytest.mark.asyncio +async def test_run_code_aborts_on_output_over_cap(): + big_line = "x" * (MAX_OUTPUT_BYTES + 1) + client = FakeHTTPClient(execute_lines=[big_line]) + handle = ContainerHandle(id="sbx_big", provider="e2b", domain="e2b.app") + handle._hidden_params = {"envd_access_token": "tok"} + with pytest.raises(ValueError, match="exceeded"): + await E2BSandboxConfig().arun_code( + container=handle, code="print('x'*999)", client=client + ) + + +# ---------- public entrypoints ---------- + + +@pytest.mark.asyncio +async def test_public_lifecycle_create_run_delete(): + client = FakeHTTPClient( + execute_lines=[json.dumps({"type": "stdout", "text": "42\n"})] + ) + container = await litellm.acreate_sandbox( + provider="e2b", api_key="e2b_key", client=client + ) + assert container.id == "sbx_123" + + result = await litellm.arun_code( + provider="e2b", + container=container, + api_key="e2b_key", + code="print(6*7)", + client=client, + ) + assert result.stdout.strip() == "42" + + assert ( + await litellm.adelete_sandbox( + provider="e2b", container=container, api_key="e2b_key", client=client + ) + is True + ) + + +@pytest.mark.asyncio +async def test_unsupported_provider_raises(): + with pytest.raises(ValueError, match="not-a-provider' is not a valid SandboxProviders"): + await litellm.acreate_sandbox(provider="not-a-provider") + + +# ---------- api_base override ---------- + + +@pytest.mark.asyncio +async def test_create_uses_api_base_override(): + client = FakeHTTPClient() + await E2BSandboxConfig().acreate_sandbox( + api_base="http://my-sandbox:8080", api_key="k", client=client + ) + _, url, _, _ = client.calls[0] + assert url == "http://my-sandbox:8080/sandboxes" + + +@pytest.mark.asyncio +async def test_create_defaults_to_e2b_api_base(): + client = FakeHTTPClient() + await E2BSandboxConfig().acreate_sandbox(api_key="k", client=client) + _, url, _, _ = client.calls[0] + assert url == "https://api.e2b.app/sandboxes" diff --git a/tests/unit/sandbox/test_opensandbox_sandbox.py b/tests/unit/sandbox/test_opensandbox_sandbox.py new file mode 100644 index 00000000000..2928dea100e --- /dev/null +++ b/tests/unit/sandbox/test_opensandbox_sandbox.py @@ -0,0 +1,647 @@ +import json + +import httpx +import pytest + +import litellm +from litellm.llms.base_llm.sandbox.transformation import ContainerHandle +from litellm.llms.opensandbox.sandbox.transformation import ( + MAX_OUTPUT_BYTES, + OPEN_SANDBOX_DEFAULT_TEMPLATE, + OpenSandboxSandboxConfig, +) +from litellm.utils import ProviderConfigManager + +TEST_API_BASE = "https://sandbox.test/v1" + + +def http_status_error(status_code, url="http://test"): + return httpx.HTTPStatusError( + f"status {status_code}", + request=httpx.Request("GET", url), + response=httpx.Response(status_code), + ) + + +def sse(data): + return f"data: {json.dumps(data)}" + + +class FakeResponse: + def __init__(self, *, json_data=None, lines=None, status_code=200): + self._json = json_data + self._lines = lines or [] + self.status_code = status_code + + def json(self): + return self._json + + def raise_for_status(self): + if self.status_code >= 400: + raise http_status_error(self.status_code) + + async def aiter_lines(self): + for line in self._lines: + yield line + + +class FakeHTTPClient: + def __init__( + self, + *, + create_json=None, + sandbox_states=None, + endpoint_json=None, + endpoint_responses=None, + execute_lines=None, + delete_status=204, + execute_raises=None, + ): + self.create_json = create_json or { + "id": "osb_123", + "status": {"state": "Running"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/opt/code-interpreter/code-interpreter.sh"], + } + self.sandbox_states = list( + sandbox_states + or [ + { + "id": "osb_123", + "status": {"state": "Running"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/opt/code-interpreter/code-interpreter.sh"], + } + ] + ) + self.endpoint_json = endpoint_json or { + "endpoint": "execd.local:44772", + "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, + } + self.endpoint_responses = ( + list(endpoint_responses) if endpoint_responses is not None else None + ) + self.execute_lines = execute_lines or [] + self.delete_status = delete_status + self.execute_raises = execute_raises + self.calls = [] + + async def post(self, url, headers=None, json=None, stream=False, **kwargs): + self.calls.append(("POST", url, headers, json, {"stream": stream})) + if url.endswith("/sandboxes"): + return FakeResponse(json_data=self.create_json) + if url.endswith("/code"): + if self.execute_raises is not None: + raise self.execute_raises + return FakeResponse(lines=self.execute_lines) + raise AssertionError(f"unexpected POST {url}") + + async def get(self, url, headers=None, params=None, **kwargs): + self.calls.append(("GET", url, headers, None, params)) + if "/endpoints/44772" in url: + if self.endpoint_responses is not None and self.endpoint_responses: + response = self.endpoint_responses.pop(0) + if isinstance(response, Exception): + raise response + if isinstance(response, FakeResponse): + return response + return FakeResponse(json_data=response) + return FakeResponse(json_data=self.endpoint_json) + if "/sandboxes/" in url: + state = self.sandbox_states.pop(0) + return FakeResponse(json_data=state) + raise AssertionError(f"unexpected GET {url}") + + async def delete(self, url, headers=None, **kwargs): + self.calls.append(("DELETE", url, headers, None, None)) + if not (200 <= self.delete_status < 300): + raise http_status_error(self.delete_status, url) + return FakeResponse(status_code=self.delete_status) + + +def test_parse_sse_lines_maps_output_result_count_and_error(): + lines = [ + sse({"type": "stdout", "text": "hello\n"}), + sse({"type": "stderr", "text": "warn\n"}), + sse({"type": "result", "results": {"text/plain": "4"}}), + sse({"type": "execution_count", "execution_count": 7}), + sse( + { + "type": "error", + "error": { + "ename": "ValueError", + "evalue": "bad", + "traceback": ["Traceback"], + }, + } + ), + ] + + result = OpenSandboxSandboxConfig._parse_lines(lines) + + assert result.stdout == "hello\n" + assert result.stderr == "warn\n" + assert result.results == [{"text/plain": "4"}] + assert result.execution_count == 7 + assert result.error == { + "name": "ValueError", + "value": "bad", + "traceback": ["Traceback"], + } + + +def test_parse_sse_lines_skips_non_json_and_control_lines(): + lines = [ + "event: message", + "not-json", + "", + sse({"type": "stdout", "text": "ok\n"}), + ] + + result = OpenSandboxSandboxConfig._parse_lines(lines) + + assert result.stdout == "ok\n" + assert result.error is None + + +def test_parse_sse_lines_maps_fallback_shapes(): + lines = [ + "data:", + sse(["not-a-dict"]), + sse({"code": "BadRequest", "message": "nope"}), + sse({"type": "result", "text/plain": "4"}), + sse({"type": "error", "name": "RuntimeError", "text": "boom"}), + sse({"type": "execution_count", "execution_count": "8"}), + ] + + result = OpenSandboxSandboxConfig._parse_lines(lines) + + assert result.results == [{"text/plain": "4"}] + assert result.execution_count == 8 + assert result.error == { + "name": "BadRequest", + "value": "nope", + "traceback": [], + } + fallback_error = OpenSandboxSandboxConfig._parse_lines( + [sse({"type": "error", "name": "RuntimeError", "text": "boom"})] + ) + assert fallback_error.error == { + "name": "RuntimeError", + "value": "boom", + "traceback": [], + } + empty_string_error = OpenSandboxSandboxConfig._parse_lines( + [ + sse( + { + "type": "error", + "error": { + "ename": "", + "name": "FallbackName", + "evalue": "", + "value": "fallback value", + "traceback": [], + }, + } + ) + ] + ) + assert empty_string_error.error == { + "name": "", + "value": "", + "traceback": [], + } + + +def test_static_helpers_cover_defaults_and_fallbacks(monkeypatch): + def fake_secret(key): + if key == "OPEN_SANDBOX_API_KEY": + return "env-key" + if key == "OPEN_SANDBOX_API_BASE": + return TEST_API_BASE + return None + + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.get_secret_str", + fake_secret, + ) + config = OpenSandboxSandboxConfig() + handle = ContainerHandle(id="osb", provider="opensandbox", domain="http://x/v1") + + assert config.validate_environment() == "env-key" + assert config.validate_environment(api_key="") == "" + assert config._api_key(api_key=None, handle=handle) == "env-key" + + handle._hidden_params = {"api_key": "stored-key"} + assert config._api_key(api_key=None, handle=handle) == "stored-key" + assert config._http(None) is not None + + body = config._create_body( + template=None, + timeout=None, + allow_internet_access=False, + metadata=None, + env_vars=None, + resource_limits=None, + resource_requests=None, + entrypoint=None, + network_policy={"egress": [{"domain": "example.com"}]}, + secure_access=True, + ) + assert body["networkPolicy"] == {"egress": [{"domain": "example.com"}]} + assert body["secureAccess"] is True + + other_body = config._create_body( + template=None, + timeout=None, + allow_internet_access=False, + metadata=None, + env_vars=None, + resource_limits=None, + resource_requests=None, + entrypoint=None, + network_policy=None, + secure_access=False, + ) + assert body["resourceLimits"] is not other_body["resourceLimits"] + + assert config._sandbox_state(None) is None + assert config._sandbox_state({"status": "Running"}) is None + assert config._as_str_dict(None) == {} + assert config._endpoint_base_url("http://execd.local", "https://api/v1") == ( + "http://execd.local" + ) + assert config._api_base(None) == TEST_API_BASE + assert config._api_base("https://direct.test/v1/") == "https://direct.test/v1" + assert config._as_int("9") == 9 + assert config._as_int("nope") is None + assert config._as_int(None) is None + assert isinstance( + ProviderConfigManager.get_provider_sandbox_config("opensandbox"), + OpenSandboxSandboxConfig, + ) + + +def test_api_base_requires_kwarg_or_env(monkeypatch): + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.get_secret_str", + lambda key: None, + ) + + with pytest.raises(ValueError, match="api_base is required"): + OpenSandboxSandboxConfig._api_base(None) + + +@pytest.mark.asyncio +async def test_create_posts_default_body_and_omits_empty_api_key(): + client = FakeHTTPClient() + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, client=client + ) + + method, url, headers, body, _ = client.calls[0] + assert method == "POST" + assert url == f"{TEST_API_BASE}/sandboxes" + assert "OPEN-SANDBOX-API-KEY" not in headers + assert body["image"] == {"uri": OPEN_SANDBOX_DEFAULT_TEMPLATE} + assert body["entrypoint"] == ["/opt/code-interpreter/code-interpreter.sh"] + assert body["timeout"] == 300 + assert body["resourceLimits"] == {"cpu": "1", "memory": "2Gi"} + assert body["networkPolicy"] == {"defaultAction": "deny", "egress": []} + assert handle.id == "osb_123" + assert handle._hidden_params["execd_endpoint"] == "execd.local:44772" + + +@pytest.mark.asyncio +async def test_create_can_opt_into_internet_access(): + client = FakeHTTPClient() + + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + allow_internet_access=True, + client=client, + ) + + _, _, _, body, _ = client.calls[0] + assert "networkPolicy" not in body + + +@pytest.mark.asyncio +async def test_create_custom_options_poll_and_endpoint_resolution(): + client = FakeHTTPClient( + create_json={ + "id": "osb_pending", + "status": {"state": "Pending"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/bin/sh"], + }, + sandbox_states=[ + { + "id": "osb_pending", + "status": {"state": "Running"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/bin/sh"], + } + ], + ) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + template="custom/image:latest", + timeout=600, + allow_internet_access=False, + api_key="osb-key", + api_base="https://sandbox.example/v1", + metadata={"suite": "unit"}, + env_vars={"PYTHONUNBUFFERED": "1"}, + resource_limits={"cpu": "500m", "memory": "512Mi"}, + resource_requests={"cpu": "250m", "memory": "256Mi"}, + entrypoint=["/bin/sh", "-lc", "sleep 3600"], + use_server_proxy=True, + client=client, + ) + + _, create_url, create_headers, body, _ = client.calls[0] + _, poll_url, poll_headers, _, _ = client.calls[1] + _, endpoint_url, endpoint_headers, _, endpoint_params = client.calls[2] + + assert create_url == "https://sandbox.example/v1/sandboxes" + assert create_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" + assert body["image"] == {"uri": "custom/image:latest"} + assert body["entrypoint"] == ["/bin/sh", "-lc", "sleep 3600"] + assert body["metadata"] == {"suite": "unit"} + assert body["env"] == {"PYTHONUNBUFFERED": "1"} + assert body["resourceLimits"] == {"cpu": "500m", "memory": "512Mi"} + assert body["resourceRequests"] == {"cpu": "250m", "memory": "256Mi"} + assert body["networkPolicy"] == {"defaultAction": "deny", "egress": []} + assert poll_url == "https://sandbox.example/v1/sandboxes/osb_pending" + assert poll_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" + assert endpoint_url.endswith("/sandboxes/osb_pending/endpoints/44772") + assert endpoint_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" + assert endpoint_params == {"use_server_proxy": True} + assert handle.id == "osb_pending" + + +@pytest.mark.asyncio +async def test_create_waits_across_pending_state(monkeypatch): + client = FakeHTTPClient( + create_json={ + "id": "osb_pending", + "status": {"state": "Pending"}, + "createdAt": "2026-01-01T00:00:00Z", + }, + sandbox_states=[ + {"id": "osb_pending", "status": {"state": "Pending"}}, + {"id": "osb_pending", "status": {"state": "Running"}}, + ], + ) + sleeps = [] + + async def fake_sleep(interval): + sleeps.append(interval) + + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.asyncio.sleep", fake_sleep + ) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + ready_timeout=1, + poll_interval=0.01, + client=client, + ) + + assert handle.id == "osb_pending" + assert sleeps == [0.01] + + +@pytest.mark.asyncio +async def test_create_raises_for_terminal_state(): + client = FakeHTTPClient( + create_json={"id": "osb_failed", "status": {"state": "Pending"}}, + sandbox_states=[ + {"id": "osb_failed", "status": {"state": "Failed"}}, + ], + ) + + with pytest.raises(ValueError, match="entered Failed"): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, client=client + ) + + +@pytest.mark.asyncio +async def test_create_times_out_waiting_for_running(): + client = FakeHTTPClient( + create_json={"id": "osb_slow", "status": {"state": "Pending"}}, + sandbox_states=[ + {"id": "osb_slow", "status": {"state": "Pending"}}, + ], + ) + + with pytest.raises(TimeoutError, match="was not Running"): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + ready_timeout=0, + poll_interval=0, + client=client, + ) + + +@pytest.mark.asyncio +async def test_create_waits_for_endpoint_resolution(monkeypatch): + client = FakeHTTPClient( + endpoint_responses=[ + http_status_error(404, f"{TEST_API_BASE}/sandboxes/osb_123"), + { + "endpoint": "execd.local:44772", + "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, + }, + ], + ) + sleeps = [] + + async def fake_sleep(interval): + sleeps.append(interval) + + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.asyncio.sleep", fake_sleep + ) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + ready_timeout=1, + poll_interval=0.01, + client=client, + ) + + endpoint_calls = [call for call in client.calls if "/endpoints/44772" in call[1]] + assert handle._hidden_params["execd_endpoint"] == "execd.local:44772" + assert len(endpoint_calls) == 2 + assert sleeps == [0.01] + + +@pytest.mark.asyncio +async def test_create_raises_when_endpoint_is_missing(): + client = FakeHTTPClient(endpoint_json={"headers": {"X": "y"}}) + + with pytest.raises(TimeoutError, match=r"execd endpoint.*not ready"): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, ready_timeout=0, client=client + ) + + +@pytest.mark.asyncio +async def test_create_reraises_non_404_endpoint_error(): + client = FakeHTTPClient(endpoint_responses=[http_status_error(500)]) + + with pytest.raises(httpx.HTTPStatusError): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, client=client + ) + + +@pytest.mark.asyncio +async def test_run_code_resolves_bare_id_and_posts_sse_request(): + client = FakeHTTPClient( + execute_lines=[ + sse({"type": "stdout", "text": "42\n"}), + ] + ) + + result = await OpenSandboxSandboxConfig().arun_code( + container="osb_bare", + code="print(6*7)", + language="python", + api_key="", + api_base="http://sandbox.local/v1", + client=client, + ) + + endpoint_call = client.calls[0] + run_call = client.calls[1] + assert endpoint_call[0] == "GET" + assert ( + endpoint_call[1] == "http://sandbox.local/v1/sandboxes/osb_bare/endpoints/44772" + ) + assert run_call[0] == "POST" + assert run_call[1] == "http://execd.local:44772/code" + assert run_call[2]["X-EXECD-ACCESS-TOKEN"] == "execd-token" + assert run_call[3] == { + "code": "print(6*7)", + "context": {"language": "python"}, + } + assert run_call[4] == {"stream": True} + assert result.stdout == "42\n" + + +@pytest.mark.asyncio +async def test_run_code_uses_https_for_scheme_less_endpoint_when_api_base_is_https(): + client = FakeHTTPClient() + handle = ContainerHandle( + id="osb_https", provider="opensandbox", domain="https://sandbox.example/v1" + ) + handle._hidden_params = { + "execd_endpoint": "execd.example/route/44772", + "execd_headers": {}, + } + + await OpenSandboxSandboxConfig().arun_code( + container=handle, code="print(1)", client=client + ) + + assert client.calls[0][1] == "https://execd.example/route/44772/code" + + +@pytest.mark.asyncio +async def test_run_code_aborts_on_output_over_cap(): + client = FakeHTTPClient(execute_lines=["x" * (MAX_OUTPUT_BYTES + 1)]) + handle = ContainerHandle(id="osb_big", provider="opensandbox", domain="http://x/v1") + handle._hidden_params = {"execd_endpoint": "execd.local:44772", "execd_headers": {}} + + with pytest.raises(ValueError, match="exceeded"): + await OpenSandboxSandboxConfig().arun_code( + container=handle, code="print('x')", client=client + ) + + +@pytest.mark.asyncio +async def test_delete_returns_false_on_404(): + client = FakeHTTPClient(delete_status=404) + + ok = await OpenSandboxSandboxConfig().adelete_sandbox( + container="osb_gone", + api_key="", + api_base="http://sandbox.local/v1", + client=client, + ) + + assert ok is False + + +@pytest.mark.asyncio +async def test_delete_reraises_non_404_http_error(): + client = FakeHTTPClient(delete_status=500) + + with pytest.raises(httpx.HTTPStatusError): + await OpenSandboxSandboxConfig().adelete_sandbox( + container="osb_err", + api_key="", + api_base="http://sandbox.local/v1", + client=client, + ) + + +@pytest.mark.asyncio +async def test_public_lifecycle_create_run_delete(): + client = FakeHTTPClient( + execute_lines=[ + sse({"type": "stdout", "text": "42\n"}), + ] + ) + + container = await litellm.acreate_sandbox( + provider="opensandbox", api_key="", api_base=TEST_API_BASE, client=client + ) + result = await litellm.arun_code( + provider="opensandbox", + container=container, + code="print(6*7)", + api_key="", + client=client, + ) + ok = await litellm.adelete_sandbox( + provider="opensandbox", + container=container, + api_key="", + client=client, + ) + + assert container.id == "osb_123" + assert result.stdout == "42\n" + assert ok is True + + +@pytest.mark.asyncio +async def test_code_interpreter_tool_deletes_even_when_run_raises(): + client = FakeHTTPClient(execute_raises=RuntimeError("boom")) + + with pytest.raises(RuntimeError, match="boom"): + await litellm.acode_interpreter_tool( + provider="opensandbox", + code="1/0", + api_key="", + api_base=TEST_API_BASE, + client=client, + ) + + assert [call[0] for call in client.calls] == ["POST", "GET", "POST", "DELETE"] + assert client.calls[0][1].endswith("/sandboxes") + assert client.calls[1][1].endswith("/endpoints/44772") + assert client.calls[2][1].endswith("/code") + assert client.calls[3][1].endswith("/sandboxes/osb_123") diff --git a/tests/unit/sandbox/test_sandbox_tools.py b/tests/unit/sandbox/test_sandbox_tools.py new file mode 100644 index 00000000000..06136534b13 --- /dev/null +++ b/tests/unit/sandbox/test_sandbox_tools.py @@ -0,0 +1,181 @@ +"""Unit tests for the sandbox-tool registry.""" + +from litellm.sandbox import sandbox_tools + + +def _reset(): + sandbox_tools.clear_sandbox_tools() + + +def test_register_resolves_provider_key_and_base(): + _reset() + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": { + "sandbox_provider": "e2b", + "api_key": "sk-literal", + "api_base": "https://sandbox.internal", + }, + } + ] + ) + + resolved = sandbox_tools.resolve_sandbox_tool("e2b_default") + assert resolved == { + "sandbox_provider": "e2b", + "api_key": "sk-literal", + "api_base": "https://sandbox.internal", + } + _reset() + + +def test_register_clears_stale_entries_on_reload(): + """A tool removed from the config must not survive a re-registration.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "old", + "litellm_params": {"sandbox_provider": "e2b"}, + } + ] + ) + assert sandbox_tools.resolve_sandbox_tool("old") is not None + + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "new", + "litellm_params": {"sandbox_provider": "e2b"}, + } + ] + ) + + assert sandbox_tools.resolve_sandbox_tool("new") is not None + assert ( + sandbox_tools.resolve_sandbox_tool("old") is None + ), "stale tool must be gone after the config is reloaded" + _reset() + + +def test_register_empty_list_clears_removed_tools(): + """Reloading a config with sandbox_tools removed (the proxy passes an empty + list) must drop previously registered credentials from the process.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": {"sandbox_provider": "e2b", "api_key": "sk-x"}, + } + ] + ) + assert sandbox_tools.resolve_sandbox_tool("e2b_default") is not None + + sandbox_tools.register_sandbox_tools([]) + + assert ( + sandbox_tools.resolve_sandbox_tool("e2b_default") is None + ), "removing sandbox_tools from config must clear stale credentials" + _reset() + + +def test_register_resolves_secret_from_env(monkeypatch): + _reset() + monkeypatch.setenv("MY_SANDBOX_KEY", "sk-from-env") + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": { + "sandbox_provider": "e2b", + "api_key": "os.environ/MY_SANDBOX_KEY", + }, + } + ] + ) + + resolved = sandbox_tools.resolve_sandbox_tool("e2b_default") + assert resolved is not None + assert resolved["api_key"] == "sk-from-env" + assert resolved["api_base"] is None + _reset() + + +def test_resolve_unknown_returns_none(): + _reset() + assert sandbox_tools.resolve_sandbox_tool("nope") is None + + +def test_register_skips_malformed_entries_without_crashing(): + """A single malformed entry (missing sandbox_tool_name, or not a dict) must + not crash registration during proxy startup/hot-reload; valid entries in the + same list must still register.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + {"litellm_params": {"sandbox_provider": "e2b"}}, # missing name + "not-a-dict", # wrong type + {"sandbox_tool_name": "", "litellm_params": {}}, # empty name + { + "sandbox_tool_name": "good", + "litellm_params": {"sandbox_provider": "e2b"}, + }, + ] + ) + + assert sandbox_tools.resolve_sandbox_tool("good") is not None + assert sandbox_tools.resolve_sandbox_tool("") is None + assert set(sandbox_tools._SANDBOX_TOOL_REGISTRY) == {"good"} + _reset() + + +def test_register_skips_entry_missing_sandbox_provider(): + """An entry with a name but no sandbox_provider must be skipped at + registration so it cannot later resolve and call acreate_sandbox(provider=None), + which fails with a cryptic runtime error instead of a clear startup warning.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + {"sandbox_tool_name": "no_provider", "litellm_params": {"api_key": "sk-x"}}, + { + "sandbox_tool_name": "null_provider", + "litellm_params": {"sandbox_provider": None}, + }, + { + "sandbox_tool_name": "good", + "litellm_params": {"sandbox_provider": "e2b"}, + }, + ] + ) + + assert sandbox_tools.resolve_sandbox_tool("no_provider") is None + assert sandbox_tools.resolve_sandbox_tool("null_provider") is None + assert set(sandbox_tools._SANDBOX_TOOL_REGISTRY) == {"good"} + _reset() + + +def test_register_swaps_registry_atomically(): + """register_sandbox_tools must replace the registry in one rebind so a + concurrent resolve never observes a half-populated or transiently empty + registry between clearing and repopulating.""" + _reset() + sandbox_tools.register_sandbox_tools( + [{"sandbox_tool_name": "a", "litellm_params": {"sandbox_provider": "e2b"}}] + ) + before = sandbox_tools._SANDBOX_TOOL_REGISTRY + + sandbox_tools.register_sandbox_tools( + [ + {"sandbox_tool_name": "b", "litellm_params": {"sandbox_provider": "e2b"}}, + {"sandbox_tool_name": "c", "litellm_params": {"sandbox_provider": "e2b"}}, + ] + ) + after = sandbox_tools._SANDBOX_TOOL_REGISTRY + + assert after is not before, "the registry must be replaced, not mutated in place" + assert set(after) == {"b", "c"} + assert "a" not in after + _reset() diff --git a/tests/unit/skills/test_skills_main.py b/tests/unit/skills/test_skills_main.py new file mode 100644 index 00000000000..e1c66c8d9ea --- /dev/null +++ b/tests/unit/skills/test_skills_main.py @@ -0,0 +1,57 @@ +from unittest.mock import MagicMock + +import litellm.skills.main as skills_main +from litellm.types.utils import LlmProviders + + +def test_create_skill_forwards_description_and_instructions_from_top_level_kwargs( + monkeypatch, +) -> None: + """The REST /v1/skills form endpoint passes description/instructions as top-level + kwargs (not extra_body). Regression for a bug where the litellm_proxy dispatch + branch of create_skill() dropped both, so every LiteLLM-hosted skill was created + with description=None and instructions=None regardless of what the caller sent.""" + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill( + display_title="Document Translator", + description="Converts files from one language into another", + instructions="Take an uploaded document and produce it in the target language", + custom_llm_provider=LlmProviders.LITELLM_PROXY.value, + ) + + assert handler.create_skill_handler.call_args.kwargs["description"] == ( + "Converts files from one language into another" + ) + assert handler.create_skill_handler.call_args.kwargs["instructions"] == ( + "Take an uploaded document and produce it in the target language" + ) + + +def test_create_skill_forwards_description_and_instructions_from_extra_body(monkeypatch) -> None: + """The SDK convention (see tests/proxy_unit_tests/test_skills_db.py) nests them under + extra_body instead of passing them as top-level kwargs; both paths must reach the DB.""" + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill( + display_title="Warehouse SQL Analyst", + extra_body={"description": "Runs SQL against the inventory database", "instructions": "Summarize results"}, + custom_llm_provider=LlmProviders.LITELLM_PROXY.value, + ) + + assert handler.create_skill_handler.call_args.kwargs["description"] == ( + "Runs SQL against the inventory database" + ) + assert handler.create_skill_handler.call_args.kwargs["instructions"] == "Summarize results" + + +def test_create_skill_without_description_or_instructions_passes_none(monkeypatch) -> None: + handler = MagicMock() + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + skills_main.create_skill(display_title="Bare Skill", custom_llm_provider=LlmProviders.LITELLM_PROXY.value) + + assert handler.create_skill_handler.call_args.kwargs["description"] is None + assert handler.create_skill_handler.call_args.kwargs["instructions"] is None diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py new file mode 100644 index 00000000000..7577064b7f9 --- /dev/null +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -0,0 +1,468 @@ +""" +Tests for enforce_model_rate_limits feature. + +This feature allows users to enforce TPM/RPM limits set on model deployments +regardless of the routing strategy being used. +""" + +import asyncio +from datetime import timedelta +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import litellm +from litellm import Router +from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( + ModelRateLimitingCheck, +) + +TPM_DEPLOYMENT = { + "tpm": 1000, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "replica-test-id"}, + "model_name": "test-model", +} + + +def _dual_cache_with_local_tpm(local_tpm: int, redis_cache: MagicMock | None) -> DualCache: + dual_cache = DualCache(redis_cache=redis_cache) + check = ModelRateLimitingCheck(dual_cache=dual_cache) + now = litellm.utils.get_utc_datetime() + for minute in (now, now + timedelta(minutes=1)): + tpm_key, _ = check._get_cache_keys(TPM_DEPLOYMENT, minute.strftime("%H-%M")) + dual_cache.set_cache(key=tpm_key, value=local_tpm, local_only=True) + return dual_cache + + +class TestModelRateLimitingCheck: + """Test the ModelRateLimitingCheck class directly.""" + + def test_get_deployment_limits_from_top_level(self): + """Test extracting limits from top-level deployment config.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "tpm": 1000, + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm == 1000 + assert rpm == 10 + + def test_get_deployment_limits_from_litellm_params(self): + """Test extracting limits from litellm_params.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4", "tpm": 2000, "rpm": 20}, + "model_info": {"id": "test-id"}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm == 2000 + assert rpm == 20 + + def test_get_deployment_limits_from_model_info(self): + """Test extracting limits from model_info.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id", "tpm": 3000, "rpm": 30}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm == 3000 + assert rpm == 30 + + def test_get_deployment_limits_none_when_not_set(self): + """Test that None is returned when limits are not set.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm is None + assert rpm is None + + def test_pre_call_check_allows_request_when_no_limits(self): + """Test that requests are allowed when no limits are set.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + result = check.pre_call_check(deployment) + assert result == deployment + + def test_pre_call_check_raises_rate_limit_error_when_over_rpm(self): + """Test that RateLimitError is raised when RPM limit is exceeded.""" + mock_cache = MagicMock() + mock_cache.increment_cache.return_value = 11 # Over limit after increment + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(deployment) + + assert "RPM limit=10" in str(exc_info.value) + assert "current usage=11" in str(exc_info.value) + + def test_pre_call_check_allows_request_under_limit(self): + """Test that requests are allowed when under the limit.""" + mock_cache = MagicMock() + mock_cache.increment_cache.return_value = 6 + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + result = check.pre_call_check(deployment) + assert result == deployment + + def test_pre_call_check_raises_rate_limit_error_when_over_tpm(self): + """Test that RateLimitError is raised when TPM limit is exceeded.""" + mock_cache = MagicMock() + mock_cache.get_cache.return_value = 1000 # Already at limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "tpm": 1000, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(deployment) + + assert "TPM limit=1000" in str(exc_info.value) + assert "current usage=1000" in str(exc_info.value) + + def test_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): + redis_cache = MagicMock() + redis_cache.get_cache.return_value = 1000 + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(TPM_DEPLOYMENT) + + assert "current usage=1000" in str(exc_info.value) + + @pytest.mark.parametrize( + "redis_get", [MagicMock(return_value=None), MagicMock(side_effect=RedisCircuitBreakerOpenError())] + ) + def test_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): + redis_cache = MagicMock() + redis_cache.get_cache = redis_get + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(TPM_DEPLOYMENT) + + assert "current usage=1000" in str(exc_info.value) + + def test_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open(self): + redis_cache = MagicMock() + redis_cache.get_cache.side_effect = RedisCircuitBreakerOpenError() + redis_cache.increment_cache.return_value = 2 + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) + + assert "RPM limit=1" in str(exc_info.value) + + def test_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) + deployment = {**TPM_DEPLOYMENT, "rpm": 1} + + assert check.pre_call_check(deployment) == deployment + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(deployment) + + assert "RPM limit=1" in str(exc_info.value) + + def test_log_success_event_increments_cache(self): + """Test that log_success_event correctly increments the cache.""" + mock_cache = MagicMock() + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + kwargs = { + "standard_logging_object": { + "model_id": "test-id", + "total_tokens": 50, + "hidden_params": {"litellm_model_name": "gpt-4"}, + } + } + + check.log_success_event(kwargs, None, None, None) + + # Verify increment_cache was called + mock_cache.increment_cache.assert_called_once() + _, kwarg_params = mock_cache.increment_cache.call_args + assert "test-id:gpt-4:tpm:" in kwarg_params["key"] + assert kwarg_params["value"] == 50 + + +class TestModelRateLimitingCheckAsync: + """Test async methods of ModelRateLimitingCheck.""" + + @pytest.mark.asyncio + async def test_async_pre_call_check_allows_request_when_no_limits(self): + """Test that requests are allowed when no limits are set (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + result = await check.async_pre_call_check(deployment) + assert result == deployment + + @pytest.mark.asyncio + async def test_async_pre_call_check_raises_rate_limit_error_when_over_rpm(self): + """Test that RateLimitError is raised when RPM limit is exceeded (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_increment_cache = AsyncMock(return_value=11) # Over limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "RPM limit=10" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_pre_call_check_allows_request_under_limit(self): + """Test that requests are allowed when under the limit (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_increment_cache = AsyncMock(return_value=6) + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + result = await check.async_pre_call_check(deployment) + assert result == deployment + + @pytest.mark.asyncio + async def test_async_pre_call_check_raises_rate_limit_error_when_over_tpm(self): + """Test that RateLimitError is raised when TPM limit is exceeded (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=1000) # Already at limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "tpm": 1000, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "TPM limit=1000" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): + redis_cache = MagicMock() + redis_cache.async_get_cache = AsyncMock(return_value=1000) + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(TPM_DEPLOYMENT) + + assert "current usage=1000" in str(exc_info.value) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "redis_get", [AsyncMock(return_value=None), AsyncMock(side_effect=RedisCircuitBreakerOpenError())] + ) + async def test_async_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): + redis_cache = MagicMock() + redis_cache.async_get_cache = redis_get + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(TPM_DEPLOYMENT) + + assert "current usage=1000" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open( + self, + ): + redis_cache = MagicMock() + redis_cache.async_get_cache = AsyncMock(side_effect=RedisCircuitBreakerOpenError()) + redis_cache.async_increment = AsyncMock(return_value=2) + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) + + assert "RPM limit=1" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): + check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) + deployment = {**TPM_DEPLOYMENT, "rpm": 1} + + assert await check.async_pre_call_check(deployment) == deployment + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "RPM limit=1" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_log_success_event_increments_cache(self): + """Test that async_log_success_event correctly increments the cache.""" + mock_cache = MagicMock() + mock_cache.async_increment_cache = AsyncMock() + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + kwargs = { + "standard_logging_object": { + "model_id": "test-id", + "total_tokens": 50, + "hidden_params": {"litellm_model_name": "gpt-4"}, + } + } + + await check.async_log_success_event(kwargs, None, None, None) + + # Verify async_increment_cache was called + mock_cache.async_increment_cache.assert_called_once() + _, kwarg_params = mock_cache.async_increment_cache.call_args + assert "test-id:gpt-4:tpm:" in kwarg_params["key"] + assert kwarg_params["value"] == 50 + + +class TestRouterWithEnforceModelRateLimits: + """Test Router integration with enforce_model_rate_limits.""" + + def test_router_initializes_with_enforce_model_rate_limits(self): + """Test that Router properly initializes the ModelRateLimitingCheck.""" + model_list = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "test"}, + "rpm": 10, + } + ] + + router = Router( + model_list=model_list, + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + # Check that the callback was added + assert router.optional_callbacks is not None + assert len(router.optional_callbacks) == 1 + assert isinstance(router.optional_callbacks[0], ModelRateLimitingCheck) + + def test_router_optional_callbacks_contains_model_rate_limiting(self): + """Test that ModelRateLimitingCheck is in the callbacks list.""" + model_list = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "test"}, + "rpm": 10, + } + ] + + Router( + model_list=model_list, + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + # Find the ModelRateLimitingCheck in litellm.callbacks + found = False + for callback in litellm.callbacks: + if isinstance(callback, ModelRateLimitingCheck): + found = True + break + + assert found, "ModelRateLimitingCheck should be in litellm.callbacks" + + +class TestModelRateLimitConcurrency: + """Test that RPM rate limiting is atomic under concurrent requests.""" + + @pytest.mark.asyncio + async def test_concurrent_requests_respect_rpm_limit(self): + """ + Fire 4 concurrent async requests with RPM limit of 2. + Exactly 2 should succeed and 2 should raise RateLimitError. + + This test validates the atomic increment-first pattern: + the old check-then-increment pattern would let 3+ through + due to a race condition on the local cache read. + """ + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + + deployment = { + "rpm": 2, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "concurrent-test-id"}, + "model_name": "test-model", + } + + async def attempt_request(): + return await check.async_pre_call_check(deployment) + + results = await asyncio.gather( + *[attempt_request() for _ in range(4)], + return_exceptions=True, + ) + + successes = [r for r in results if not isinstance(r, Exception)] + failures = [r for r in results if isinstance(r, litellm.RateLimitError)] + + assert len(successes) == 2, f"Expected 2 successes, got {len(successes)}" + assert len(failures) == 2, f"Expected 2 rate limit errors, got {len(failures)}" diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py new file mode 100644 index 00000000000..a5a68271111 --- /dev/null +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -0,0 +1,1041 @@ +""" +Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits). +""" + +import asyncio + +import pytest + +import litellm +from litellm import Router +from litellm.caching.dual_cache import DualCache +from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( + ITPM_CACHE_KEY, + ITPM_RESERVED_KEY, + OTPM_CACHE_KEY, + OTPM_RESERVED_KEY, + _reservation_value, + _resolve_max_tokens, + async_io_token_pre_call_check, + async_io_token_reconcile_success, + build_io_token_rate_limit_headers, + deployment_has_io_token_limits, + get_io_token_rate_limit_request_kwargs, + io_token_reconcile_success, + io_token_refund_failure, + refund_stale_reservation_before_retry, + set_io_token_rate_limit_request_kwargs, +) +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( + ModelRateLimitingCheck, +) +from litellm.types.utils import ModelResponse, Usage + + +class TestIOTokenRateLimitHelpers: + def test_deployment_has_io_token_limits(self): + assert deployment_has_io_token_limits({"litellm_params": {"itpm": 100, "otpm": 50}}) + assert not deployment_has_io_token_limits({"litellm_params": {"model": "x"}}) + + def test_reservation_value_minimal_when_estimate_fails(self): + # A failed/empty estimate (0) must reserve a minimal slot, not the + # entire limit - otherwise one request whose estimate failed fills + # the whole bucket and blocks every concurrent request until it + # completes and reconciles. + assert _reservation_value(0, 100) == 1 + assert _reservation_value(0, 1) == 1 + # A real non-zero estimate is reserved as-is. + assert _reservation_value(42, 100) == 42 + + def test_resolve_max_tokens_respects_explicit_zero(self): + deployment = {"litellm_params": {"model": "openai/gpt-4o-mini"}} + # An explicit max_tokens=0 is honored, not replaced by the model default. + assert _resolve_max_tokens({"max_tokens": 0}, deployment) == 0 + # max_completion_tokens is the fallback only when max_tokens is absent. + assert _resolve_max_tokens({"max_completion_tokens": 12}, deployment) == 12 + assert _resolve_max_tokens({"max_output_tokens": 9}, deployment) == 9 + + def test_build_io_token_rate_limit_headers(self): + headers = build_io_token_rate_limit_headers( + itpm_limit=200, + otpm_limit=40, + current_itpm=15, + current_otpm=4, + ) + assert headers["x-ratelimit-limit-input-tokens"] == 200 + assert headers["x-ratelimit-remaining-input-tokens"] == 185 + assert headers["x-ratelimit-limit-output-tokens"] == 40 + assert headers["x-ratelimit-remaining-output-tokens"] == 36 + + +class TestModelRateLimitingCheckIOTokens: + @pytest.mark.asyncio + async def test_itpm_reservation_and_reconcile(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 100, + "otpm": 50, + }, + "model_info": {"id": "io-test-id"}, + "model_name": "opus", + } + + request_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 10, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + await check.async_pre_call_check(deployment) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + otpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" + + kwargs = { + "standard_logging_object": { + "model_id": "io-test-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, + "metadata": dict(request_kwargs["metadata"]), + }, + "metadata": request_kwargs["metadata"], + } + response = ModelResponse( + choices=[ + { + "message": {"role": "assistant", "content": "hi"}, + "index": 0, + "finish_reason": "stop", + } + ], + usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8), + ) + await check.async_log_success_event(kwargs, response, None, None) + + current_itpm = await dual_cache.async_get_cache(key=itpm_key) + current_otpm = await dual_cache.async_get_cache(key=otpm_key) + # ITPM tracks input tokens only (billable prompt tokens), not output. + assert current_itpm == 5 + assert current_otpm == 3 + + @pytest.mark.asyncio + async def test_itpm_limit_raises_429(self): + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 5, + }, + "model_info": {"id": "io-limit-id"}, + "model_name": "opus", + } + + # ITPM enforces input tokens only; the prompt alone must exceed the limit, + # a large max_tokens must not contribute to the ITPM reservation. + set_io_token_rate_limit_request_kwargs( + { + "messages": [ + { + "role": "user", + "content": "hello world this is a longer prompt that exceeds the tiny itpm limit", + } + ], + "max_tokens": 10, + "metadata": {}, + } + ) + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "ITPM limit=5" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_otpm_atomic_reservation_no_overshoot_under_concurrency(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + otpm_limit = 10 + max_tokens = 4 + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "otpm": otpm_limit, + }, + "model_info": {"id": "io-otpm-race-id"}, + "model_name": "opus", + } + + set_io_token_rate_limit_request_kwargs( + { + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": max_tokens, + "metadata": {}, + } + ) + + async def _attempt(): + try: + await check.async_pre_call_check(deployment) + return True + except litellm.RateLimitError: + return False + + results = await asyncio.gather(*[_attempt() for _ in range(8)]) + successes = sum(1 for r in results if r) + + minute = get_utc_datetime().strftime("%H-%M") + otpm_key = f"global_router:io-otpm-race-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" + current_otpm = await dual_cache.async_get_cache(key=otpm_key) + + # Atomic reservation must never let concurrent requests overshoot the limit. + assert current_otpm is not None + assert current_otpm <= otpm_limit + assert successes == otpm_limit // max_tokens + assert current_otpm == successes * max_tokens + + @pytest.mark.asyncio + async def test_itpm_estimate_failure_reserves_minimal_not_full_limit(self): + """ + When input-token estimation yields 0 (no messages/prompt/input field, + unsupported model, tokenizer error), the reservation must be a + minimal 1 token, not the entire itpm limit. Otherwise the first + request whose estimate fails fills the whole bucket and every + concurrent request is rejected until it completes - effectively + serializing traffic to the deployment. + """ + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + itpm_limit = 5 + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": itpm_limit, + }, + "model_info": {"id": "io-itpm-estimate-fail-id"}, + "model_name": "opus", + } + + # No messages/prompt/input field -> _estimate_input_tokens returns 0. + set_io_token_rate_limit_request_kwargs( + { + "max_tokens": 5, + "metadata": {}, + } + ) + + async def _attempt(): + try: + await check.async_pre_call_check(deployment) + return True + except litellm.RateLimitError: + return False + + results = await asyncio.gather(*[_attempt() for _ in range(8)]) + successes = sum(1 for r in results if r) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-itpm-estimate-fail-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + current_itpm = await dual_cache.async_get_cache(key=itpm_key) + + # A minimal 1-token reservation per request lets itpm_limit concurrent + # requests through, instead of a single request starving the rest. + assert current_itpm is not None + assert current_itpm <= itpm_limit + assert successes == itpm_limit + + @pytest.mark.asyncio + async def test_reservation_read_prefers_top_level_metadata_over_litellm_params(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-lp-id:bedrock_mantle/test:itpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) + + # Production kwargs commonly carry litellm_params.metadata; the stashed + # reservation lives in the top-level metadata and must still be found. + kwargs = { + "standard_logging_object": { + "model_id": "io-lp-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, + "metadata": {}, + }, + "metadata": {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key}, + "litellm_params": {"metadata": {"user_api_key_hash": "abc123"}}, + } + await check.async_log_failure_event(kwargs, None, None, None) + + current = await dual_cache.async_get_cache(key=itpm_key) + assert current == 0 + + @pytest.mark.asyncio + async def test_reconcile_tracks_actual_usage_when_estimate_zero(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 100, + }, + "model_info": {"id": "io-zero-est-id"}, + "model_name": "opus", + } + + request_kwargs = {"max_tokens": 5, "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + await check.async_pre_call_check(deployment) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-zero-est-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + # A failed/zero estimate reserves a minimal 1 token, not the full + # itpm limit, so it doesn't starve concurrent requests. + assert await dual_cache.async_get_cache(key=itpm_key) == 1 + + kwargs = { + "standard_logging_object": { + "model_id": "io-zero-est-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, + "metadata": dict(request_kwargs["metadata"]), + }, + "metadata": request_kwargs["metadata"], + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=7, completion_tokens=0, total_tokens=7), + ) + await check.async_log_success_event(kwargs, response, None, None) + + assert await dual_cache.async_get_cache(key=itpm_key) == 7 + + @pytest.mark.asyncio + async def test_zero_estimate_reserves_minimal_capacity_before_reconcile(self): + """ + A zero/failed estimate reserves a minimal 1 token rather than the + full itpm limit, so up to itpm_limit such calls are allowed + concurrently instead of the first one claiming the entire bucket. + """ + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 2, + }, + "model_info": {"id": "io-zero-cap-id"}, + "model_name": "opus", + } + request_kwargs = {"max_tokens": 5, "metadata": {}} + set_io_token_rate_limit_request_kwargs(request_kwargs) + await check.async_pre_call_check(deployment) + + # Second zero-estimate call still fits within the itpm=2 limit. + set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) + await check.async_pre_call_check(deployment) + + # A third exceeds the limit and is rejected. + set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) + with pytest.raises(litellm.RateLimitError): + await check.async_pre_call_check(deployment) + + @pytest.mark.asyncio + async def test_explicit_zero_max_tokens_does_not_reserve_otpm(self): + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "otpm": 5, + }, + "model_info": {"id": "io-zero-output-id"}, + "model_name": "opus", + } + zero_output_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 0, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(zero_output_kwargs) + await check.async_pre_call_check(deployment) + + zero_output_otpm_key = zero_output_kwargs["metadata"][OTPM_CACHE_KEY] + assert zero_output_kwargs["metadata"][OTPM_RESERVED_KEY] == 0 + assert (await dual_cache.async_get_cache(key=zero_output_otpm_key) or 0) == 0 + + normal_output_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 5, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(normal_output_kwargs) + await check.async_pre_call_check(deployment) + + normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY] + assert await dual_cache.async_get_cache(key=normal_output_otpm_key) == 5 + + def test_sync_io_pre_call_reserves_and_reconciles(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 100, + "otpm": 50, + }, + "model_info": {"id": "io-sync-id"}, + "model_name": "opus", + } + request_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 10, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + check.pre_call_check(deployment) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + otpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" + kwargs = { + "standard_logging_object": { + "model_id": "io-sync-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, + "metadata": dict(request_kwargs["metadata"]), + }, + "metadata": request_kwargs["metadata"], + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8), + ) + check.log_success_event(kwargs, response, None, None) + + assert dual_cache.get_cache(key=itpm_key) == 5 + assert dual_cache.get_cache(key=otpm_key) == 3 + + @pytest.mark.asyncio + async def test_reconcile_runs_via_success_event_without_model_id(self): + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + itpm_key = "global_router:io-noid:bedrock_mantle/test:itpm:12-34" + await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + + # standard_logging_object has no model_id (only the TPM path needs it); + # IO reconciliation must still run off the stashed cache key. + kwargs = { + "standard_logging_object": { + "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, + "metadata": {}, + }, + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + ) + await check.async_log_success_event(kwargs, response, None, None) + + assert await dual_cache.async_get_cache(key=itpm_key) == 3 + + @pytest.mark.asyncio + async def test_failure_clears_reservation_so_retry_is_not_poisoned(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-first:bedrock_mantle/test:itpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) + + # Shared request metadata carrying the first (IO) deployment's reservation. + metadata = {ITPM_RESERVED_KEY: 8, ITPM_CACHE_KEY: itpm_key} + fail_kwargs = { + "metadata": metadata, + "standard_logging_object": { + "model_id": "io-first", + "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, + "metadata": {}, + }, + } + await check.async_log_failure_event(fail_kwargs, None, None, None) + + assert await dual_cache.async_get_cache(key=itpm_key) == 0 + assert ITPM_RESERVED_KEY not in metadata + assert ITPM_CACHE_KEY not in metadata + + # Retry succeeds on a non-IO fallback deployment reusing the same metadata. + retry_kwargs = { + "metadata": metadata, + "standard_logging_object": { + "model_id": "non-io-second", + "hidden_params": {"litellm_model_name": "openai/gpt-4o-mini"}, + "metadata": {}, + "total_tokens": 12, + }, + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=6, completion_tokens=6, total_tokens=12), + ) + await check.async_log_success_event(retry_kwargs, response, None, None) + + # The first deployment's ITPM counter is not driven negative... + assert await dual_cache.async_get_cache(key=itpm_key) == 0 + # ...and the non-IO deployment's TPM usage is tracked normally. + tpm_key = f"non-io-second:openai/gpt-4o-mini:tpm:{minute}" + assert await dual_cache.async_get_cache(key=tpm_key) == 12 + + @pytest.mark.asyncio + async def test_stale_reservation_refunded_before_retry_overwrites_it(self): + """ + A retry reuses the same mutable kwargs dict for the next deployment. + If deployment A's failure event hasn't run yet (e.g. it was scheduled + as a background task) when the retry calls + set_io_token_rate_limit_request_kwargs for deployment B, the router + must first synchronously refund + clear A's reservation via + refund_stale_reservation_before_retry - otherwise A's counter stays + elevated by the reservation until its TTL expires, and the + now-orphaned sentinels must not leak into B's accounting either. + """ + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + minute = get_utc_datetime().strftime("%H-%M") + itpm_key_a = f"global_router:io-retry-a:bedrock_mantle/test-a:itpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key_a, value=9, ttl=60) + + # Deployment A's still-unreconciled reservation, stashed on the shared + # kwargs dict the retry loop reuses. + shared_kwargs = {"metadata": {ITPM_RESERVED_KEY: 9, ITPM_CACHE_KEY: itpm_key_a}} + + # Router calls this before overwriting kwargs for deployment B's attempt - + # simulating the fix landing ahead of set_io_token_rate_limit_request_kwargs. + refund_stale_reservation_before_retry(dual_cache, shared_kwargs) + + # A's reservation is refunded immediately, not left stranded for a + # background failure task that may run arbitrarily later (or never, + # if the sentinels get cleared out from under it first). + assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 + assert ITPM_RESERVED_KEY not in shared_kwargs["metadata"] + assert ITPM_CACHE_KEY not in shared_kwargs["metadata"] + + # A's own (now-late) failure event finds nothing left to refund and + # is a safe no-op, since the sentinels were already cleared above. + io_token_refund_failure(dual_cache, shared_kwargs) + assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 + + # The retry proceeds to stash deployment B's own reservation on the + # same dict; it starts clean, unaffected by A's cleared sentinels. + set_io_token_rate_limit_request_kwargs(shared_kwargs) + itpm_key_b = f"global_router:io-retry-b:bedrock_mantle/test-b:itpm:{minute}" + shared_kwargs["metadata"][ITPM_RESERVED_KEY] = 4 + shared_kwargs["metadata"][ITPM_CACHE_KEY] = itpm_key_b + await dual_cache.async_increment_cache(key=itpm_key_b, value=4, ttl=60) + assert await dual_cache.async_get_cache(key=itpm_key_b) == 4 + + @pytest.mark.asyncio + async def test_client_supplied_reservation_keys_are_stripped(self): + # metadata is caller-controlled; the server-only reservation sentinels + # must be removed before the router captures the request kwargs. + forged = { + "metadata": {ITPM_RESERVED_KEY: 999999, ITPM_CACHE_KEY: "attacker:key:itpm:00-00"}, + "litellm_metadata": {OTPM_RESERVED_KEY: 7}, + "litellm_params": {"metadata": {OTPM_CACHE_KEY: "attacker:key:otpm:00-00"}}, + } + set_io_token_rate_limit_request_kwargs(forged) + stored = get_io_token_rate_limit_request_kwargs() + + assert ITPM_RESERVED_KEY not in stored["metadata"] + assert ITPM_CACHE_KEY not in stored["metadata"] + assert OTPM_RESERVED_KEY not in stored["litellm_metadata"] + assert OTPM_CACHE_KEY not in stored["litellm_params"]["metadata"] + + @pytest.mark.asyncio + async def test_forged_reservation_cannot_decrement_counter(self): + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + victim_key = "global_router:victim:model:itpm:00-00" + await dual_cache.async_increment_cache(key=victim_key, value=100, ttl=60) + + # A caller forges a reservation pointing at another deployment's counter. + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 100, ITPM_CACHE_KEY: victim_key}, + "standard_logging_object": { + "model_id": "m", + "hidden_params": {"litellm_model_name": "model"}, + "metadata": {}, + "total_tokens": 2, + }, + } + # The router sanitizes the request kwargs before the call runs. + set_io_token_rate_limit_request_kwargs(kwargs) + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ) + await check.async_log_success_event(kwargs, response, None, None) + + # The forged reservation was stripped, so the victim counter is untouched. + assert await dual_cache.async_get_cache(key=victim_key) == 100 + + @pytest.mark.asyncio + async def test_otpm_reservation_error_rolls_back_itpm(self): + from litellm.utils import get_utc_datetime + + class _OtpmFailCache(DualCache): + async def async_increment_cache(self, key, **kwargs): + if ":otpm:" in key: + raise RuntimeError("transient cache error") + return await super().async_increment_cache(key=key, **kwargs) + + dual_cache = _OtpmFailCache() + deployment = { + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 1000, + "otpm": 1000, + }, + "model_info": {"id": "io-rollback-id"}, + "model_name": "opus", + } + set_io_token_rate_limit_request_kwargs( + { + "messages": [{"role": "user", "content": "hello world"}], + "max_tokens": 5, + "metadata": {}, + } + ) + + with pytest.raises(RuntimeError): + await async_io_token_pre_call_check(dual_cache, deployment) + + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-rollback-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + # A transient OTPM error must release the ITPM reservation, not leak it. + assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 + + @pytest.mark.asyncio + async def test_reconcile_clears_stash_even_when_increment_errors(self): + class _ItpmFailCache(DualCache): + async def async_increment_cache(self, key, **kwargs): + if ":itpm:" in key: + raise RuntimeError("transient cache error") + return await super().async_increment_cache(key=key, **kwargs) + + dual_cache = _ItpmFailCache() + metadata = {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: "global_router:x:model:itpm:00-00"} + kwargs = {"metadata": metadata} + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + ) + + with pytest.raises(RuntimeError): + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + # The stash is cleared even though reconciliation raised, so a duplicate + # success event can't re-process it. + assert ITPM_RESERVED_KEY not in metadata + assert ITPM_CACHE_KEY not in metadata + + @pytest.mark.asyncio + async def test_io_conflict_warning_not_collapsed_for_missing_model_id(self, caplog): + import logging + + check = ModelRateLimitingCheck(dual_cache=DualCache()) + deployment = { + "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100, "tpm": 1000}, + "model_info": {}, + } + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + check._warn_io_token_and_tpm_rpm_coexist_once(deployment) + check._warn_io_token_and_tpm_rpm_coexist_once(deployment) + + warnings = [r for r in caplog.records if "both limit types are enforced" in r.message] + # id-less deployments are not collapsed onto a single dedup key. + assert len(warnings) == 2 + + @pytest.mark.asyncio + async def test_missing_deployment_id_skips_io_reservation(self): + dual_cache = DualCache() + deployment = { + "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100}, + "model_info": {}, # no id -> cannot build a per-deployment cache key + "model_name": "opus", + } + request_kwargs = { + "messages": [{"role": "user", "content": "hello world"}], + "max_tokens": 5, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + + result = await async_io_token_pre_call_check(dual_cache, deployment) + + assert result is deployment + # No reservation is stashed, so nothing lands in a shared None:None bucket. + assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] + + @pytest.mark.asyncio + async def test_reconcile_uses_reservation_minute_key(self): + dual_cache = DualCache() + # Reservation was made on a fixed minute key; a call that finishes in a + # later minute must reconcile against that same key, never a key built + # from the response-time minute. + itpm_key = "global_router:io-min-id:bedrock_mantle/test:itpm:99-99" + await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + + kwargs = {"metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}} + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=4, completion_tokens=0, total_tokens=4), + ) + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + assert await dual_cache.async_get_cache(key=itpm_key) == 4 + + @pytest.mark.asyncio + async def test_reconcile_missing_usage_keeps_reservation(self): + dual_cache = DualCache() + itpm_key = "global_router:io-missing-usage:bedrock_mantle/test:itpm:00-00" + otpm_key = "global_router:io-missing-usage:bedrock_mantle/test:otpm:00-00" + await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) + await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) + + kwargs = { + "metadata": { + ITPM_RESERVED_KEY: 8, + OTPM_RESERVED_KEY: 5, + ITPM_CACHE_KEY: itpm_key, + OTPM_CACHE_KEY: otpm_key, + } + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], + ) + + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + assert await dual_cache.async_get_cache(key=itpm_key) == 8 + assert await dual_cache.async_get_cache(key=otpm_key) == 5 + assert ITPM_RESERVED_KEY not in kwargs["metadata"] + + @pytest.mark.asyncio + async def test_reconcile_total_tokens_only_keeps_reservation(self): + """ + A response usage object with only total_tokens (no prompt/completion + breakdown) can't be split into input/output, so it must be treated the + same as missing usage: keep the reservation instead of resolving to + (0, 0) and refunding it in full. + """ + dual_cache = DualCache() + itpm_key = "global_router:io-total-only:bedrock_mantle/test:itpm:00-00" + otpm_key = "global_router:io-total-only:bedrock_mantle/test:otpm:00-00" + await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) + await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) + + kwargs = { + "metadata": { + ITPM_RESERVED_KEY: 8, + OTPM_RESERVED_KEY: 5, + ITPM_CACHE_KEY: itpm_key, + OTPM_CACHE_KEY: otpm_key, + } + } + response = {"type": "message", "usage": {"total_tokens": 13}} + + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + assert await dual_cache.async_get_cache(key=itpm_key) == 8 + assert await dual_cache.async_get_cache(key=otpm_key) == 5 + + def test_reconcile_standard_logging_total_tokens_only_keeps_reservation(self): + dual_cache = DualCache() + itpm_key = "global_router:io-slo-total-only:bedrock_mantle/test:itpm:00-00" + dual_cache.set_cache(key=itpm_key, value=10, ttl=60) + + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + "standard_logging_object": {"total_tokens": 4}, + } + response = {"type": "message", "role": "assistant", "content": []} + + io_token_reconcile_success(dual_cache, kwargs, response) + + assert dual_cache.get_cache(key=itpm_key) == 10 + + @pytest.mark.asyncio + async def test_reconcile_falls_back_to_standard_logging_object(self): + dual_cache = DualCache() + itpm_key = "global_router:io-slo-fallback:bedrock_mantle/test:itpm:00-00" + await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + "standard_logging_object": { + "prompt_tokens": 4, + "completion_tokens": 0, + "total_tokens": 4, + }, + } + response = {"type": "message", "role": "assistant", "content": []} + + await async_io_token_reconcile_success(dual_cache, kwargs, response) + + assert await dual_cache.async_get_cache(key=itpm_key) == 4 + + def test_sync_reconcile_anthropic_dict_usage(self): + dual_cache = DualCache() + itpm_key = "global_router:io-anthropic:bedrock_mantle/test:itpm:00-00" + otpm_key = "global_router:io-anthropic:bedrock_mantle/test:otpm:00-00" + dual_cache.set_cache(key=itpm_key, value=6, ttl=60) + dual_cache.set_cache(key=otpm_key, value=4, ttl=60) + + kwargs = { + "metadata": { + ITPM_RESERVED_KEY: 6, + OTPM_RESERVED_KEY: 4, + ITPM_CACHE_KEY: itpm_key, + OTPM_CACHE_KEY: otpm_key, + } + } + response = { + "type": "message", + "usage": {"input_tokens": 3, "output_tokens": 2, "cache_read_input_tokens": 1}, + } + + io_token_reconcile_success(dual_cache, kwargs, response) + + assert dual_cache.get_cache(key=itpm_key) == 2 + assert dual_cache.get_cache(key=otpm_key) == 2 + + @pytest.mark.asyncio + async def test_io_and_tpm_rpm_limits_both_enforced_with_warning(self, caplog): + import logging + + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + model_id = "io-mixed-id" + deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" + deployment = { + "litellm_params": { + "model": deployment_name, + "itpm": 100, + "rpm": 1, + }, + "model_info": {"id": model_id}, + "model_name": "opus", + } + + minute = get_utc_datetime().strftime("%H-%M") + rpm_key = f"{model_id}:{deployment_name}:rpm:{minute}" + itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" + await dual_cache.async_increment_cache(key=rpm_key, value=5, ttl=60) + + request_kwargs = { + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 5, + "metadata": {}, + } + set_io_token_rate_limit_request_kwargs(request_kwargs) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + with pytest.raises(litellm.RateLimitError): + await check.async_pre_call_check(deployment) + + assert await dual_cache.async_get_cache(key=rpm_key) == 6 + assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 + assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] + assert any("both limit types are enforced" in record.message for record in caplog.records) + + @pytest.mark.asyncio + async def test_io_success_still_tracks_tpm_for_mixed_deployment(self): + """ + A deployment with itpm/otpm AND tpm/rpm must have BOTH counters updated on + success, otherwise the tpm_key the pre-call check reads is never written + and the tpm_limit can never be enforced. + """ + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + model_id = "io-tpm-mixed-id" + deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" + tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key, value=5, ttl=60) + + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, + "standard_logging_object": { + "model_id": model_id, + "total_tokens": 7, + "hidden_params": {"litellm_model_name": deployment_name}, + }, + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + ) + + await check.async_log_success_event(kwargs, response, None, None) + + # ITPM reconciled down from the 5-token reservation to actual usage (3). + assert await dual_cache.async_get_cache(key=itpm_key) == 3 + # TPM tracking must still run so the tpm/rpm pre-call path can enforce it. + assert await dual_cache.async_get_cache(key=tpm_key) == 7 + + def test_io_success_still_tracks_tpm_for_mixed_deployment_sync(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + model_id = "io-tpm-mixed-sync-id" + deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" + tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" + dual_cache.set_cache(key=itpm_key, value=5, ttl=60) + + kwargs = { + "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, + "standard_logging_object": { + "model_id": model_id, + "total_tokens": 7, + "hidden_params": {"litellm_model_name": deployment_name}, + }, + } + response = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + ) + + check.log_success_event(kwargs, response, None, None) + + assert dual_cache.get_cache(key=itpm_key) == 3 + assert dual_cache.get_cache(key=tpm_key) == 7 + + @pytest.mark.asyncio + async def test_failure_refunds_itpm_reservation(self): + from litellm.utils import get_utc_datetime + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + minute = get_utc_datetime().strftime("%H-%M") + itpm_key = f"global_router:io-refund-id:bedrock_mantle/test:itpm:{minute}" + await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) + + reservation = {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key} + kwargs = { + "standard_logging_object": { + "model_id": "io-refund-id", + "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, + "metadata": dict(reservation), + }, + "metadata": dict(reservation), + } + await check.async_log_failure_event(kwargs, None, None, None) + + current = await dual_cache.async_get_cache(key=itpm_key) + assert current == 0 + + +class TestRouterIOTokenIntegration: + @pytest.mark.asyncio + async def test_model_group_info_aggregates_io_limits(self): + router = Router( + model_list=[ + { + "model_name": "opus", + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-opus-4-7", + "itpm": 100, + "otpm": 20, + }, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + info = router.get_model_group_info("opus") + assert info is not None + assert info.itpm == 100 + assert info.otpm == 20 + + +class TestContextSlotRetention: + def test_setter_stores_kwargs_only_for_io_limited_deployments(self): + """ + The context slot pins the entire request kwargs (messages included) + for the lifetime of the surrounding asyncio context, and pooled + resources created mid-request (e.g. redis connections) capture that + context, extending the pin far past the request. Only ITPM/OTPM + pre-call checks read the slot, so the setter must store None for + deployments without io token limits and still clear reservation + sentinels from kwargs either way. + """ + kwargs = { + "messages": [{"role": "user", "content": "x" * 1000}], + "metadata": {ITPM_RESERVED_KEY: 999, ITPM_CACHE_KEY: "forged"}, + } + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) + assert get_io_token_rate_limit_request_kwargs() is None + assert ITPM_RESERVED_KEY not in kwargs["metadata"] + assert ITPM_CACHE_KEY not in kwargs["metadata"] + + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=True) + assert get_io_token_rate_limit_request_kwargs() is kwargs + + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) + assert get_io_token_rate_limit_request_kwargs() is None + + @pytest.mark.asyncio + async def test_router_does_not_pin_kwargs_without_io_limits(self): + router = Router( + model_list=[ + { + "model_name": "plain", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + } + ] + ) + set_io_token_rate_limit_request_kwargs(None) + kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + deployment = router.get_deployment_by_model_group_name("plain") + assert deployment is not None + router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) + assert get_io_token_rate_limit_request_kwargs() is None + + @pytest.mark.asyncio + async def test_router_pins_kwargs_for_io_limited_deployment(self): + router = Router( + model_list=[ + { + "model_name": "limited", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "itpm": 100}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + set_io_token_rate_limit_request_kwargs(None) + kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + deployment = router.get_deployment_by_model_group_name("limited") + assert deployment is not None + router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) + assert get_io_token_rate_limit_request_kwargs() is kwargs diff --git a/tests/unit/types/llms/test_types_llms_bedrock.py b/tests/unit/types/llms/test_types_llms_bedrock.py new file mode 100644 index 00000000000..a5ad882e775 --- /dev/null +++ b/tests/unit/types/llms/test_types_llms_bedrock.py @@ -0,0 +1,46 @@ +import pytest +from pydantic import ValidationError + +from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams + + +def test_model_validate_keeps_auth_params_and_ignores_request_params(): + auth_params = AwsAuthParams.model_validate( + { + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-role", + "aws_session_name": "litellm-session", + "aws_external_id": "litellm-external-id", + "aws_region_name": "us-west-2", + "aws_bedrock_runtime_endpoint": "https://bedrock.example.com", + "model": "anthropic.claude-haiku-4-5-20251001-v1:0", + "temperature": 0.1, + "messages": [{"role": "user", "content": "hi"}], + } + ) + + assert auth_params.aws_role_name == "arn:aws:iam::999999999999:role/litellm-role" + assert auth_params.aws_session_name == "litellm-session" + assert auth_params.aws_external_id == "litellm-external-id" + assert auth_params.aws_access_key_id is None + assert set(auth_params.model_dump()) == set(AWS_AUTH_PARAM_KEYS) + assert not set(AWS_AUTH_PARAM_KEYS) & {"aws_region_name", "aws_bedrock_runtime_endpoint", "model", "temperature"} + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("aws_role_name", 1234), + ("aws_session_name", ["litellm-session"]), + ("aws_external_id", {"id": "x"}), + ], +) +def test_model_validate_rejects_non_string_credentials(field, value): + with pytest.raises(ValidationError): + AwsAuthParams.model_validate({field: value}) + + +def test_frozen_struct_rejects_field_assignment(): + auth_params = AwsAuthParams(aws_role_name="arn:aws:iam::999999999999:role/litellm-role") + + with pytest.raises(ValidationError): + auth_params.aws_role_name = "arn:aws:iam::999999999999:role/other-role" diff --git a/tests/unit/types/llms/test_types_llms_openai.py b/tests/unit/types/llms/test_types_llms_openai.py new file mode 100644 index 00000000000..64ec09838e8 --- /dev/null +++ b/tests/unit/types/llms/test_types_llms_openai.py @@ -0,0 +1,591 @@ +import asyncio +from typing import Optional +from unittest.mock import AsyncMock, patch + +import pytest + +import json + +import litellm +from litellm.types.llms.openai import HttpxBinaryResponseContent + + +@pytest.mark.parametrize("stream", (False, True)) +def test_completion_response_reasoning_summary_round_trip(stream: bool) -> None: + from typing import Final + + from litellm.types.llms.openai import ( + ChatCompletionReasoningItem, + ChatCompletionReasoningSummaryTextBlock, + ) + from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, + ) + + reasoning_item: Final = ChatCompletionReasoningItem( + type="reasoning", + id="rs_123", + encrypted_content="encrypted", + summary=[ChatCompletionReasoningSummaryTextBlock(type="summary_text", text="Reasoning summary")], + ) + response: Final = ( + ModelResponseStream(choices=[StreamingChoices(delta=Delta(reasoning_items=[reasoning_item]))]) + if stream + else ModelResponse(choices=[Choices(message=Message(reasoning_items=[reasoning_item]))]) + ) + message_key: Final = "delta" if stream else "message" + assert response.model_dump()["choices"][0][message_key]["reasoning_items"] == [reasoning_item] + + restored: Final = type(response).model_validate_json(response.model_dump_json()) + assert restored.model_dump()["choices"][0][message_key]["reasoning_items"] == [reasoning_item] + + +def test_generic_event(): + from litellm.types.llms.openai import GenericEvent + + event = {"type": "test", "test": "test"} + event = GenericEvent(**event) + assert event.type == "test" + assert event.test == "test" + + +def test_output_item_added_event(): + from litellm.types.llms.openai import OutputItemAddedEvent + + event = { + "type": "response.output_item.added", + "sequence_number": 4, + "output_index": 1, + "item": None, + } + event = OutputItemAddedEvent(**event) + assert event.type == "response.output_item.added" + assert event.sequence_number == 4 + assert event.output_index == 1 + assert event.item is None + + +class TestResponsesAPIResponseOutputText: + """Tests for the output_text property on ResponsesAPIResponse""" + + def test_output_text_with_single_message(self): + """Test output_text with a single message containing text output""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Hello, world!", + } + ], + } + ], + ) + + assert response.output_text == "Hello, world!" + + def test_output_text_with_multiple_messages(self): + """Test output_text with multiple messages aggregates all text""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "First part. ", + } + ], + }, + { + "type": "message", + "id": "msg_2", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Second part.", + } + ], + }, + ], + ) + + assert response.output_text == "First part. Second part." + + def test_output_text_with_no_text_content(self): + """Test output_text returns empty string when no output_text content exists""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "function_call", + "id": "call_123", + "status": "completed", + "name": "get_weather", + "arguments": "{}", + } + ], + ) + + assert response.output_text == "" + + def test_output_text_with_mixed_content(self): + """Test output_text only aggregates output_text type content""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "The weather is sunny. ", + }, + { + "type": "refusal", + "refusal": "I cannot do that.", + }, + ], + }, + { + "type": "function_call", + "id": "call_123", + "status": "completed", + "name": "get_weather", + "arguments": "{}", + }, + ], + ) + + assert response.output_text == "The weather is sunny. " + + def test_output_text_with_empty_output(self): + """Test output_text returns empty string with empty output list""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[], + ) + + assert response.output_text == "" + + +class TestAssistantMessageImageUrlContent: + """ + Regression tests for image_url blocks in assistant message content. + + Bug: ChatCompletionAssistantMessage.content did not include + ChatCompletionImageObject in its union, so Pydantic v2 silently dropped + image_url blocks (content → []) when serialising via AllMessageValues. + This affects users who store conversation history as JSON (e.g. in a DB) + and read it back typed as list[AllMessageValues]. + """ + + ASSISTANT_MESSAGE_WITH_IMAGE = { + "role": "assistant", + "content": [ + {"type": "text", "text": "Here is the image you requested:"}, + { + "type": "image_url", + "image_url": { + "url": ( + "data:image/png;base64," + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAA" + "DUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + ) + }, + }, + ], + } + + def test_assistant_message_image_url_preserved_single(self): + """ + TypeAdapter(ChatCompletionAssistantMessage): image_url block must survive + validate_python → dump_python without being dropped or raising an error. + """ + from typing import List + + from pydantic import TypeAdapter + + from litellm.types.llms.openai import ChatCompletionAssistantMessage + + adapter = TypeAdapter(ChatCompletionAssistantMessage) + validated = adapter.validate_python(self.ASSISTANT_MESSAGE_WITH_IMAGE) + dumped = adapter.dump_python(validated) + + raw_content = dumped.get("content") + # Pydantic may return a lazy SerializationIterator for Iterable fields; + # convert to list to consume it — this must not raise ValidationError. + content_blocks = list(raw_content) if raw_content is not None else [] + + assert ( + len(content_blocks) == 2 + ), f"Expected 2 content blocks (text + image_url), got {len(content_blocks)}: {content_blocks}" + types = [b.get("type") for b in content_blocks if isinstance(b, dict)] + assert ( + "image_url" in types + ), f"image_url block was silently dropped; blocks: {content_blocks}" + + def test_assistant_message_image_url_preserved_in_all_message_values(self): + """ + TypeAdapter(List[AllMessageValues]) DB round-trip: image_url blocks in an + assistant message must not be silently dropped during dump_python(mode='json'). + + This is the primary failing path: conversation history stored as JSON in a + database and read back typed as list[AllMessageValues]. + """ + from typing import List + + from pydantic import TypeAdapter + + from litellm.types.llms.openai import AllMessageValues + + conversation = [ + { + "role": "user", + "content": "Generate an image of a banana wearing a LiteLLM costume", + }, + self.ASSISTANT_MESSAGE_WITH_IMAGE, + ] + + adapter = TypeAdapter(List[AllMessageValues]) + validated = adapter.validate_python(conversation) + dumped = adapter.dump_python(validated, mode="json") + + assistant = next((m for m in dumped if m.get("role") == "assistant"), None) + assert assistant is not None, "Assistant message missing after serialisation" + + content = assistant.get("content", []) + assert isinstance( + content, list + ), f"content should be a list, got {type(content)}" + assert ( + len(content) == 2 + ), f"Expected 2 content blocks (text + image_url), got {len(content)}: {content}" + types = [b.get("type") for b in content if isinstance(b, dict)] + assert ( + "image_url" in types + ), f"image_url block was silently dropped during AllMessageValues serialisation; blocks: {content}" + + +class TestResponsesAPIReasoningNullFields: + """ + Tests for issue #16824: reasoning output items should not include null + status/content/encrypted_content fields. + + When a provider returns reasoning items without these fields, LiteLLM's + Pydantic parsing adds them as Optional defaults (None). Serializing them + as null breaks downstream SDKs (e.g., the OpenAI C# SDK crashes on + status=null). + + The fix uses a field_serializer on ResponsesAPIResponse.output that + mirrors the request-side filtering in + OpenAIResponsesAPIConfig._handle_reasoning_item(). + """ + + def _make_response(self, output): + from litellm.types.llms.openai import ResponsesAPIResponse + + return ResponsesAPIResponse( + id="resp_test", + created_at=1741476542, + model="gpt-5-mini", + object="response", + status="completed", + output=output, + ) + + def test_reasoning_item_null_fields_removed_model_dump(self): + """Null status/content/encrypted_content should be absent from model_dump.""" + response = self._make_response( + output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] + ) + dumped = response.model_dump() + reasoning = dumped["output"][0] + assert "status" not in reasoning + assert "content" not in reasoning + assert "encrypted_content" not in reasoning + + def test_reasoning_item_null_fields_removed_model_dump_json(self): + """Null fields should also be absent from model_dump_json.""" + response = self._make_response( + output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] + ) + parsed = json.loads(response.model_dump_json()) + reasoning = parsed["output"][0] + assert "status" not in reasoning + assert "content" not in reasoning + assert "encrypted_content" not in reasoning + + def test_reasoning_item_non_null_values_preserved(self): + """Non-null values on reasoning items should be kept.""" + response = self._make_response( + output=[ + { + "id": "rs_abc", + "type": "reasoning", + "summary": [], + "status": "completed", + "encrypted_content": "gAAAA...", + } + ] + ) + dumped = response.model_dump() + reasoning = dumped["output"][0] + assert reasoning["status"] == "completed" + assert reasoning["encrypted_content"] == "gAAAA..." + + def test_message_item_not_affected(self): + """Non-reasoning output items should keep all their fields.""" + response = self._make_response( + output=[ + { + "id": "msg_abc", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "Hello!", + "annotations": [], + } + ], + } + ] + ) + dumped = response.model_dump() + message = dumped["output"][0] + assert message["status"] == "completed" + assert message["type"] == "message" + assert len(message["content"]) == 1 + + def test_mixed_output_reasoning_and_message(self): + """Reasoning items cleaned, message items untouched in same response.""" + response = self._make_response( + output=[ + {"id": "rs_abc", "type": "reasoning", "summary": []}, + { + "id": "msg_abc", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "Answer", + "annotations": [], + } + ], + }, + ] + ) + dumped = response.model_dump() + reasoning = [ + o + for o in dumped["output"] + if isinstance(o, dict) and o.get("type") == "reasoning" + ][0] + message = [ + o + for o in dumped["output"] + if isinstance(o, dict) and o.get("type") == "message" + ][0] + assert "status" not in reasoning + assert "content" not in reasoning + assert message["status"] == "completed" + assert len(message["content"]) == 1 + + def test_reasoning_core_fields_preserved(self): + """id, type, summary should always be present on reasoning items.""" + response = self._make_response( + output=[{"id": "rs_abc", "type": "reasoning", "summary": ["thinking..."]}] + ) + dumped = response.model_dump() + reasoning = dumped["output"][0] + assert reasoning["id"] == "rs_abc" + assert reasoning["type"] == "reasoning" + assert reasoning["summary"] == ["thinking..."] + + def test_top_level_null_fields_unaffected(self): + """Top-level response fields with None should not be affected.""" + response = self._make_response( + output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] + ) + dumped = response.model_dump() + assert "error" in dumped + assert dumped["error"] is None + assert "instructions" in dumped + assert dumped["instructions"] is None + + +def test_normalize_fine_tuning_job_dict_maps_azure_pending(): + from litellm.llms.openai.fine_tuning.handler import _normalize_fine_tuning_job_dict + + out = _normalize_fine_tuning_job_dict( + {"organization_id": None, "result_files": None, "status": "pending"}, + is_azure=True, + ) + assert out["organization_id"] == "" + assert out["result_files"] == [] + assert out["status"] == "queued" + + +def test_normalize_fine_tuning_job_dict_openai_unchanged(): + from litellm.llms.openai.fine_tuning.handler import _normalize_fine_tuning_job_dict + + data = {"organization_id": None, "result_files": None, "status": "pending"} + out = _normalize_fine_tuning_job_dict(data, is_azure=False) + assert out is data + + +def test_openai_file_object_accepts_pending_status(): + from litellm.types.llms.openai import OpenAIFileObject + + file_obj = OpenAIFileObject( + id="file-123", + bytes=1024, + created_at=1677610602, + filename="train.jsonl", + object="file", + purpose="fine-tune", + status="pending", + ) + assert file_obj.status == "pending" + + +class TestOpenAIFileObjectBatchGuardrailSerialization: + """The proxy-only `litellm_batch_guardrail` key must reach the wire only when something set it.""" + + @staticmethod + def _file_object(**overrides): + from litellm.types.llms.openai import OpenAIFileObject + + return OpenAIFileObject( + id="file-123", + object="file", + bytes=1024, + created_at=1677610602, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + **overrides, + ) + + @staticmethod + def _report(): + from litellm.types.llms.openai import BatchGuardrailRecord, BatchGuardrailReport + + return BatchGuardrailReport( + submitted_records=3, + modified_records=(BatchGuardrailRecord(line=2, custom_id="dirty", action="redacted"),), + ) + + @pytest.mark.parametrize("mode", ["python", "json"]) + def test_key_absent_when_unset(self, mode): + assert "litellm_batch_guardrail" not in self._file_object().model_dump(mode=mode) + + @pytest.mark.parametrize("mode", ["python", "json"]) + def test_key_present_when_set(self, mode): + dumped = self._file_object(litellm_batch_guardrail=self._report()).model_dump(mode=mode) + assert dumped["litellm_batch_guardrail"]["submitted_records"] == 3 + + def test_nested_nulls_of_a_set_report_survive(self): + """`exclude_none=True` was rejected as the fix because it would strip these.""" + dumped = self._file_object(litellm_batch_guardrail=self._report()).model_dump(mode="json") + assert dumped["litellm_batch_guardrail"]["modified_records"] == [ + {"line": 2, "custom_id": "dirty", "action": "redacted", "guardrail": None} + ] + + def test_by_alias_dump_also_omits_the_key(self): + """Tripwire: the serializer filters a literal key name, which an added alias would bypass.""" + assert "litellm_batch_guardrail" not in self._file_object().model_dump(mode="json", by_alias=True) + + def test_other_optional_fields_still_serialize_as_null(self): + dumped = self._file_object().model_dump(mode="json") + assert dumped["expires_at"] is None + assert dumped["status_details"] is None + + def test_round_trip_of_a_set_report_is_lossless(self): + from litellm.types.llms.openai import OpenAIFileObject + + original = self._file_object(litellm_batch_guardrail=self._report()) + assert OpenAIFileObject(**original.model_dump()) == original + + def test_serialization_json_schema_still_describes_the_model(self): + """A return annotation on the wrap serializer would collapse this to a bare object.""" + from litellm.types.llms.openai import OpenAIFileObject + + schema = OpenAIFileObject.model_json_schema(mode="serialization") + assert "litellm_batch_guardrail" in schema["properties"] + + def test_key_omitted_inside_a_file_list_page(self): + from litellm.types.llms.openai import FileListPage + + page = FileListPage(object="list", data=[self._file_object()], has_more=False) + assert "litellm_batch_guardrail" not in page.model_dump(mode="json")["data"][0] + + +def _binary_content(payload: bytes) -> HttpxBinaryResponseContent: + import httpx + + return HttpxBinaryResponseContent(httpx.Response(200, content=payload)) + + +def test_httpx_binary_response_content_hidden_params_are_per_instance(): + first = _binary_content(b"first") + second = _binary_content(b"second") + + first._hidden_params["response_cost"] = 0.5 + + assert second._hidden_params == {} + + +def test_set_response_cost_none_leaves_hidden_params_empty(): + binary_response = _binary_content(b"audio") + + binary_response.set_response_cost(None) + + assert "response_cost" not in binary_response._hidden_params + + binary_response.set_response_cost(0.25) + + assert binary_response._hidden_params["response_cost"] == 0.25 + + binary_response.set_response_cost(None) + + assert "response_cost" not in binary_response._hidden_params diff --git a/tests/unit/types/proxy/policy_engine/__init__.py b/tests/unit/types/proxy/policy_engine/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/types/proxy/policy_engine/test_pipeline_types.py b/tests/unit/types/proxy/policy_engine/test_pipeline_types.py new file mode 100644 index 00000000000..2e5986d3ef8 --- /dev/null +++ b/tests/unit/types/proxy/policy_engine/test_pipeline_types.py @@ -0,0 +1,168 @@ +""" +Tests for pipeline type definitions. +""" + +import pytest +from pydantic import ValidationError + +from litellm.types.proxy.policy_engine.pipeline_types import ( + GuardrailPipeline, + PipelineExecutionResult, + PipelineStep, + PipelineStepResult, +) +from litellm.types.proxy.policy_engine.policy_types import ( + Policy, + PolicyGuardrails, +) + + +def test_pipeline_step_defaults(): + step = PipelineStep(guardrail="my-guard") + assert step.on_fail == "block" + assert step.on_pass == "allow" + assert step.on_error is None + assert step.pass_data is False + assert step.modify_response_message is None + + +def test_pipeline_step_valid_actions(): + step = PipelineStep(guardrail="my-guard", on_fail="next", on_pass="next") + assert step.on_fail == "next" + assert step.on_pass == "next" + + +def test_pipeline_step_all_action_types(): + for action in ("allow", "block", "next", "modify_response"): + step = PipelineStep( + guardrail="g", on_fail=action, on_pass=action, on_error=action + ) + assert step.on_fail == action + assert step.on_pass == action + assert step.on_error == action + + +def test_pipeline_step_invalid_action_rejected(): + with pytest.raises(ValidationError): + PipelineStep(guardrail="my-guard", on_fail="invalid_action") + + +def test_pipeline_step_invalid_on_pass_rejected(): + with pytest.raises(ValidationError): + PipelineStep(guardrail="my-guard", on_pass="skip") + + +def test_pipeline_step_on_error_valid(): + step = PipelineStep( + guardrail="g", on_error="next", on_fail="block", on_pass="allow" + ) + assert step.on_error == "next" + + +def test_pipeline_step_invalid_on_error_rejected(): + with pytest.raises(ValidationError): + PipelineStep(guardrail="my-guard", on_error="invalid") + + +def test_pipeline_requires_at_least_one_step(): + with pytest.raises(ValidationError): + GuardrailPipeline(mode="pre_call", steps=[]) + + +def test_pipeline_invalid_mode_rejected(): + with pytest.raises(ValidationError): + GuardrailPipeline( + mode="during_call", + steps=[PipelineStep(guardrail="g")], + ) + + +def test_pipeline_valid_modes(): + for mode in ("pre_call", "post_call"): + pipeline = GuardrailPipeline( + mode=mode, + steps=[PipelineStep(guardrail="g")], + ) + assert pipeline.mode == mode + + +def test_pipeline_with_multiple_steps(): + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[ + PipelineStep(guardrail="g1", on_fail="next", on_pass="allow"), + PipelineStep(guardrail="g2", on_fail="block", on_pass="allow"), + ], + ) + assert len(pipeline.steps) == 2 + assert pipeline.steps[0].guardrail == "g1" + assert pipeline.steps[1].guardrail == "g2" + + +def test_policy_with_pipeline_parses(): + policy = Policy( + guardrails=PolicyGuardrails(add=["g1", "g2"]), + pipeline=GuardrailPipeline( + mode="pre_call", + steps=[ + PipelineStep(guardrail="g1", on_fail="next"), + PipelineStep(guardrail="g2"), + ], + ), + ) + assert policy.pipeline is not None + assert len(policy.pipeline.steps) == 2 + + +def test_policy_without_pipeline(): + policy = Policy( + guardrails=PolicyGuardrails(add=["g1"]), + ) + assert policy.pipeline is None + + +def test_pipeline_step_result(): + result = PipelineStepResult( + guardrail_name="g1", + outcome="fail", + action_taken="next", + error_detail="Content policy violation", + duration_seconds=0.05, + ) + assert result.outcome == "fail" + assert result.action_taken == "next" + + +def test_pipeline_execution_result(): + result = PipelineExecutionResult( + terminal_action="block", + step_results=[ + PipelineStepResult( + guardrail_name="g1", + outcome="fail", + action_taken="next", + ), + PipelineStepResult( + guardrail_name="g2", + outcome="fail", + action_taken="block", + ), + ], + error_message="Content blocked", + ) + assert result.terminal_action == "block" + assert len(result.step_results) == 2 + + +def test_pipeline_step_extra_fields_rejected(): + with pytest.raises(ValidationError): + PipelineStep(guardrail="g", unknown_field="value") + + +def test_pipeline_extra_fields_rejected(): + with pytest.raises(ValidationError): + GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="g")], + unknown="value", + ) diff --git a/tests/unit/types/proxy/policy_engine/test_policy_types.py b/tests/unit/types/proxy/policy_engine/test_policy_types.py new file mode 100644 index 00000000000..bcd6d39aa4d --- /dev/null +++ b/tests/unit/types/proxy/policy_engine/test_policy_types.py @@ -0,0 +1,15 @@ +import pytest +from pydantic import ValidationError + +from litellm.types.proxy.policy_engine.policy_types import PolicyAttachment + + +@pytest.mark.parametrize("priority", [-2147483648, 2147483647]) +def test_policy_attachment_accepts_int32_priority(priority: int): + assert PolicyAttachment(policy="p", priority=priority).priority == priority + + +@pytest.mark.parametrize("priority", [-2147483649, 2147483648]) +def test_policy_attachment_rejects_priority_outside_int32(priority: int): + with pytest.raises(ValidationError): + PolicyAttachment(policy="p", priority=priority) diff --git a/tests/unit/types/proxy/policy_engine/test_resolver_types.py b/tests/unit/types/proxy/policy_engine/test_resolver_types.py new file mode 100644 index 00000000000..f31b9d7e873 --- /dev/null +++ b/tests/unit/types/proxy/policy_engine/test_resolver_types.py @@ -0,0 +1,115 @@ +""" +Tests for pipeline field on policy CRUD types (resolver_types.py). +""" + +import pytest +from pydantic import ValidationError + +from litellm.types.proxy.policy_engine.resolver_types import ( + PolicyAttachmentCreateRequest, + PolicyCreateRequest, + PolicyDBResponse, + PolicyUpdateRequest, +) + + +def test_policy_create_request_with_pipeline(): + pipeline_data = { + "mode": "pre_call", + "steps": [ + {"guardrail": "g1", "on_fail": "next", "on_pass": "allow"}, + {"guardrail": "g2", "on_fail": "block", "on_pass": "allow"}, + ], + } + req = PolicyCreateRequest( + policy_name="test-policy", + guardrails_add=["g1", "g2"], + pipeline=pipeline_data, + ) + assert req.pipeline is not None + assert req.pipeline["mode"] == "pre_call" + assert len(req.pipeline["steps"]) == 2 + + +def test_policy_create_request_without_pipeline(): + req = PolicyCreateRequest( + policy_name="test-policy", + guardrails_add=["g1"], + ) + assert req.pipeline is None + + +def test_policy_update_request_with_pipeline(): + pipeline_data = { + "mode": "pre_call", + "steps": [ + {"guardrail": "g1", "on_fail": "block", "on_pass": "allow"}, + ], + } + req = PolicyUpdateRequest(pipeline=pipeline_data) + assert req.pipeline is not None + assert req.pipeline["steps"][0]["guardrail"] == "g1" + + +def test_policy_db_response_with_pipeline(): + pipeline_data = { + "mode": "pre_call", + "steps": [ + {"guardrail": "g1", "on_fail": "next", "on_pass": "allow"}, + {"guardrail": "g2", "on_fail": "block", "on_pass": "allow"}, + ], + } + resp = PolicyDBResponse( + policy_id="test-id", + policy_name="test-policy", + guardrails_add=["g1", "g2"], + pipeline=pipeline_data, + ) + assert resp.pipeline is not None + assert resp.pipeline["mode"] == "pre_call" + dumped = resp.model_dump() + assert dumped["pipeline"]["steps"][0]["guardrail"] == "g1" + + +def test_policy_db_response_without_pipeline(): + resp = PolicyDBResponse( + policy_id="test-id", + policy_name="test-policy", + ) + assert resp.pipeline is None + dumped = resp.model_dump() + assert dumped["pipeline"] is None + + +def test_policy_create_request_roundtrip(): + pipeline_data = { + "mode": "post_call", + "steps": [ + { + "guardrail": "g1", + "on_fail": "modify_response", + "on_pass": "next", + "pass_data": True, + "modify_response_message": "custom msg", + }, + ], + } + req = PolicyCreateRequest( + policy_name="roundtrip-test", + guardrails_add=["g1"], + pipeline=pipeline_data, + ) + dumped = req.model_dump() + restored = PolicyCreateRequest(**dumped) + assert restored.pipeline == pipeline_data + + +@pytest.mark.parametrize("priority", [-2147483648, 2147483647]) +def test_policy_attachment_create_request_accepts_int32_priority(priority: int): + assert PolicyAttachmentCreateRequest(policy_name="p", priority=priority).priority == priority + + +@pytest.mark.parametrize("priority", [-2147483649, 2147483648]) +def test_policy_attachment_create_request_rejects_priority_outside_int32(priority: int): + with pytest.raises(ValidationError): + PolicyAttachmentCreateRequest(policy_name="p", priority=priority) diff --git a/tests/unit/videos/__init__.py b/tests/unit/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/videos/test_main.py b/tests/unit/videos/test_main.py new file mode 100644 index 00000000000..22e1e5c05eb --- /dev/null +++ b/tests/unit/videos/test_main.py @@ -0,0 +1,455 @@ +""" +Dispatch-contract tests for litellm/videos/main.py + +Each public video operation is a pair: a sync `video_*` worker (decorated with +@client) that resolves the provider, fetches the provider config, logs, and then +forwards to exactly one `base_llm_http_handler.video_*_handler`; and an async +`avideo_*` wrapper that delegates to the sync worker in an executor. + +This file locks the contract of that layer so a regression fails loudly: + + 1. DISPATCH - the one correct handler fired and every sibling video handler + asserted NOT called. A copy-paste that calls the wrong handler + (e.g. remix -> edit) flips this. + 2. RESULT - the handler's return value is propagated by identity. + 3. PROVIDER - custom_llm_provider is decoded from an encoded video id when not + passed (status/content/remix/edit/extension), or defaults to + "openai" (list/create_character/get_character). This is the exact + surface of the historical "content defaulted to openai" bug. + 4. PAYLOAD - the provider config object and the operation's identifying args + (video_id/prompt/name/...) reach the handler; _is_async is False + on the sync path. + 5. SHORT-CIRCUIT - mock_response returns a typed object without any handler call. + 6. UNSUPPORTED - a None provider config raises before any handler fires. + 7. DELEGATION - avideo_* returns the sync worker's result untouched, sets + async_call=True, and pre-resolves the provider where it must. + +Seams mocked: the http handler (network), the provider-config registry lookup, +get_llm_provider, and the video-generation optional-param builders. The id decode +helper runs for real against genuinely-encoded ids, so the provider assertions +reflect production. +""" + +from contextlib import ExitStack +from dataclasses import dataclass +from typing import Any, Dict +from unittest.mock import MagicMock, patch + +import pytest + + +import litellm +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.types.videos.main import CharacterObject, VideoObject +from litellm.types.videos.utils import encode_video_id_with_provider +from litellm.videos import main as videos_main + +# A real model-encoded video id: decodes (for real) to provider "azure". Used to +# prove the sync workers derive custom_llm_provider from the id, not a hardcode. +AZURE_VIDEO_ID = encode_video_id_with_provider("video_raw", "azure", "deployment-1") + +# The nine sync handlers on base_llm_http_handler. Dispatch tests assert exactly +# one fired and the other eight did not. +SYNC_HANDLERS = ( + "video_generation_handler", + "video_content_handler", + "video_remix_handler", + "video_create_character_handler", + "video_get_character_handler", + "video_edit_handler", + "video_extension_handler", + "video_list_handler", + "video_status_handler", +) + +GEN_OPTIONAL_PARAMS = {"seconds": "8", "size": "720x1280"} + + +@dataclass +class Seams: + handler: MagicMock + get_config: MagicMock + config: MagicMock + + def kwargs_of(self, handler_name: str) -> Dict[str, Any]: + method = getattr(self.handler, handler_name) + assert method.call_count == 1 + return dict(method.call_args.kwargs) + + def assert_only(self, handler_name: str) -> None: + for name in SYNC_HANDLERS: + method = getattr(self.handler, name) + if name == handler_name: + method.assert_called_once() + else: + method.assert_not_called() + + +@pytest.fixture +def seams(): + handler = MagicMock(spec=BaseLLMHTTPHandler) + config = MagicMock(name="provider_video_config") + get_config = MagicMock(return_value=config) + + with ExitStack() as stack: + stack.enter_context(patch.object(videos_main, "base_llm_http_handler", handler)) + stack.enter_context( + patch.object( + videos_main.ProviderConfigManager, + "get_provider_video_config", + get_config, + ) + ) + # video_generation resolves model+provider through get_llm_provider and + # builds optional params; mock those so the dispatch payload is deterministic. + stack.enter_context( + patch.object( + videos_main, + "get_llm_provider", + MagicMock(return_value=("sora-2", "openai", None, None)), + ) + ) + stack.enter_context( + patch.object( + videos_main.VideoGenerationRequestUtils, + "get_requested_video_generation_optional_param", + MagicMock(return_value={"seconds": "8"}), + ) + ) + stack.enter_context( + patch.object( + videos_main.VideoGenerationRequestUtils, + "get_optional_params_video_generation", + MagicMock(return_value=dict(GEN_OPTIONAL_PARAMS)), + ) + ) + yield Seams(handler=handler, get_config=get_config, config=config) + + +# =========================================================================== # +# Dispatch contract - one rich test per sync worker. +# =========================================================================== # + + +def test_video_generation__dispatch(seams): + result = videos_main.video_generation(prompt="a sunset", model="sora-2") + + seams.assert_only("video_generation_handler") + assert result is seams.handler.video_generation_handler.return_value + kw = seams.kwargs_of("video_generation_handler") + assert kw["model"] == "sora-2" + assert kw["prompt"] == "a sunset" + assert kw["custom_llm_provider"] == "openai" + assert kw["video_generation_provider_config"] is seams.config + assert kw["video_generation_optional_request_params"] == GEN_OPTIONAL_PARAMS + assert kw["_is_async"] is False + + +def test_video_status__dispatch_and_provider_from_id(seams): + result = videos_main.video_status(video_id=AZURE_VIDEO_ID) + + seams.assert_only("video_status_handler") + assert result is seams.handler.video_status_handler.return_value + kw = seams.kwargs_of("video_status_handler") + assert kw["video_id"] == AZURE_VIDEO_ID + assert kw["custom_llm_provider"] == "azure" # decoded from the id, not openai + assert kw["video_status_provider_config"] is seams.config + assert kw["_is_async"] is False + # provider config requested for the decoded provider, not a hardcode. + assert seams.get_config.call_args.kwargs["provider"] == litellm.LlmProviders.AZURE + + +def test_video_content__dispatch_and_provider_from_id(seams): + result = videos_main.video_content(video_id=AZURE_VIDEO_ID, variant="thumbnail") + + seams.assert_only("video_content_handler") + assert result is seams.handler.video_content_handler.return_value + kw = seams.kwargs_of("video_content_handler") + assert kw["video_id"] == AZURE_VIDEO_ID + assert kw["custom_llm_provider"] == "azure" + assert kw["variant"] == "thumbnail" + assert kw["video_content_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_content__plain_id_defaults_to_openai(seams): + videos_main.video_content(video_id="video_plain") + + assert seams.kwargs_of("video_content_handler")["custom_llm_provider"] == "openai" + + +def test_video_remix__dispatch_and_provider_from_id(seams): + result = videos_main.video_remix(video_id=AZURE_VIDEO_ID, prompt="new colors") + + seams.assert_only("video_remix_handler") + assert result is seams.handler.video_remix_handler.return_value + kw = seams.kwargs_of("video_remix_handler") + assert kw["video_id"] == AZURE_VIDEO_ID + assert kw["prompt"] == "new colors" + assert kw["custom_llm_provider"] == "azure" + assert kw["video_remix_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_edit__dispatch_and_provider_from_id(seams): + result = videos_main.video_edit(video_id=AZURE_VIDEO_ID, prompt="brighter") + + seams.assert_only("video_edit_handler") + assert result is seams.handler.video_edit_handler.return_value + kw = seams.kwargs_of("video_edit_handler") + assert kw["video_id"] == AZURE_VIDEO_ID + assert kw["prompt"] == "brighter" + assert kw["custom_llm_provider"] == "azure" + assert kw["video_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_extension__dispatch_and_provider_from_id(seams): + result = videos_main.video_extension( + video_id=AZURE_VIDEO_ID, prompt="continue", seconds="5" + ) + + seams.assert_only("video_extension_handler") + assert result is seams.handler.video_extension_handler.return_value + kw = seams.kwargs_of("video_extension_handler") + assert kw["video_id"] == AZURE_VIDEO_ID + assert kw["prompt"] == "continue" + assert kw["seconds"] == "5" + assert kw["custom_llm_provider"] == "azure" + assert kw["video_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_list__dispatch_defaults_to_openai(seams): + result = videos_main.video_list(after="cur", limit=5, order="desc") + + seams.assert_only("video_list_handler") + assert result is seams.handler.video_list_handler.return_value + kw = seams.kwargs_of("video_list_handler") + assert kw["after"] == "cur" + assert kw["limit"] == 5 + assert kw["order"] == "desc" + assert kw["custom_llm_provider"] == "openai" + assert kw["video_list_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_create_character__dispatch_defaults_to_openai(seams): + video = MagicMock(name="video_upload") + result = videos_main.video_create_character(name="hero", video=video) + + seams.assert_only("video_create_character_handler") + assert result is seams.handler.video_create_character_handler.return_value + kw = seams.kwargs_of("video_create_character_handler") + assert kw["name"] == "hero" + assert kw["video"] is video + assert kw["custom_llm_provider"] == "openai" + assert kw["video_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_video_get_character__dispatch_defaults_to_openai(seams): + result = videos_main.video_get_character(character_id="char_1") + + seams.assert_only("video_get_character_handler") + assert result is seams.handler.video_get_character_handler.return_value + kw = seams.kwargs_of("video_get_character_handler") + assert kw["character_id"] == "char_1" + assert kw["custom_llm_provider"] == "openai" + assert kw["video_provider_config"] is seams.config + assert kw["_is_async"] is False + + +def test_explicit_provider_beats_decoded_id(seams): + """An explicit custom_llm_provider wins over the one encoded in the id.""" + videos_main.video_status(video_id=AZURE_VIDEO_ID, custom_llm_provider="vertex_ai") + + assert seams.kwargs_of("video_status_handler")["custom_llm_provider"] == "vertex_ai" + + +# =========================================================================== # +# mock_response short-circuit - returns a typed object, no handler call. +# =========================================================================== # + + +def test_generation__mock_response_short_circuits(seams): + resp = videos_main.video_generation( + prompt="x", + model="sora-2", + mock_response={"id": "v1", "object": "video", "status": "queued"}, + ) + + assert isinstance(resp, VideoObject) + assert resp.id == "v1" + seams.handler.video_generation_handler.assert_not_called() + + +def test_list__mock_response_short_circuits(seams): + resp = videos_main.video_list( + mock_response=[{"id": "v1", "object": "video", "status": "completed"}] + ) + + assert isinstance(resp, list) + assert resp[0].id == "v1" + seams.handler.video_list_handler.assert_not_called() + + +def test_get_character__mock_response_short_circuits(seams): + resp = videos_main.video_get_character( + character_id="char_1", + mock_response={ + "id": "char_1", + "object": "character", + "created_at": 1, + "name": "hero", + }, + ) + + assert isinstance(resp, CharacterObject) + assert resp.id == "char_1" + seams.handler.video_get_character_handler.assert_not_called() + + +# =========================================================================== # +# Unsupported provider - a None provider config raises before any dispatch. +# =========================================================================== # + + +def test_unsupported_provider_raises_without_dispatch(seams): + seams.get_config.return_value = None + + with pytest.raises(litellm.APIConnectionError): + videos_main.video_status(video_id=AZURE_VIDEO_ID) + + seams.handler.video_status_handler.assert_not_called() + + +# =========================================================================== # +# Async-wrapper delegation - representative coverage. +# =========================================================================== # + + +@pytest.mark.asyncio +async def test_avideo_generation__delegates_with_async_flag(): + sentinel = VideoObject(id="v-async", object="video", status="queued") + with ( + patch.object( + videos_main, "video_generation", MagicMock(return_value=sentinel) + ) as sync, + patch.object( + litellm, + "get_llm_provider", + MagicMock(return_value=("sora-2", "openai", None, None)), + ), + ): + result = await videos_main.avideo_generation(prompt="x", model="sora-2") + + assert result is sentinel + assert sync.call_args.kwargs["async_call"] is True + assert sync.call_args.kwargs["custom_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_avideo_status__delegates_untouched(): + sentinel = VideoObject(id="v-async", object="video", status="queued") + with patch.object( + videos_main, "video_status", MagicMock(return_value=sentinel) + ) as sync: + result = await videos_main.avideo_status(video_id="video_plain") + + assert result is sentinel + assert sync.call_args.kwargs["async_call"] is True + assert sync.call_args.kwargs["video_id"] == "video_plain" + + +@pytest.mark.asyncio +async def test_avideo_content__pre_decodes_provider_before_delegating(): + """avideo_content resolves the provider from the encoded id itself before + handing off, so the sync worker receives the decoded provider, not None.""" + sentinel = b"mp4-bytes" + with patch.object( + videos_main, "video_content", MagicMock(return_value=sentinel) + ) as sync: + result = await videos_main.avideo_content(video_id=AZURE_VIDEO_ID) + + assert result is sentinel + assert sync.call_args.kwargs["async_call"] is True + assert sync.call_args.kwargs["custom_llm_provider"] == "azure" + + +# =========================================================================== # +# Credential passthrough - DB/YAML model-config credentials the router injects +# via kwargs must reach the provider call for EVERY video handler, carried in +# litellm_params. Distinct per-field values catch a cross-wired field. +# =========================================================================== # + +DB_YAML_CREDS = { + "api_key": "sk-db-credential", + "api_base": "https://db-resource.test", + "api_version": "2024-12-31", + "vertex_project": "db-project-xyz", +} + +CREDENTIAL_OPERATIONS = [ + ( + "video_generation_handler", + lambda: videos_main.video_generation( + prompt="p", model="sora-2", **DB_YAML_CREDS + ), + ), + ( + "video_status_handler", + lambda: videos_main.video_status(video_id=AZURE_VIDEO_ID, **DB_YAML_CREDS), + ), + ( + "video_content_handler", + lambda: videos_main.video_content(video_id=AZURE_VIDEO_ID, **DB_YAML_CREDS), + ), + ( + "video_remix_handler", + lambda: videos_main.video_remix( + video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS + ), + ), + ( + "video_edit_handler", + lambda: videos_main.video_edit( + video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS + ), + ), + ( + "video_extension_handler", + lambda: videos_main.video_extension( + video_id=AZURE_VIDEO_ID, prompt="p", seconds="5", **DB_YAML_CREDS + ), + ), + ( + "video_list_handler", + lambda: videos_main.video_list(**DB_YAML_CREDS), + ), + ( + "video_create_character_handler", + lambda: videos_main.video_create_character( + name="hero", video=MagicMock(name="vid"), **DB_YAML_CREDS + ), + ), + ( + "video_get_character_handler", + lambda: videos_main.video_get_character(character_id="char_1", **DB_YAML_CREDS), + ), +] + + +@pytest.mark.parametrize( + "handler_name,invoke", + CREDENTIAL_OPERATIONS, + ids=[op[0] for op in CREDENTIAL_OPERATIONS], +) +def test_db_yaml_credentials_reach_every_handler(seams, handler_name, invoke): + invoke() + + litellm_params = seams.kwargs_of(handler_name)["litellm_params"] + assert litellm_params.get("api_key") == DB_YAML_CREDS["api_key"] + assert litellm_params.get("api_base") == DB_YAML_CREDS["api_base"] + assert litellm_params.get("api_version") == DB_YAML_CREDS["api_version"] + assert litellm_params.get("vertex_project") == DB_YAML_CREDS["vertex_project"] diff --git a/tests/unit/videos/test_utils.py b/tests/unit/videos/test_utils.py new file mode 100644 index 00000000000..728644cdda5 --- /dev/null +++ b/tests/unit/videos/test_utils.py @@ -0,0 +1,181 @@ +""" +Pure-logic contract tests for litellm/videos/main.py's request utils +(litellm/videos/utils.py: VideoGenerationRequestUtils). + +These lock the exact param-shaping behavior so a mutation that drops a filter, +flips a precedence, or stops removing a key fails loudly. The only seam is the +provider config's map_openai_params (a provider boundary); filter_out_litellm_params +runs for real, so the "litellm-internal params get stripped" assertions reflect +production. Every test asserts the exact resulting dict, never "ran without error". +""" + +from unittest.mock import MagicMock + + + +import litellm +from litellm.videos.utils import VideoGenerationRequestUtils + +get_requested = ( + VideoGenerationRequestUtils.get_requested_video_generation_optional_param +) +get_optional = VideoGenerationRequestUtils.get_optional_params_video_generation + + +# =========================================================================== # +# get_requested_video_generation_optional_param +# +# Receives the caller's full local_vars; must return only the API-bound optional +# params. filter_out_litellm_params strips known internal keys for real; the +# values used below were chosen against the live set: seconds/size/user/foo_param/ +# vertex_project/extra/a/b survive, api_key/metadata/litellm_* are stripped. +# =========================================================================== # + + +def test_requested__drops_none_and_excluded_keys(): + result = get_requested( + { + "seconds": "8", + "size": None, # None -> dropped + "prompt": "a sunset", # excluded + "model": "sora-2", # excluded + "user": "u1", + } + ) + assert result == {"seconds": "8", "user": "u1"} + + +def test_requested__strips_litellm_internal_params(): + result = get_requested( + { + "seconds": "8", + "api_key": "sk-secret", + "metadata": {"x": 1}, + "litellm_call_id": "id-123", + } + ) + assert result == {"seconds": "8"} + + +def test_requested__timeout_always_removed(): + # timeout is NOT a litellm-internal param, so only the explicit pop removes it. + result = get_requested({"seconds": "8", "timeout": 30}) + assert result == {"seconds": "8"} + + +def test_requested__nested_kwargs_merge_and_override_base(): + result = get_requested( + {"seconds": "8", "kwargs": {"size": "720x1280", "seconds": "override"}} + ) + # nested kwargs win over the top-level base params on collision. + assert result == {"seconds": "override", "size": "720x1280"} + + +def test_requested__non_dict_kwargs_treated_as_empty(): + result = get_requested({"seconds": "8", "kwargs": "not-a-dict"}) + assert result == {"seconds": "8"} + + +def test_requested__none_input_returns_empty(): + assert get_requested(None) == {} + + +def test_requested__top_level_extra_body_spread_and_preserved(): + result = get_requested( + {"seconds": "8", "extra_body": {"vertex_project": "proj", "foo_param": "bar"}} + ) + # extra_body keys are both spread at top level AND kept under "extra_body". + assert result == { + "seconds": "8", + "vertex_project": "proj", + "foo_param": "bar", + "extra_body": {"vertex_project": "proj", "foo_param": "bar"}, + } + + +def test_requested__extra_body_kwargs_overrides_top_level(): + result = get_requested( + { + "extra_body": {"a": "top", "b": "top_b"}, + "kwargs": {"extra_body": {"a": "kw"}}, + } + ) + # kwargs' extra_body wins over the top-level extra_body on collision; the + # non-colliding top-level key survives. + assert result == { + "a": "kw", + "b": "top_b", + "extra_body": {"a": "kw", "b": "top_b"}, + } + + +def test_requested__extra_body_strips_litellm_internal_params(): + result = get_requested({"extra_body": {"api_key": "sk", "foo_param": "bar"}}) + # api_key filtered out of extra_body; only foo_param remains (and is spread). + assert result == {"foo_param": "bar", "extra_body": {"foo_param": "bar"}} + + +def test_requested__empty_extra_body_not_added(): + result = get_requested({"seconds": "8", "extra_body": {}}) + assert result == {"seconds": "8"} + assert "extra_body" not in result + + +# =========================================================================== # +# get_optional_params_video_generation +# +# Delegates mapping to the provider config (the seam) then folds extra_body in. +# =========================================================================== # + + +def _config(map_return): + config = MagicMock() + config.map_openai_params.return_value = map_return + return config + + +def test_optional__delegates_to_map_openai_params_with_drop_params(): + config = _config({"seconds": "8"}) + optional_params = {"seconds": "8"} + + result = get_optional( + model="sora-2", + video_generation_provider_config=config, + video_generation_optional_params=optional_params, + ) + + assert result == {"seconds": "8"} + config.map_openai_params.assert_called_once_with( + video_create_optional_params=optional_params, + model="sora-2", + drop_params=litellm.drop_params, + ) + + +def test_optional__extra_body_overrides_mapped_and_is_removed(): + # mapped output carries a leftover extra_body that must be popped; the input + # extra_body overrides a colliding mapped key and is spread in. + config = _config({"seconds": "8", "size": "mapped", "extra_body": {"leftover": 1}}) + + result = get_optional( + model="sora-2", + video_generation_provider_config=config, + video_generation_optional_params={ + "extra_body": {"size": "override", "extra": "x"} + }, + ) + + assert result == {"seconds": "8", "size": "override", "extra": "x"} + assert "extra_body" not in result + + +def test_optional__non_dict_extra_body_ignored(): + config = _config({"seconds": "8"}) + + result = get_optional( + model="sora-2", + video_generation_provider_config=config, + video_generation_optional_params={"seconds": "8", "extra_body": None}, + ) + + assert result == {"seconds": "8"} From 924ad6e57118b7e42f2f483c9a028c41f90f25f3 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 10:59:22 +0000 Subject: [PATCH 27/47] test: remove phase 16 legacy test files from tests/test_litellm Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../rust_bridge/ocr/test_route_host.py | 85 -- .../rust_bridge/responses/test_route_host.py | 57 - .../test_litellm/sandbox/test_e2b_sandbox.py | 318 ----- .../sandbox/test_opensandbox_sandbox.py | 647 ---------- .../sandbox/test_sandbox_tools.py | 181 --- tests/test_litellm/skills/test_skills_main.py | 57 - .../test_enforce_model_rate_limits.py | 468 -------- .../test_router/test_io_token_rate_limits.py | 1069 ----------------- .../types/llms/test_types_llms_bedrock.py | 46 - .../types/llms/test_types_llms_openai.py | 591 --------- .../policy_engine/test_pipeline_types.py | 168 --- .../proxy/policy_engine/test_policy_types.py | 15 - .../policy_engine/test_resolver_types.py | 115 -- tests/test_litellm/videos/test_main.py | 455 ------- tests/test_litellm/videos/test_utils.py | 193 --- 15 files changed, 4465 deletions(-) delete mode 100644 tests/test_litellm/rust_bridge/ocr/test_route_host.py delete mode 100644 tests/test_litellm/rust_bridge/responses/test_route_host.py delete mode 100644 tests/test_litellm/sandbox/test_e2b_sandbox.py delete mode 100644 tests/test_litellm/sandbox/test_opensandbox_sandbox.py delete mode 100644 tests/test_litellm/sandbox/test_sandbox_tools.py delete mode 100644 tests/test_litellm/skills/test_skills_main.py delete mode 100644 tests/test_litellm/test_router/test_enforce_model_rate_limits.py delete mode 100644 tests/test_litellm/test_router/test_io_token_rate_limits.py delete mode 100644 tests/test_litellm/types/llms/test_types_llms_bedrock.py delete mode 100644 tests/test_litellm/types/llms/test_types_llms_openai.py delete mode 100644 tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py delete mode 100644 tests/test_litellm/types/proxy/policy_engine/test_policy_types.py delete mode 100644 tests/test_litellm/types/proxy/policy_engine/test_resolver_types.py delete mode 100644 tests/test_litellm/videos/test_main.py delete mode 100644 tests/test_litellm/videos/test_utils.py diff --git a/tests/test_litellm/rust_bridge/ocr/test_route_host.py b/tests/test_litellm/rust_bridge/ocr/test_route_host.py deleted file mode 100644 index 699492e4424..00000000000 --- a/tests/test_litellm/rust_bridge/ocr/test_route_host.py +++ /dev/null @@ -1,85 +0,0 @@ -from typing import Final - -import pytest - -import litellm -from litellm.rust_bridge.ocr.route_host import UpstreamFailure, map_failure -from litellm.rust_bridge.ocr.route_host import response as build_ocr_response -from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest - -REQUEST: Final = LiteLLMOcrRequest( - model="mistral/mistral-ocr-latest", - document={"type": "document_url", "document_url": "https://example.com/file.pdf"}, - api_key="test-key", - api_base=None, - timeout=None, - custom_llm_provider=None, - extra_headers=None, - kwargs={"req_format": "markdown"}, -) - - -class RustUpstreamError(Exception): - def __init__(self, status: int, body: str, headers: tuple[tuple[str, str], ...]) -> None: - super().__init__(status, body) - self.headers: Final = list(headers) - - -class RustFormatError(Exception): - ocr_request_format_error: Final = True - - -def test_rust_ocr_response_retains_provider_native_response(): - provider_response = {"status": "succeeded", "analyzeResult": {"content": "native"}} - response = build_ocr_response( - { - "pages": [], - "model": "prebuilt-layout", - "document_annotation": None, - "usage_info": {"pages_processed": 0}, - "object": "ocr", - "provider_native_response": provider_response, - } - ) - - assert response.get_provider_native_response() == provider_response - assert response.model_dump().get("provider_native_response") is None - - -def test_map_failure_builds_public_error_from_upstream_status_and_headers() -> None: - error: Final = RustUpstreamError(429, '{"message": "slow down"}', (("retry-after", "7"),)) - - public_error: Final = map_failure(error, REQUEST, "mistral") - - assert isinstance(public_error, litellm.RateLimitError) - assert public_error.status_code == 429 - assert public_error.response.headers["retry-after"] == "7" - assert public_error.response.text == '{"message": "slow down"}' - assert public_error.__context__ is error - assert public_error.llm_provider == "mistral" - - -def test_map_failure_maps_upstream_401_to_authentication_error() -> None: - error: Final = RustUpstreamError(401, '{"message": "Unauthorized"}', ()) - - public_error: Final = map_failure(error, REQUEST, "mistral") - - assert isinstance(public_error, litellm.AuthenticationError) - assert public_error.status_code == 401 - assert public_error.response.text == '{"message": "Unauthorized"}' - assert public_error.__context__ is error - - -def test_map_failure_leaves_non_upstream_errors_unwrapped() -> None: - error: Final = RuntimeError("bridge exploded") - - public_error: Final = map_failure(error, REQUEST, "mistral") - - assert not isinstance(public_error, UpstreamFailure) - assert isinstance(public_error, litellm.APIConnectionError) - assert "bridge exploded" in str(public_error) - - -def test_map_failure_reports_invalid_request_format_as_unsupported_params() -> None: - with pytest.raises(litellm.UnsupportedParamsError, match="Invalid `req_format`: 'markdown'"): - raise map_failure(RustFormatError(), REQUEST, "mistral") diff --git a/tests/test_litellm/rust_bridge/responses/test_route_host.py b/tests/test_litellm/rust_bridge/responses/test_route_host.py deleted file mode 100644 index 49bf19e7d8a..00000000000 --- a/tests/test_litellm/rust_bridge/responses/test_route_host.py +++ /dev/null @@ -1,57 +0,0 @@ -from types import MappingProxyType -from typing import Final - -import pytest -from pydantic import ValidationError - -from litellm.rust_bridge.responses.route_host import arguments, response -from litellm.rust_bridge.responses.entrypoints import LiteLLMResponsesRequest -from litellm.types.llms.openai import ResponsesAPIResponse - - -def test_response_validates_into_the_public_responses_model() -> None: - built: Final = response( - MappingProxyType( - { - "id": "resp_native", - "object": "response", - "created_at": 1, - "model": "gpt-4o", - "status": "completed", - "output": [ - { - "type": "message", - "id": "msg_native", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": "native", "annotations": []}], - } - ], - } - ) - ) - - assert isinstance(built, ResponsesAPIResponse) - assert built.id == "resp_native" - assert built.output[0].content[0].text == "native" - - -def test_response_rejects_a_payload_missing_required_fields() -> None: - with pytest.raises(ValidationError): - response(MappingProxyType({"object": "response"})) - - -def test_arguments_are_the_public_kwargs_view() -> None: - kwargs: Final = MappingProxyType({"litellm_metadata": {"user_id": "u"}}) - request: Final = LiteLLMResponsesRequest( - model="gpt-4o", - input="hi", - stream=None, - api_key=None, - api_base=None, - custom_llm_provider="openai", - extra_headers=None, - kwargs=kwargs, - ) - - assert arguments(request) is kwargs diff --git a/tests/test_litellm/sandbox/test_e2b_sandbox.py b/tests/test_litellm/sandbox/test_e2b_sandbox.py deleted file mode 100644 index e01b9120416..00000000000 --- a/tests/test_litellm/sandbox/test_e2b_sandbox.py +++ /dev/null @@ -1,318 +0,0 @@ -""" -Tests for the e2b code execution sandbox primitive. - -Unit tests inject a fake async HTTP client (dependency injection, no -monkeypatching) and assert request shapes and result mapping. Real-network -integration tests live in tests/integration/sandbox/test_e2b_sandbox.py. -""" - -import json - -import httpx -import pytest - -import litellm -from litellm.llms.base_llm.sandbox.transformation import ContainerHandle -from litellm.llms.e2b.sandbox.transformation import ( - MAX_OUTPUT_BYTES, - E2BSandboxConfig, -) - - -class FakeResponse: - def __init__(self, *, json_data=None, lines=None, status_code=200): - self._json = json_data - self._lines = lines or [] - self.status_code = status_code - - def json(self): - return self._json - - async def aiter_lines(self): - for line in self._lines: - yield line - - -class FakeHTTPClient: - """Records outbound requests and returns canned responses keyed by URL.""" - - def __init__( - self, - *, - create_json=None, - execute_lines=None, - delete_status=204, - execute_raises=None, - ): - self.create_json = create_json or { - "sandboxID": "sbx_123", - "domain": "e2b.app", - "envdAccessToken": "tok_abc", - } - self.execute_lines = execute_lines or [] - self.delete_status = delete_status - self.execute_raises = execute_raises - self.calls = [] - - async def post(self, url, headers=None, json=None, stream=False, **kwargs): - self.calls.append(("POST", url, headers, json)) - if url.endswith("/sandboxes"): - return FakeResponse(json_data=self.create_json) - if url.endswith("/execute"): - if self.execute_raises is not None: - raise self.execute_raises - return FakeResponse(lines=self.execute_lines) - raise AssertionError(f"unexpected POST {url}") - - async def delete(self, url, headers=None, **kwargs): - self.calls.append(("DELETE", url, headers, None)) - if not (200 <= self.delete_status < 300): - raise httpx.HTTPStatusError( - f"status {self.delete_status}", - request=httpx.Request("DELETE", url), - response=httpx.Response(self.delete_status), - ) - return FakeResponse(status_code=self.delete_status) - - -# ---------- pure parser ---------- - - -def test_parse_lines_stdout_and_count(): - lines = [ - json.dumps({"type": "stdout", "text": "6\n", "timestamp": 1}), - json.dumps({"type": "number_of_executions", "execution_count": 1}), - ] - result = E2BSandboxConfig._parse_lines(lines) - assert result.stdout == "6\n" - assert result.execution_count == 1 - assert result.error is None - - -def test_parse_lines_error_surfaces_name_and_traceback(): - lines = [ - json.dumps( - { - "type": "error", - "name": "ZeroDivisionError", - "value": "division by zero", - "traceback": "Traceback (most recent call last): ...", - } - ) - ] - result = E2BSandboxConfig._parse_lines(lines) - assert result.error["name"] == "ZeroDivisionError" - assert "Traceback" in result.error["traceback"] - - -def test_parse_lines_result_carries_png(): - lines = [ - json.dumps({"type": "result", "png": "BASE64DATA", "is_main_result": True}) - ] - result = E2BSandboxConfig._parse_lines(lines) - assert result.results and result.results[0]["png"] == "BASE64DATA" - assert "type" not in result.results[0] - - -# ---------- request shapes ---------- - - -@pytest.mark.asyncio -async def test_template_flows_into_create_request_as_templateID(): - client = FakeHTTPClient() - cfg = E2BSandboxConfig() - handle = await cfg.acreate_sandbox( - template="my-custom-template", api_key="e2b_key", client=client - ) - - method, url, headers, body = client.calls[0] - assert method == "POST" - assert url.endswith("/sandboxes") - assert body["templateID"] == "my-custom-template" # not "template" - assert body["secure"] is True - assert headers["X-API-Key"] == "e2b_key" - assert handle.id == "sbx_123" - assert handle._hidden_params["envd_access_token"] == "tok_abc" - - -@pytest.mark.asyncio -async def test_create_defaults_template_when_omitted(): - client = FakeHTTPClient() - await E2BSandboxConfig().acreate_sandbox(api_key="e2b_key", client=client) - _, _, _, body = client.calls[0] - assert body["templateID"] == "code-interpreter-v1" - - -@pytest.mark.asyncio -async def test_run_code_targets_jupyter_host_with_access_token(): - client = FakeHTTPClient( - execute_lines=[json.dumps({"type": "stdout", "text": "42\n", "timestamp": 1})] - ) - handle = ContainerHandle(id="sbx_xyz", provider="e2b", domain="e2b.app") - handle._hidden_params = {"envd_access_token": "tok_run"} - - result = await E2BSandboxConfig().arun_code( - container=handle, code="print(6*7)", client=client - ) - - method, url, headers, body = client.calls[0] - assert url == "https://49999-sbx_xyz.e2b.app/execute" - assert headers["X-Access-Token"] == "tok_run" - assert body["code"] == "print(6*7)" - assert result.stdout.strip() == "42" - - -@pytest.mark.asyncio -async def test_delete_issues_delete_to_sandbox_id(): - client = FakeHTTPClient(delete_status=204) - handle = ContainerHandle(id="sbx_del", provider="e2b", domain="e2b.app") - handle._hidden_params = {"api_key": "e2b_key"} - - ok = await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) - - method, url, headers, _ = client.calls[0] - assert method == "DELETE" - assert url.endswith("/sandboxes/sbx_del") - assert ok is True - - -@pytest.mark.asyncio -async def test_delete_returns_false_on_404(): - client = FakeHTTPClient(delete_status=404) - handle = ContainerHandle(id="sbx_gone", provider="e2b", domain="e2b.app") - handle._hidden_params = {"api_key": "e2b_key"} - ok = await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) - assert ok is False - - -# ---------- ephemeral teardown ---------- - - -@pytest.mark.asyncio -async def test_code_interpreter_tool_deletes_even_when_run_raises(): - client = FakeHTTPClient(execute_raises=RuntimeError("boom")) - - with pytest.raises(RuntimeError, match="boom"): - await litellm.acode_interpreter_tool( - provider="e2b", code="1/0", api_key="e2b_key", client=client - ) - - methods = [c[0] for c in client.calls] - urls = [c[1] for c in client.calls] - assert methods == ["POST", "POST", "DELETE"] # create, run(raises), delete - assert urls[0].endswith("/sandboxes") - assert urls[1].endswith("/execute") - assert urls[2].endswith("/sandboxes/sbx_123") - - -# ---------- correctness guards ---------- - - -@pytest.mark.asyncio -async def test_delete_reraises_non_404_http_error(): - client = FakeHTTPClient(delete_status=500) - handle = ContainerHandle(id="sbx_err", provider="e2b", domain="e2b.app") - handle._hidden_params = {"api_key": "e2b_key"} - with pytest.raises(httpx.HTTPStatusError): - await E2BSandboxConfig().adelete_sandbox(container=handle, client=client) - - -@pytest.mark.asyncio -async def test_create_preserves_explicit_zero_timeout(): - client = FakeHTTPClient() - await E2BSandboxConfig().acreate_sandbox( - timeout=0, api_key="e2b_key", client=client - ) - _, _, _, body = client.calls[0] - assert body["timeout"] == 0 - - -@pytest.mark.asyncio -async def test_run_code_rejects_bare_id_without_access_token(): - client = FakeHTTPClient() - with pytest.raises(ValueError, match="access token"): - await E2BSandboxConfig().arun_code( - container="sbx_no_token", code="print(1)", client=client - ) - assert client.calls == [] # never reached the network - - -def test_parse_lines_skips_non_json_lines(): - lines = [ - "not-json-heartbeat", - json.dumps({"type": "stdout", "text": "ok\n"}), - "", - "{partial", - ] - result = E2BSandboxConfig._parse_lines(lines) - assert result.stdout == "ok\n" - assert result.error is None - - -@pytest.mark.asyncio -async def test_run_code_aborts_on_output_over_cap(): - big_line = "x" * (MAX_OUTPUT_BYTES + 1) - client = FakeHTTPClient(execute_lines=[big_line]) - handle = ContainerHandle(id="sbx_big", provider="e2b", domain="e2b.app") - handle._hidden_params = {"envd_access_token": "tok"} - with pytest.raises(ValueError, match="exceeded"): - await E2BSandboxConfig().arun_code( - container=handle, code="print('x'*999)", client=client - ) - - -# ---------- public entrypoints ---------- - - -@pytest.mark.asyncio -async def test_public_lifecycle_create_run_delete(): - client = FakeHTTPClient( - execute_lines=[json.dumps({"type": "stdout", "text": "42\n"})] - ) - container = await litellm.acreate_sandbox( - provider="e2b", api_key="e2b_key", client=client - ) - assert container.id == "sbx_123" - - result = await litellm.arun_code( - provider="e2b", - container=container, - api_key="e2b_key", - code="print(6*7)", - client=client, - ) - assert result.stdout.strip() == "42" - - assert ( - await litellm.adelete_sandbox( - provider="e2b", container=container, api_key="e2b_key", client=client - ) - is True - ) - - -@pytest.mark.asyncio -async def test_unsupported_provider_raises(): - with pytest.raises(ValueError, match="not-a-provider' is not a valid SandboxProviders"): - await litellm.acreate_sandbox(provider="not-a-provider") - - -# ---------- api_base override ---------- - - -@pytest.mark.asyncio -async def test_create_uses_api_base_override(): - client = FakeHTTPClient() - await E2BSandboxConfig().acreate_sandbox( - api_base="http://my-sandbox:8080", api_key="k", client=client - ) - _, url, _, _ = client.calls[0] - assert url == "http://my-sandbox:8080/sandboxes" - - -@pytest.mark.asyncio -async def test_create_defaults_to_e2b_api_base(): - client = FakeHTTPClient() - await E2BSandboxConfig().acreate_sandbox(api_key="k", client=client) - _, url, _, _ = client.calls[0] - assert url == "https://api.e2b.app/sandboxes" diff --git a/tests/test_litellm/sandbox/test_opensandbox_sandbox.py b/tests/test_litellm/sandbox/test_opensandbox_sandbox.py deleted file mode 100644 index 2928dea100e..00000000000 --- a/tests/test_litellm/sandbox/test_opensandbox_sandbox.py +++ /dev/null @@ -1,647 +0,0 @@ -import json - -import httpx -import pytest - -import litellm -from litellm.llms.base_llm.sandbox.transformation import ContainerHandle -from litellm.llms.opensandbox.sandbox.transformation import ( - MAX_OUTPUT_BYTES, - OPEN_SANDBOX_DEFAULT_TEMPLATE, - OpenSandboxSandboxConfig, -) -from litellm.utils import ProviderConfigManager - -TEST_API_BASE = "https://sandbox.test/v1" - - -def http_status_error(status_code, url="http://test"): - return httpx.HTTPStatusError( - f"status {status_code}", - request=httpx.Request("GET", url), - response=httpx.Response(status_code), - ) - - -def sse(data): - return f"data: {json.dumps(data)}" - - -class FakeResponse: - def __init__(self, *, json_data=None, lines=None, status_code=200): - self._json = json_data - self._lines = lines or [] - self.status_code = status_code - - def json(self): - return self._json - - def raise_for_status(self): - if self.status_code >= 400: - raise http_status_error(self.status_code) - - async def aiter_lines(self): - for line in self._lines: - yield line - - -class FakeHTTPClient: - def __init__( - self, - *, - create_json=None, - sandbox_states=None, - endpoint_json=None, - endpoint_responses=None, - execute_lines=None, - delete_status=204, - execute_raises=None, - ): - self.create_json = create_json or { - "id": "osb_123", - "status": {"state": "Running"}, - "createdAt": "2026-01-01T00:00:00Z", - "entrypoint": ["/opt/code-interpreter/code-interpreter.sh"], - } - self.sandbox_states = list( - sandbox_states - or [ - { - "id": "osb_123", - "status": {"state": "Running"}, - "createdAt": "2026-01-01T00:00:00Z", - "entrypoint": ["/opt/code-interpreter/code-interpreter.sh"], - } - ] - ) - self.endpoint_json = endpoint_json or { - "endpoint": "execd.local:44772", - "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, - } - self.endpoint_responses = ( - list(endpoint_responses) if endpoint_responses is not None else None - ) - self.execute_lines = execute_lines or [] - self.delete_status = delete_status - self.execute_raises = execute_raises - self.calls = [] - - async def post(self, url, headers=None, json=None, stream=False, **kwargs): - self.calls.append(("POST", url, headers, json, {"stream": stream})) - if url.endswith("/sandboxes"): - return FakeResponse(json_data=self.create_json) - if url.endswith("/code"): - if self.execute_raises is not None: - raise self.execute_raises - return FakeResponse(lines=self.execute_lines) - raise AssertionError(f"unexpected POST {url}") - - async def get(self, url, headers=None, params=None, **kwargs): - self.calls.append(("GET", url, headers, None, params)) - if "/endpoints/44772" in url: - if self.endpoint_responses is not None and self.endpoint_responses: - response = self.endpoint_responses.pop(0) - if isinstance(response, Exception): - raise response - if isinstance(response, FakeResponse): - return response - return FakeResponse(json_data=response) - return FakeResponse(json_data=self.endpoint_json) - if "/sandboxes/" in url: - state = self.sandbox_states.pop(0) - return FakeResponse(json_data=state) - raise AssertionError(f"unexpected GET {url}") - - async def delete(self, url, headers=None, **kwargs): - self.calls.append(("DELETE", url, headers, None, None)) - if not (200 <= self.delete_status < 300): - raise http_status_error(self.delete_status, url) - return FakeResponse(status_code=self.delete_status) - - -def test_parse_sse_lines_maps_output_result_count_and_error(): - lines = [ - sse({"type": "stdout", "text": "hello\n"}), - sse({"type": "stderr", "text": "warn\n"}), - sse({"type": "result", "results": {"text/plain": "4"}}), - sse({"type": "execution_count", "execution_count": 7}), - sse( - { - "type": "error", - "error": { - "ename": "ValueError", - "evalue": "bad", - "traceback": ["Traceback"], - }, - } - ), - ] - - result = OpenSandboxSandboxConfig._parse_lines(lines) - - assert result.stdout == "hello\n" - assert result.stderr == "warn\n" - assert result.results == [{"text/plain": "4"}] - assert result.execution_count == 7 - assert result.error == { - "name": "ValueError", - "value": "bad", - "traceback": ["Traceback"], - } - - -def test_parse_sse_lines_skips_non_json_and_control_lines(): - lines = [ - "event: message", - "not-json", - "", - sse({"type": "stdout", "text": "ok\n"}), - ] - - result = OpenSandboxSandboxConfig._parse_lines(lines) - - assert result.stdout == "ok\n" - assert result.error is None - - -def test_parse_sse_lines_maps_fallback_shapes(): - lines = [ - "data:", - sse(["not-a-dict"]), - sse({"code": "BadRequest", "message": "nope"}), - sse({"type": "result", "text/plain": "4"}), - sse({"type": "error", "name": "RuntimeError", "text": "boom"}), - sse({"type": "execution_count", "execution_count": "8"}), - ] - - result = OpenSandboxSandboxConfig._parse_lines(lines) - - assert result.results == [{"text/plain": "4"}] - assert result.execution_count == 8 - assert result.error == { - "name": "BadRequest", - "value": "nope", - "traceback": [], - } - fallback_error = OpenSandboxSandboxConfig._parse_lines( - [sse({"type": "error", "name": "RuntimeError", "text": "boom"})] - ) - assert fallback_error.error == { - "name": "RuntimeError", - "value": "boom", - "traceback": [], - } - empty_string_error = OpenSandboxSandboxConfig._parse_lines( - [ - sse( - { - "type": "error", - "error": { - "ename": "", - "name": "FallbackName", - "evalue": "", - "value": "fallback value", - "traceback": [], - }, - } - ) - ] - ) - assert empty_string_error.error == { - "name": "", - "value": "", - "traceback": [], - } - - -def test_static_helpers_cover_defaults_and_fallbacks(monkeypatch): - def fake_secret(key): - if key == "OPEN_SANDBOX_API_KEY": - return "env-key" - if key == "OPEN_SANDBOX_API_BASE": - return TEST_API_BASE - return None - - monkeypatch.setattr( - "litellm.llms.opensandbox.sandbox.transformation.get_secret_str", - fake_secret, - ) - config = OpenSandboxSandboxConfig() - handle = ContainerHandle(id="osb", provider="opensandbox", domain="http://x/v1") - - assert config.validate_environment() == "env-key" - assert config.validate_environment(api_key="") == "" - assert config._api_key(api_key=None, handle=handle) == "env-key" - - handle._hidden_params = {"api_key": "stored-key"} - assert config._api_key(api_key=None, handle=handle) == "stored-key" - assert config._http(None) is not None - - body = config._create_body( - template=None, - timeout=None, - allow_internet_access=False, - metadata=None, - env_vars=None, - resource_limits=None, - resource_requests=None, - entrypoint=None, - network_policy={"egress": [{"domain": "example.com"}]}, - secure_access=True, - ) - assert body["networkPolicy"] == {"egress": [{"domain": "example.com"}]} - assert body["secureAccess"] is True - - other_body = config._create_body( - template=None, - timeout=None, - allow_internet_access=False, - metadata=None, - env_vars=None, - resource_limits=None, - resource_requests=None, - entrypoint=None, - network_policy=None, - secure_access=False, - ) - assert body["resourceLimits"] is not other_body["resourceLimits"] - - assert config._sandbox_state(None) is None - assert config._sandbox_state({"status": "Running"}) is None - assert config._as_str_dict(None) == {} - assert config._endpoint_base_url("http://execd.local", "https://api/v1") == ( - "http://execd.local" - ) - assert config._api_base(None) == TEST_API_BASE - assert config._api_base("https://direct.test/v1/") == "https://direct.test/v1" - assert config._as_int("9") == 9 - assert config._as_int("nope") is None - assert config._as_int(None) is None - assert isinstance( - ProviderConfigManager.get_provider_sandbox_config("opensandbox"), - OpenSandboxSandboxConfig, - ) - - -def test_api_base_requires_kwarg_or_env(monkeypatch): - monkeypatch.setattr( - "litellm.llms.opensandbox.sandbox.transformation.get_secret_str", - lambda key: None, - ) - - with pytest.raises(ValueError, match="api_base is required"): - OpenSandboxSandboxConfig._api_base(None) - - -@pytest.mark.asyncio -async def test_create_posts_default_body_and_omits_empty_api_key(): - client = FakeHTTPClient() - - handle = await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", api_base=TEST_API_BASE, client=client - ) - - method, url, headers, body, _ = client.calls[0] - assert method == "POST" - assert url == f"{TEST_API_BASE}/sandboxes" - assert "OPEN-SANDBOX-API-KEY" not in headers - assert body["image"] == {"uri": OPEN_SANDBOX_DEFAULT_TEMPLATE} - assert body["entrypoint"] == ["/opt/code-interpreter/code-interpreter.sh"] - assert body["timeout"] == 300 - assert body["resourceLimits"] == {"cpu": "1", "memory": "2Gi"} - assert body["networkPolicy"] == {"defaultAction": "deny", "egress": []} - assert handle.id == "osb_123" - assert handle._hidden_params["execd_endpoint"] == "execd.local:44772" - - -@pytest.mark.asyncio -async def test_create_can_opt_into_internet_access(): - client = FakeHTTPClient() - - await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", - api_base=TEST_API_BASE, - allow_internet_access=True, - client=client, - ) - - _, _, _, body, _ = client.calls[0] - assert "networkPolicy" not in body - - -@pytest.mark.asyncio -async def test_create_custom_options_poll_and_endpoint_resolution(): - client = FakeHTTPClient( - create_json={ - "id": "osb_pending", - "status": {"state": "Pending"}, - "createdAt": "2026-01-01T00:00:00Z", - "entrypoint": ["/bin/sh"], - }, - sandbox_states=[ - { - "id": "osb_pending", - "status": {"state": "Running"}, - "createdAt": "2026-01-01T00:00:00Z", - "entrypoint": ["/bin/sh"], - } - ], - ) - - handle = await OpenSandboxSandboxConfig().acreate_sandbox( - template="custom/image:latest", - timeout=600, - allow_internet_access=False, - api_key="osb-key", - api_base="https://sandbox.example/v1", - metadata={"suite": "unit"}, - env_vars={"PYTHONUNBUFFERED": "1"}, - resource_limits={"cpu": "500m", "memory": "512Mi"}, - resource_requests={"cpu": "250m", "memory": "256Mi"}, - entrypoint=["/bin/sh", "-lc", "sleep 3600"], - use_server_proxy=True, - client=client, - ) - - _, create_url, create_headers, body, _ = client.calls[0] - _, poll_url, poll_headers, _, _ = client.calls[1] - _, endpoint_url, endpoint_headers, _, endpoint_params = client.calls[2] - - assert create_url == "https://sandbox.example/v1/sandboxes" - assert create_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" - assert body["image"] == {"uri": "custom/image:latest"} - assert body["entrypoint"] == ["/bin/sh", "-lc", "sleep 3600"] - assert body["metadata"] == {"suite": "unit"} - assert body["env"] == {"PYTHONUNBUFFERED": "1"} - assert body["resourceLimits"] == {"cpu": "500m", "memory": "512Mi"} - assert body["resourceRequests"] == {"cpu": "250m", "memory": "256Mi"} - assert body["networkPolicy"] == {"defaultAction": "deny", "egress": []} - assert poll_url == "https://sandbox.example/v1/sandboxes/osb_pending" - assert poll_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" - assert endpoint_url.endswith("/sandboxes/osb_pending/endpoints/44772") - assert endpoint_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" - assert endpoint_params == {"use_server_proxy": True} - assert handle.id == "osb_pending" - - -@pytest.mark.asyncio -async def test_create_waits_across_pending_state(monkeypatch): - client = FakeHTTPClient( - create_json={ - "id": "osb_pending", - "status": {"state": "Pending"}, - "createdAt": "2026-01-01T00:00:00Z", - }, - sandbox_states=[ - {"id": "osb_pending", "status": {"state": "Pending"}}, - {"id": "osb_pending", "status": {"state": "Running"}}, - ], - ) - sleeps = [] - - async def fake_sleep(interval): - sleeps.append(interval) - - monkeypatch.setattr( - "litellm.llms.opensandbox.sandbox.transformation.asyncio.sleep", fake_sleep - ) - - handle = await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", - api_base=TEST_API_BASE, - ready_timeout=1, - poll_interval=0.01, - client=client, - ) - - assert handle.id == "osb_pending" - assert sleeps == [0.01] - - -@pytest.mark.asyncio -async def test_create_raises_for_terminal_state(): - client = FakeHTTPClient( - create_json={"id": "osb_failed", "status": {"state": "Pending"}}, - sandbox_states=[ - {"id": "osb_failed", "status": {"state": "Failed"}}, - ], - ) - - with pytest.raises(ValueError, match="entered Failed"): - await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", api_base=TEST_API_BASE, client=client - ) - - -@pytest.mark.asyncio -async def test_create_times_out_waiting_for_running(): - client = FakeHTTPClient( - create_json={"id": "osb_slow", "status": {"state": "Pending"}}, - sandbox_states=[ - {"id": "osb_slow", "status": {"state": "Pending"}}, - ], - ) - - with pytest.raises(TimeoutError, match="was not Running"): - await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", - api_base=TEST_API_BASE, - ready_timeout=0, - poll_interval=0, - client=client, - ) - - -@pytest.mark.asyncio -async def test_create_waits_for_endpoint_resolution(monkeypatch): - client = FakeHTTPClient( - endpoint_responses=[ - http_status_error(404, f"{TEST_API_BASE}/sandboxes/osb_123"), - { - "endpoint": "execd.local:44772", - "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, - }, - ], - ) - sleeps = [] - - async def fake_sleep(interval): - sleeps.append(interval) - - monkeypatch.setattr( - "litellm.llms.opensandbox.sandbox.transformation.asyncio.sleep", fake_sleep - ) - - handle = await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", - api_base=TEST_API_BASE, - ready_timeout=1, - poll_interval=0.01, - client=client, - ) - - endpoint_calls = [call for call in client.calls if "/endpoints/44772" in call[1]] - assert handle._hidden_params["execd_endpoint"] == "execd.local:44772" - assert len(endpoint_calls) == 2 - assert sleeps == [0.01] - - -@pytest.mark.asyncio -async def test_create_raises_when_endpoint_is_missing(): - client = FakeHTTPClient(endpoint_json={"headers": {"X": "y"}}) - - with pytest.raises(TimeoutError, match=r"execd endpoint.*not ready"): - await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", api_base=TEST_API_BASE, ready_timeout=0, client=client - ) - - -@pytest.mark.asyncio -async def test_create_reraises_non_404_endpoint_error(): - client = FakeHTTPClient(endpoint_responses=[http_status_error(500)]) - - with pytest.raises(httpx.HTTPStatusError): - await OpenSandboxSandboxConfig().acreate_sandbox( - api_key="", api_base=TEST_API_BASE, client=client - ) - - -@pytest.mark.asyncio -async def test_run_code_resolves_bare_id_and_posts_sse_request(): - client = FakeHTTPClient( - execute_lines=[ - sse({"type": "stdout", "text": "42\n"}), - ] - ) - - result = await OpenSandboxSandboxConfig().arun_code( - container="osb_bare", - code="print(6*7)", - language="python", - api_key="", - api_base="http://sandbox.local/v1", - client=client, - ) - - endpoint_call = client.calls[0] - run_call = client.calls[1] - assert endpoint_call[0] == "GET" - assert ( - endpoint_call[1] == "http://sandbox.local/v1/sandboxes/osb_bare/endpoints/44772" - ) - assert run_call[0] == "POST" - assert run_call[1] == "http://execd.local:44772/code" - assert run_call[2]["X-EXECD-ACCESS-TOKEN"] == "execd-token" - assert run_call[3] == { - "code": "print(6*7)", - "context": {"language": "python"}, - } - assert run_call[4] == {"stream": True} - assert result.stdout == "42\n" - - -@pytest.mark.asyncio -async def test_run_code_uses_https_for_scheme_less_endpoint_when_api_base_is_https(): - client = FakeHTTPClient() - handle = ContainerHandle( - id="osb_https", provider="opensandbox", domain="https://sandbox.example/v1" - ) - handle._hidden_params = { - "execd_endpoint": "execd.example/route/44772", - "execd_headers": {}, - } - - await OpenSandboxSandboxConfig().arun_code( - container=handle, code="print(1)", client=client - ) - - assert client.calls[0][1] == "https://execd.example/route/44772/code" - - -@pytest.mark.asyncio -async def test_run_code_aborts_on_output_over_cap(): - client = FakeHTTPClient(execute_lines=["x" * (MAX_OUTPUT_BYTES + 1)]) - handle = ContainerHandle(id="osb_big", provider="opensandbox", domain="http://x/v1") - handle._hidden_params = {"execd_endpoint": "execd.local:44772", "execd_headers": {}} - - with pytest.raises(ValueError, match="exceeded"): - await OpenSandboxSandboxConfig().arun_code( - container=handle, code="print('x')", client=client - ) - - -@pytest.mark.asyncio -async def test_delete_returns_false_on_404(): - client = FakeHTTPClient(delete_status=404) - - ok = await OpenSandboxSandboxConfig().adelete_sandbox( - container="osb_gone", - api_key="", - api_base="http://sandbox.local/v1", - client=client, - ) - - assert ok is False - - -@pytest.mark.asyncio -async def test_delete_reraises_non_404_http_error(): - client = FakeHTTPClient(delete_status=500) - - with pytest.raises(httpx.HTTPStatusError): - await OpenSandboxSandboxConfig().adelete_sandbox( - container="osb_err", - api_key="", - api_base="http://sandbox.local/v1", - client=client, - ) - - -@pytest.mark.asyncio -async def test_public_lifecycle_create_run_delete(): - client = FakeHTTPClient( - execute_lines=[ - sse({"type": "stdout", "text": "42\n"}), - ] - ) - - container = await litellm.acreate_sandbox( - provider="opensandbox", api_key="", api_base=TEST_API_BASE, client=client - ) - result = await litellm.arun_code( - provider="opensandbox", - container=container, - code="print(6*7)", - api_key="", - client=client, - ) - ok = await litellm.adelete_sandbox( - provider="opensandbox", - container=container, - api_key="", - client=client, - ) - - assert container.id == "osb_123" - assert result.stdout == "42\n" - assert ok is True - - -@pytest.mark.asyncio -async def test_code_interpreter_tool_deletes_even_when_run_raises(): - client = FakeHTTPClient(execute_raises=RuntimeError("boom")) - - with pytest.raises(RuntimeError, match="boom"): - await litellm.acode_interpreter_tool( - provider="opensandbox", - code="1/0", - api_key="", - api_base=TEST_API_BASE, - client=client, - ) - - assert [call[0] for call in client.calls] == ["POST", "GET", "POST", "DELETE"] - assert client.calls[0][1].endswith("/sandboxes") - assert client.calls[1][1].endswith("/endpoints/44772") - assert client.calls[2][1].endswith("/code") - assert client.calls[3][1].endswith("/sandboxes/osb_123") diff --git a/tests/test_litellm/sandbox/test_sandbox_tools.py b/tests/test_litellm/sandbox/test_sandbox_tools.py deleted file mode 100644 index 06136534b13..00000000000 --- a/tests/test_litellm/sandbox/test_sandbox_tools.py +++ /dev/null @@ -1,181 +0,0 @@ -"""Unit tests for the sandbox-tool registry.""" - -from litellm.sandbox import sandbox_tools - - -def _reset(): - sandbox_tools.clear_sandbox_tools() - - -def test_register_resolves_provider_key_and_base(): - _reset() - sandbox_tools.register_sandbox_tools( - [ - { - "sandbox_tool_name": "e2b_default", - "litellm_params": { - "sandbox_provider": "e2b", - "api_key": "sk-literal", - "api_base": "https://sandbox.internal", - }, - } - ] - ) - - resolved = sandbox_tools.resolve_sandbox_tool("e2b_default") - assert resolved == { - "sandbox_provider": "e2b", - "api_key": "sk-literal", - "api_base": "https://sandbox.internal", - } - _reset() - - -def test_register_clears_stale_entries_on_reload(): - """A tool removed from the config must not survive a re-registration.""" - _reset() - sandbox_tools.register_sandbox_tools( - [ - { - "sandbox_tool_name": "old", - "litellm_params": {"sandbox_provider": "e2b"}, - } - ] - ) - assert sandbox_tools.resolve_sandbox_tool("old") is not None - - sandbox_tools.register_sandbox_tools( - [ - { - "sandbox_tool_name": "new", - "litellm_params": {"sandbox_provider": "e2b"}, - } - ] - ) - - assert sandbox_tools.resolve_sandbox_tool("new") is not None - assert ( - sandbox_tools.resolve_sandbox_tool("old") is None - ), "stale tool must be gone after the config is reloaded" - _reset() - - -def test_register_empty_list_clears_removed_tools(): - """Reloading a config with sandbox_tools removed (the proxy passes an empty - list) must drop previously registered credentials from the process.""" - _reset() - sandbox_tools.register_sandbox_tools( - [ - { - "sandbox_tool_name": "e2b_default", - "litellm_params": {"sandbox_provider": "e2b", "api_key": "sk-x"}, - } - ] - ) - assert sandbox_tools.resolve_sandbox_tool("e2b_default") is not None - - sandbox_tools.register_sandbox_tools([]) - - assert ( - sandbox_tools.resolve_sandbox_tool("e2b_default") is None - ), "removing sandbox_tools from config must clear stale credentials" - _reset() - - -def test_register_resolves_secret_from_env(monkeypatch): - _reset() - monkeypatch.setenv("MY_SANDBOX_KEY", "sk-from-env") - sandbox_tools.register_sandbox_tools( - [ - { - "sandbox_tool_name": "e2b_default", - "litellm_params": { - "sandbox_provider": "e2b", - "api_key": "os.environ/MY_SANDBOX_KEY", - }, - } - ] - ) - - resolved = sandbox_tools.resolve_sandbox_tool("e2b_default") - assert resolved is not None - assert resolved["api_key"] == "sk-from-env" - assert resolved["api_base"] is None - _reset() - - -def test_resolve_unknown_returns_none(): - _reset() - assert sandbox_tools.resolve_sandbox_tool("nope") is None - - -def test_register_skips_malformed_entries_without_crashing(): - """A single malformed entry (missing sandbox_tool_name, or not a dict) must - not crash registration during proxy startup/hot-reload; valid entries in the - same list must still register.""" - _reset() - sandbox_tools.register_sandbox_tools( - [ - {"litellm_params": {"sandbox_provider": "e2b"}}, # missing name - "not-a-dict", # wrong type - {"sandbox_tool_name": "", "litellm_params": {}}, # empty name - { - "sandbox_tool_name": "good", - "litellm_params": {"sandbox_provider": "e2b"}, - }, - ] - ) - - assert sandbox_tools.resolve_sandbox_tool("good") is not None - assert sandbox_tools.resolve_sandbox_tool("") is None - assert set(sandbox_tools._SANDBOX_TOOL_REGISTRY) == {"good"} - _reset() - - -def test_register_skips_entry_missing_sandbox_provider(): - """An entry with a name but no sandbox_provider must be skipped at - registration so it cannot later resolve and call acreate_sandbox(provider=None), - which fails with a cryptic runtime error instead of a clear startup warning.""" - _reset() - sandbox_tools.register_sandbox_tools( - [ - {"sandbox_tool_name": "no_provider", "litellm_params": {"api_key": "sk-x"}}, - { - "sandbox_tool_name": "null_provider", - "litellm_params": {"sandbox_provider": None}, - }, - { - "sandbox_tool_name": "good", - "litellm_params": {"sandbox_provider": "e2b"}, - }, - ] - ) - - assert sandbox_tools.resolve_sandbox_tool("no_provider") is None - assert sandbox_tools.resolve_sandbox_tool("null_provider") is None - assert set(sandbox_tools._SANDBOX_TOOL_REGISTRY) == {"good"} - _reset() - - -def test_register_swaps_registry_atomically(): - """register_sandbox_tools must replace the registry in one rebind so a - concurrent resolve never observes a half-populated or transiently empty - registry between clearing and repopulating.""" - _reset() - sandbox_tools.register_sandbox_tools( - [{"sandbox_tool_name": "a", "litellm_params": {"sandbox_provider": "e2b"}}] - ) - before = sandbox_tools._SANDBOX_TOOL_REGISTRY - - sandbox_tools.register_sandbox_tools( - [ - {"sandbox_tool_name": "b", "litellm_params": {"sandbox_provider": "e2b"}}, - {"sandbox_tool_name": "c", "litellm_params": {"sandbox_provider": "e2b"}}, - ] - ) - after = sandbox_tools._SANDBOX_TOOL_REGISTRY - - assert after is not before, "the registry must be replaced, not mutated in place" - assert set(after) == {"b", "c"} - assert "a" not in after - _reset() diff --git a/tests/test_litellm/skills/test_skills_main.py b/tests/test_litellm/skills/test_skills_main.py deleted file mode 100644 index e1c66c8d9ea..00000000000 --- a/tests/test_litellm/skills/test_skills_main.py +++ /dev/null @@ -1,57 +0,0 @@ -from unittest.mock import MagicMock - -import litellm.skills.main as skills_main -from litellm.types.utils import LlmProviders - - -def test_create_skill_forwards_description_and_instructions_from_top_level_kwargs( - monkeypatch, -) -> None: - """The REST /v1/skills form endpoint passes description/instructions as top-level - kwargs (not extra_body). Regression for a bug where the litellm_proxy dispatch - branch of create_skill() dropped both, so every LiteLLM-hosted skill was created - with description=None and instructions=None regardless of what the caller sent.""" - handler = MagicMock() - monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) - - skills_main.create_skill( - display_title="Document Translator", - description="Converts files from one language into another", - instructions="Take an uploaded document and produce it in the target language", - custom_llm_provider=LlmProviders.LITELLM_PROXY.value, - ) - - assert handler.create_skill_handler.call_args.kwargs["description"] == ( - "Converts files from one language into another" - ) - assert handler.create_skill_handler.call_args.kwargs["instructions"] == ( - "Take an uploaded document and produce it in the target language" - ) - - -def test_create_skill_forwards_description_and_instructions_from_extra_body(monkeypatch) -> None: - """The SDK convention (see tests/proxy_unit_tests/test_skills_db.py) nests them under - extra_body instead of passing them as top-level kwargs; both paths must reach the DB.""" - handler = MagicMock() - monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) - - skills_main.create_skill( - display_title="Warehouse SQL Analyst", - extra_body={"description": "Runs SQL against the inventory database", "instructions": "Summarize results"}, - custom_llm_provider=LlmProviders.LITELLM_PROXY.value, - ) - - assert handler.create_skill_handler.call_args.kwargs["description"] == ( - "Runs SQL against the inventory database" - ) - assert handler.create_skill_handler.call_args.kwargs["instructions"] == "Summarize results" - - -def test_create_skill_without_description_or_instructions_passes_none(monkeypatch) -> None: - handler = MagicMock() - monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) - - skills_main.create_skill(display_title="Bare Skill", custom_llm_provider=LlmProviders.LITELLM_PROXY.value) - - assert handler.create_skill_handler.call_args.kwargs["description"] is None - assert handler.create_skill_handler.call_args.kwargs["instructions"] is None diff --git a/tests/test_litellm/test_router/test_enforce_model_rate_limits.py b/tests/test_litellm/test_router/test_enforce_model_rate_limits.py deleted file mode 100644 index 7577064b7f9..00000000000 --- a/tests/test_litellm/test_router/test_enforce_model_rate_limits.py +++ /dev/null @@ -1,468 +0,0 @@ -""" -Tests for enforce_model_rate_limits feature. - -This feature allows users to enforce TPM/RPM limits set on model deployments -regardless of the routing strategy being used. -""" - -import asyncio -from datetime import timedelta -from unittest.mock import AsyncMock, MagicMock - -import pytest - -import litellm -from litellm import Router -from litellm.caching.dual_cache import DualCache -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError -from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( - ModelRateLimitingCheck, -) - -TPM_DEPLOYMENT = { - "tpm": 1000, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "replica-test-id"}, - "model_name": "test-model", -} - - -def _dual_cache_with_local_tpm(local_tpm: int, redis_cache: MagicMock | None) -> DualCache: - dual_cache = DualCache(redis_cache=redis_cache) - check = ModelRateLimitingCheck(dual_cache=dual_cache) - now = litellm.utils.get_utc_datetime() - for minute in (now, now + timedelta(minutes=1)): - tpm_key, _ = check._get_cache_keys(TPM_DEPLOYMENT, minute.strftime("%H-%M")) - dual_cache.set_cache(key=tpm_key, value=local_tpm, local_only=True) - return dual_cache - - -class TestModelRateLimitingCheck: - """Test the ModelRateLimitingCheck class directly.""" - - def test_get_deployment_limits_from_top_level(self): - """Test extracting limits from top-level deployment config.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) - - deployment = { - "tpm": 1000, - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - } - - tpm, rpm = check._get_deployment_limits(deployment) - assert tpm == 1000 - assert rpm == 10 - - def test_get_deployment_limits_from_litellm_params(self): - """Test extracting limits from litellm_params.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) - - deployment = { - "litellm_params": {"model": "gpt-4", "tpm": 2000, "rpm": 20}, - "model_info": {"id": "test-id"}, - } - - tpm, rpm = check._get_deployment_limits(deployment) - assert tpm == 2000 - assert rpm == 20 - - def test_get_deployment_limits_from_model_info(self): - """Test extracting limits from model_info.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) - - deployment = { - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id", "tpm": 3000, "rpm": 30}, - } - - tpm, rpm = check._get_deployment_limits(deployment) - assert tpm == 3000 - assert rpm == 30 - - def test_get_deployment_limits_none_when_not_set(self): - """Test that None is returned when limits are not set.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) - - deployment = { - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - } - - tpm, rpm = check._get_deployment_limits(deployment) - assert tpm is None - assert rpm is None - - def test_pre_call_check_allows_request_when_no_limits(self): - """Test that requests are allowed when no limits are set.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) - - deployment = { - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - } - - result = check.pre_call_check(deployment) - assert result == deployment - - def test_pre_call_check_raises_rate_limit_error_when_over_rpm(self): - """Test that RateLimitError is raised when RPM limit is exceeded.""" - mock_cache = MagicMock() - mock_cache.increment_cache.return_value = 11 # Over limit after increment - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(deployment) - - assert "RPM limit=10" in str(exc_info.value) - assert "current usage=11" in str(exc_info.value) - - def test_pre_call_check_allows_request_under_limit(self): - """Test that requests are allowed when under the limit.""" - mock_cache = MagicMock() - mock_cache.increment_cache.return_value = 6 - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - result = check.pre_call_check(deployment) - assert result == deployment - - def test_pre_call_check_raises_rate_limit_error_when_over_tpm(self): - """Test that RateLimitError is raised when TPM limit is exceeded.""" - mock_cache = MagicMock() - mock_cache.get_cache.return_value = 1000 # Already at limit - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "tpm": 1000, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(deployment) - - assert "TPM limit=1000" in str(exc_info.value) - assert "current usage=1000" in str(exc_info.value) - - def test_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): - redis_cache = MagicMock() - redis_cache.get_cache.return_value = 1000 - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(TPM_DEPLOYMENT) - - assert "current usage=1000" in str(exc_info.value) - - @pytest.mark.parametrize( - "redis_get", [MagicMock(return_value=None), MagicMock(side_effect=RedisCircuitBreakerOpenError())] - ) - def test_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): - redis_cache = MagicMock() - redis_cache.get_cache = redis_get - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(TPM_DEPLOYMENT) - - assert "current usage=1000" in str(exc_info.value) - - def test_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open(self): - redis_cache = MagicMock() - redis_cache.get_cache.side_effect = RedisCircuitBreakerOpenError() - redis_cache.increment_cache.return_value = 2 - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) - - assert "RPM limit=1" in str(exc_info.value) - - def test_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) - deployment = {**TPM_DEPLOYMENT, "rpm": 1} - - assert check.pre_call_check(deployment) == deployment - with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(deployment) - - assert "RPM limit=1" in str(exc_info.value) - - def test_log_success_event_increments_cache(self): - """Test that log_success_event correctly increments the cache.""" - mock_cache = MagicMock() - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - kwargs = { - "standard_logging_object": { - "model_id": "test-id", - "total_tokens": 50, - "hidden_params": {"litellm_model_name": "gpt-4"}, - } - } - - check.log_success_event(kwargs, None, None, None) - - # Verify increment_cache was called - mock_cache.increment_cache.assert_called_once() - _, kwarg_params = mock_cache.increment_cache.call_args - assert "test-id:gpt-4:tpm:" in kwarg_params["key"] - assert kwarg_params["value"] == 50 - - -class TestModelRateLimitingCheckAsync: - """Test async methods of ModelRateLimitingCheck.""" - - @pytest.mark.asyncio - async def test_async_pre_call_check_allows_request_when_no_limits(self): - """Test that requests are allowed when no limits are set (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - } - - result = await check.async_pre_call_check(deployment) - assert result == deployment - - @pytest.mark.asyncio - async def test_async_pre_call_check_raises_rate_limit_error_when_over_rpm(self): - """Test that RateLimitError is raised when RPM limit is exceeded (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache = AsyncMock(return_value=11) # Over limit - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) - - assert "RPM limit=10" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_async_pre_call_check_allows_request_under_limit(self): - """Test that requests are allowed when under the limit (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache = AsyncMock(return_value=6) - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - result = await check.async_pre_call_check(deployment) - assert result == deployment - - @pytest.mark.asyncio - async def test_async_pre_call_check_raises_rate_limit_error_when_over_tpm(self): - """Test that RateLimitError is raised when TPM limit is exceeded (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=1000) # Already at limit - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "tpm": 1000, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) - - assert "TPM limit=1000" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_async_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(return_value=1000) - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(TPM_DEPLOYMENT) - - assert "current usage=1000" in str(exc_info.value) - - @pytest.mark.asyncio - @pytest.mark.parametrize( - "redis_get", [AsyncMock(return_value=None), AsyncMock(side_effect=RedisCircuitBreakerOpenError())] - ) - async def test_async_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): - redis_cache = MagicMock() - redis_cache.async_get_cache = redis_get - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(TPM_DEPLOYMENT) - - assert "current usage=1000" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_async_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open( - self, - ): - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(side_effect=RedisCircuitBreakerOpenError()) - redis_cache.async_increment = AsyncMock(return_value=2) - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) - - assert "RPM limit=1" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_async_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) - deployment = {**TPM_DEPLOYMENT, "rpm": 1} - - assert await check.async_pre_call_check(deployment) == deployment - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) - - assert "RPM limit=1" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_async_log_success_event_increments_cache(self): - """Test that async_log_success_event correctly increments the cache.""" - mock_cache = MagicMock() - mock_cache.async_increment_cache = AsyncMock() - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - kwargs = { - "standard_logging_object": { - "model_id": "test-id", - "total_tokens": 50, - "hidden_params": {"litellm_model_name": "gpt-4"}, - } - } - - await check.async_log_success_event(kwargs, None, None, None) - - # Verify async_increment_cache was called - mock_cache.async_increment_cache.assert_called_once() - _, kwarg_params = mock_cache.async_increment_cache.call_args - assert "test-id:gpt-4:tpm:" in kwarg_params["key"] - assert kwarg_params["value"] == 50 - - -class TestRouterWithEnforceModelRateLimits: - """Test Router integration with enforce_model_rate_limits.""" - - def test_router_initializes_with_enforce_model_rate_limits(self): - """Test that Router properly initializes the ModelRateLimitingCheck.""" - model_list = [ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "test"}, - "rpm": 10, - } - ] - - router = Router( - model_list=model_list, - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - - # Check that the callback was added - assert router.optional_callbacks is not None - assert len(router.optional_callbacks) == 1 - assert isinstance(router.optional_callbacks[0], ModelRateLimitingCheck) - - def test_router_optional_callbacks_contains_model_rate_limiting(self): - """Test that ModelRateLimitingCheck is in the callbacks list.""" - model_list = [ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "test"}, - "rpm": 10, - } - ] - - Router( - model_list=model_list, - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - - # Find the ModelRateLimitingCheck in litellm.callbacks - found = False - for callback in litellm.callbacks: - if isinstance(callback, ModelRateLimitingCheck): - found = True - break - - assert found, "ModelRateLimitingCheck should be in litellm.callbacks" - - -class TestModelRateLimitConcurrency: - """Test that RPM rate limiting is atomic under concurrent requests.""" - - @pytest.mark.asyncio - async def test_concurrent_requests_respect_rpm_limit(self): - """ - Fire 4 concurrent async requests with RPM limit of 2. - Exactly 2 should succeed and 2 should raise RateLimitError. - - This test validates the atomic increment-first pattern: - the old check-then-increment pattern would let 3+ through - due to a race condition on the local cache read. - """ - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - - deployment = { - "rpm": 2, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "concurrent-test-id"}, - "model_name": "test-model", - } - - async def attempt_request(): - return await check.async_pre_call_check(deployment) - - results = await asyncio.gather( - *[attempt_request() for _ in range(4)], - return_exceptions=True, - ) - - successes = [r for r in results if not isinstance(r, Exception)] - failures = [r for r in results if isinstance(r, litellm.RateLimitError)] - - assert len(successes) == 2, f"Expected 2 successes, got {len(successes)}" - assert len(failures) == 2, f"Expected 2 rate limit errors, got {len(failures)}" diff --git a/tests/test_litellm/test_router/test_io_token_rate_limits.py b/tests/test_litellm/test_router/test_io_token_rate_limits.py deleted file mode 100644 index 3cef1c7bb63..00000000000 --- a/tests/test_litellm/test_router/test_io_token_rate_limits.py +++ /dev/null @@ -1,1069 +0,0 @@ -""" -Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits). -""" - -import asyncio - -import pytest - -import litellm -from litellm import Router -from litellm.caching.dual_cache import DualCache -from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( - ITPM_CACHE_KEY, - ITPM_RESERVED_KEY, - OTPM_CACHE_KEY, - OTPM_RESERVED_KEY, - _reservation_value, - _resolve_max_tokens, - async_io_token_pre_call_check, - async_io_token_reconcile_success, - build_io_token_rate_limit_headers, - deployment_has_io_token_limits, - get_io_token_rate_limit_request_kwargs, - io_token_reconcile_success, - io_token_refund_failure, - refund_stale_reservation_before_retry, - set_io_token_rate_limit_request_kwargs, -) -from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( - ModelRateLimitingCheck, -) -from litellm.types.utils import ModelResponse, Usage - - -class TestIOTokenRateLimitHelpers: - def test_deployment_has_io_token_limits(self): - assert deployment_has_io_token_limits({"litellm_params": {"itpm": 100, "otpm": 50}}) - assert not deployment_has_io_token_limits({"litellm_params": {"model": "x"}}) - - def test_reservation_value_minimal_when_estimate_fails(self): - # A failed/empty estimate (0) must reserve a minimal slot, not the - # entire limit - otherwise one request whose estimate failed fills - # the whole bucket and blocks every concurrent request until it - # completes and reconciles. - assert _reservation_value(0, 100) == 1 - assert _reservation_value(0, 1) == 1 - # A real non-zero estimate is reserved as-is. - assert _reservation_value(42, 100) == 42 - - def test_resolve_max_tokens_respects_explicit_zero(self): - deployment = {"litellm_params": {"model": "openai/gpt-4o-mini"}} - # An explicit max_tokens=0 is honored, not replaced by the model default. - assert _resolve_max_tokens({"max_tokens": 0}, deployment) == 0 - # max_completion_tokens is the fallback only when max_tokens is absent. - assert _resolve_max_tokens({"max_completion_tokens": 12}, deployment) == 12 - assert _resolve_max_tokens({"max_output_tokens": 9}, deployment) == 9 - - def test_build_io_token_rate_limit_headers(self): - headers = build_io_token_rate_limit_headers( - itpm_limit=200, - otpm_limit=40, - current_itpm=15, - current_otpm=4, - ) - assert headers["x-ratelimit-limit-input-tokens"] == 200 - assert headers["x-ratelimit-remaining-input-tokens"] == 185 - assert headers["x-ratelimit-limit-output-tokens"] == 40 - assert headers["x-ratelimit-remaining-output-tokens"] == 36 - - -class TestModelRateLimitingCheckIOTokens: - @pytest.mark.asyncio - async def test_itpm_reservation_and_reconcile(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 100, - "otpm": 50, - }, - "model_info": {"id": "io-test-id"}, - "model_name": "opus", - } - - request_kwargs = { - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 10, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(request_kwargs) - await check.async_pre_call_check(deployment) - - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - otpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" - - kwargs = { - "standard_logging_object": { - "model_id": "io-test-id", - "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, - "metadata": dict(request_kwargs["metadata"]), - }, - "metadata": request_kwargs["metadata"], - } - response = ModelResponse( - choices=[ - { - "message": {"role": "assistant", "content": "hi"}, - "index": 0, - "finish_reason": "stop", - } - ], - usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8), - ) - await check.async_log_success_event(kwargs, response, None, None) - - current_itpm = await dual_cache.async_get_cache(key=itpm_key) - current_otpm = await dual_cache.async_get_cache(key=otpm_key) - # ITPM tracks input tokens only (billable prompt tokens), not output. - assert current_itpm == 5 - assert current_otpm == 3 - - @pytest.mark.asyncio - async def test_itpm_limit_raises_429(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 5, - }, - "model_info": {"id": "io-limit-id"}, - "model_name": "opus", - } - - # ITPM enforces input tokens only; the prompt alone must exceed the limit, - # a large max_tokens must not contribute to the ITPM reservation. - set_io_token_rate_limit_request_kwargs( - { - "messages": [ - { - "role": "user", - "content": "hello world this is a longer prompt that exceeds the tiny itpm limit", - } - ], - "max_tokens": 10, - "metadata": {}, - } - ) - - with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) - - assert "ITPM limit=5" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_otpm_atomic_reservation_no_overshoot_under_concurrency(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - otpm_limit = 10 - max_tokens = 4 - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "otpm": otpm_limit, - }, - "model_info": {"id": "io-otpm-race-id"}, - "model_name": "opus", - } - - set_io_token_rate_limit_request_kwargs( - { - "messages": [{"role": "user", "content": "hi"}], - "max_tokens": max_tokens, - "metadata": {}, - } - ) - - async def _attempt(): - try: - await check.async_pre_call_check(deployment) - return True - except litellm.RateLimitError: - return False - - results = await asyncio.gather(*[_attempt() for _ in range(8)]) - successes = sum(1 for r in results if r) - - minute = get_utc_datetime().strftime("%H-%M") - otpm_key = f"global_router:io-otpm-race-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" - current_otpm = await dual_cache.async_get_cache(key=otpm_key) - - # Atomic reservation must never let concurrent requests overshoot the limit. - assert current_otpm is not None - assert current_otpm <= otpm_limit - assert successes == otpm_limit // max_tokens - assert current_otpm == successes * max_tokens - - @pytest.mark.asyncio - async def test_itpm_estimate_failure_reserves_minimal_not_full_limit(self): - """ - When input-token estimation yields 0 (no messages/prompt/input field, - unsupported model, tokenizer error), the reservation must be a - minimal 1 token, not the entire itpm limit. Otherwise the first - request whose estimate fails fills the whole bucket and every - concurrent request is rejected until it completes - effectively - serializing traffic to the deployment. - """ - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - itpm_limit = 5 - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": itpm_limit, - }, - "model_info": {"id": "io-itpm-estimate-fail-id"}, - "model_name": "opus", - } - - # No messages/prompt/input field -> _estimate_input_tokens returns 0. - set_io_token_rate_limit_request_kwargs( - { - "max_tokens": 5, - "metadata": {}, - } - ) - - async def _attempt(): - try: - await check.async_pre_call_check(deployment) - return True - except litellm.RateLimitError: - return False - - results = await asyncio.gather(*[_attempt() for _ in range(8)]) - successes = sum(1 for r in results if r) - - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-itpm-estimate-fail-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - current_itpm = await dual_cache.async_get_cache(key=itpm_key) - - # A minimal 1-token reservation per request lets itpm_limit concurrent - # requests through, instead of a single request starving the rest. - assert current_itpm is not None - assert current_itpm <= itpm_limit - assert successes == itpm_limit - - @pytest.mark.asyncio - async def test_reservation_read_prefers_top_level_metadata_over_litellm_params(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-lp-id:bedrock_mantle/test:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) - - # Production kwargs commonly carry litellm_params.metadata; the stashed - # reservation lives in the top-level metadata and must still be found. - kwargs = { - "standard_logging_object": { - "model_id": "io-lp-id", - "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, - "metadata": {}, - }, - "metadata": {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key}, - "litellm_params": {"metadata": {"user_api_key_hash": "abc123"}}, - } - await check.async_log_failure_event(kwargs, None, None, None) - - current = await dual_cache.async_get_cache(key=itpm_key) - assert current == 0 - - @pytest.mark.asyncio - async def test_reconcile_tracks_actual_usage_when_estimate_zero(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 100, - }, - "model_info": {"id": "io-zero-est-id"}, - "model_name": "opus", - } - - request_kwargs = {"max_tokens": 5, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) - await check.async_pre_call_check(deployment) - - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-zero-est-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - # A failed/zero estimate reserves a minimal 1 token, not the full - # itpm limit, so it doesn't starve concurrent requests. - assert await dual_cache.async_get_cache(key=itpm_key) == 1 - - kwargs = { - "standard_logging_object": { - "model_id": "io-zero-est-id", - "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, - "metadata": dict(request_kwargs["metadata"]), - }, - "metadata": request_kwargs["metadata"], - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=7, completion_tokens=0, total_tokens=7), - ) - await check.async_log_success_event(kwargs, response, None, None) - - assert await dual_cache.async_get_cache(key=itpm_key) == 7 - - @pytest.mark.asyncio - async def test_zero_estimate_reserves_minimal_capacity_before_reconcile(self): - """ - A zero/failed estimate reserves a minimal 1 token rather than the - full itpm limit, so up to itpm_limit such calls are allowed - concurrently instead of the first one claiming the entire bucket. - """ - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 2, - }, - "model_info": {"id": "io-zero-cap-id"}, - "model_name": "opus", - } - request_kwargs = {"max_tokens": 5, "metadata": {}} - set_io_token_rate_limit_request_kwargs(request_kwargs) - await check.async_pre_call_check(deployment) - - # Second zero-estimate call still fits within the itpm=2 limit. - set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) - await check.async_pre_call_check(deployment) - - # A third exceeds the limit and is rejected. - set_io_token_rate_limit_request_kwargs({"max_tokens": 5, "metadata": {}}) - with pytest.raises(litellm.RateLimitError): - await check.async_pre_call_check(deployment) - - @pytest.mark.asyncio - async def test_explicit_zero_max_tokens_does_not_reserve_otpm(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "otpm": 5, - }, - "model_info": {"id": "io-zero-output-id"}, - "model_name": "opus", - } - zero_output_kwargs = { - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 0, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(zero_output_kwargs) - await check.async_pre_call_check(deployment) - - zero_output_otpm_key = zero_output_kwargs["metadata"][OTPM_CACHE_KEY] - assert zero_output_kwargs["metadata"][OTPM_RESERVED_KEY] == 0 - assert (await dual_cache.async_get_cache(key=zero_output_otpm_key) or 0) == 0 - - normal_output_kwargs = { - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 5, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(normal_output_kwargs) - await check.async_pre_call_check(deployment) - - normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY] - assert await dual_cache.async_get_cache(key=normal_output_otpm_key) == 5 - - def test_sync_io_pre_call_reserves_and_reconciles(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 100, - "otpm": 50, - }, - "model_info": {"id": "io-sync-id"}, - "model_name": "opus", - } - request_kwargs = { - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 10, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(request_kwargs) - check.pre_call_check(deployment) - - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - otpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" - kwargs = { - "standard_logging_object": { - "model_id": "io-sync-id", - "hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"}, - "metadata": dict(request_kwargs["metadata"]), - }, - "metadata": request_kwargs["metadata"], - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8), - ) - check.log_success_event(kwargs, response, None, None) - - assert dual_cache.get_cache(key=itpm_key) == 5 - assert dual_cache.get_cache(key=otpm_key) == 3 - - @pytest.mark.asyncio - async def test_reconcile_runs_via_success_event_without_model_id(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - itpm_key = "global_router:io-noid:bedrock_mantle/test:itpm:12-34" - await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) - - # standard_logging_object has no model_id (only the TPM path needs it); - # IO reconciliation must still run off the stashed cache key. - kwargs = { - "standard_logging_object": { - "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, - "metadata": {}, - }, - "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), - ) - await check.async_log_success_event(kwargs, response, None, None) - - assert await dual_cache.async_get_cache(key=itpm_key) == 3 - - @pytest.mark.asyncio - async def test_failure_clears_reservation_so_retry_is_not_poisoned(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-first:bedrock_mantle/test:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) - - # Shared request metadata carrying the first (IO) deployment's reservation. - metadata = {ITPM_RESERVED_KEY: 8, ITPM_CACHE_KEY: itpm_key} - fail_kwargs = { - "metadata": metadata, - "standard_logging_object": { - "model_id": "io-first", - "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, - "metadata": {}, - }, - } - await check.async_log_failure_event(fail_kwargs, None, None, None) - - assert await dual_cache.async_get_cache(key=itpm_key) == 0 - assert ITPM_RESERVED_KEY not in metadata - assert ITPM_CACHE_KEY not in metadata - - # Retry succeeds on a non-IO fallback deployment reusing the same metadata. - retry_kwargs = { - "metadata": metadata, - "standard_logging_object": { - "model_id": "non-io-second", - "hidden_params": {"litellm_model_name": "openai/gpt-4o-mini"}, - "metadata": {}, - "total_tokens": 12, - }, - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=6, completion_tokens=6, total_tokens=12), - ) - await check.async_log_success_event(retry_kwargs, response, None, None) - - # The first deployment's ITPM counter is not driven negative... - assert await dual_cache.async_get_cache(key=itpm_key) == 0 - # ...and the non-IO deployment's TPM usage is tracked normally. - tpm_key = f"non-io-second:openai/gpt-4o-mini:tpm:{minute}" - assert await dual_cache.async_get_cache(key=tpm_key) == 12 - - @pytest.mark.asyncio - async def test_stale_reservation_refunded_before_retry_overwrites_it(self): - """ - A retry reuses the same mutable kwargs dict for the next deployment. - If deployment A's failure event hasn't run yet (e.g. it was scheduled - as a background task) when the retry calls - set_io_token_rate_limit_request_kwargs for deployment B, the router - must first synchronously refund + clear A's reservation via - refund_stale_reservation_before_retry - otherwise A's counter stays - elevated by the reservation until its TTL expires, and the - now-orphaned sentinels must not leak into B's accounting either. - """ - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - minute = get_utc_datetime().strftime("%H-%M") - itpm_key_a = f"global_router:io-retry-a:bedrock_mantle/test-a:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key_a, value=9, ttl=60) - - # Deployment A's still-unreconciled reservation, stashed on the shared - # kwargs dict the retry loop reuses. - shared_kwargs = {"metadata": {ITPM_RESERVED_KEY: 9, ITPM_CACHE_KEY: itpm_key_a}} - - # Router calls this before overwriting kwargs for deployment B's attempt - - # simulating the fix landing ahead of set_io_token_rate_limit_request_kwargs. - refund_stale_reservation_before_retry(dual_cache, shared_kwargs) - - # A's reservation is refunded immediately, not left stranded for a - # background failure task that may run arbitrarily later (or never, - # if the sentinels get cleared out from under it first). - assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 - assert ITPM_RESERVED_KEY not in shared_kwargs["metadata"] - assert ITPM_CACHE_KEY not in shared_kwargs["metadata"] - - # A's own (now-late) failure event finds nothing left to refund and - # is a safe no-op, since the sentinels were already cleared above. - io_token_refund_failure(dual_cache, shared_kwargs) - assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 - - # The retry proceeds to stash deployment B's own reservation on the - # same dict; it starts clean, unaffected by A's cleared sentinels. - set_io_token_rate_limit_request_kwargs(shared_kwargs) - itpm_key_b = f"global_router:io-retry-b:bedrock_mantle/test-b:itpm:{minute}" - shared_kwargs["metadata"][ITPM_RESERVED_KEY] = 4 - shared_kwargs["metadata"][ITPM_CACHE_KEY] = itpm_key_b - await dual_cache.async_increment_cache(key=itpm_key_b, value=4, ttl=60) - assert await dual_cache.async_get_cache(key=itpm_key_b) == 4 - - @pytest.mark.asyncio - async def test_client_supplied_reservation_keys_are_stripped(self): - # metadata is caller-controlled; the server-only reservation sentinels - # must be removed before the router captures the request kwargs. - forged = { - "metadata": {ITPM_RESERVED_KEY: 999999, ITPM_CACHE_KEY: "attacker:key:itpm:00-00"}, - "litellm_metadata": {OTPM_RESERVED_KEY: 7}, - "litellm_params": {"metadata": {OTPM_CACHE_KEY: "attacker:key:otpm:00-00"}}, - } - set_io_token_rate_limit_request_kwargs(forged) - stored = get_io_token_rate_limit_request_kwargs() - - assert ITPM_RESERVED_KEY not in stored["metadata"] - assert ITPM_CACHE_KEY not in stored["metadata"] - assert OTPM_RESERVED_KEY not in stored["litellm_metadata"] - assert OTPM_CACHE_KEY not in stored["litellm_params"]["metadata"] - - @pytest.mark.asyncio - async def test_forged_reservation_cannot_decrement_counter(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - victim_key = "global_router:victim:model:itpm:00-00" - await dual_cache.async_increment_cache(key=victim_key, value=100, ttl=60) - - # A caller forges a reservation pointing at another deployment's counter. - kwargs = { - "metadata": {ITPM_RESERVED_KEY: 100, ITPM_CACHE_KEY: victim_key}, - "standard_logging_object": { - "model_id": "m", - "hidden_params": {"litellm_model_name": "model"}, - "metadata": {}, - "total_tokens": 2, - }, - } - # The router sanitizes the request kwargs before the call runs. - set_io_token_rate_limit_request_kwargs(kwargs) - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), - ) - await check.async_log_success_event(kwargs, response, None, None) - - # The forged reservation was stripped, so the victim counter is untouched. - assert await dual_cache.async_get_cache(key=victim_key) == 100 - - @pytest.mark.asyncio - async def test_otpm_reservation_error_rolls_back_itpm(self): - from litellm.utils import get_utc_datetime - - class _OtpmFailCache(DualCache): - async def async_increment_cache(self, key, **kwargs): - if ":otpm:" in key: - raise RuntimeError("transient cache error") - return await super().async_increment_cache(key=key, **kwargs) - - dual_cache = _OtpmFailCache() - deployment = { - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 1000, - "otpm": 1000, - }, - "model_info": {"id": "io-rollback-id"}, - "model_name": "opus", - } - set_io_token_rate_limit_request_kwargs( - { - "messages": [{"role": "user", "content": "hello world"}], - "max_tokens": 5, - "metadata": {}, - } - ) - - with pytest.raises(RuntimeError): - await async_io_token_pre_call_check(dual_cache, deployment) - - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-rollback-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - # A transient OTPM error must release the ITPM reservation, not leak it. - assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 - - @pytest.mark.asyncio - async def test_reconcile_clears_stash_even_when_increment_errors(self): - class _ItpmFailCache(DualCache): - async def async_increment_cache(self, key, **kwargs): - if ":itpm:" in key: - raise RuntimeError("transient cache error") - return await super().async_increment_cache(key=key, **kwargs) - - dual_cache = _ItpmFailCache() - metadata = {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: "global_router:x:model:itpm:00-00"} - kwargs = {"metadata": metadata} - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), - ) - - with pytest.raises(RuntimeError): - await async_io_token_reconcile_success(dual_cache, kwargs, response) - - # The stash is cleared even though reconciliation raised, so a duplicate - # success event can't re-process it. - assert ITPM_RESERVED_KEY not in metadata - assert ITPM_CACHE_KEY not in metadata - - @pytest.mark.asyncio - async def test_io_conflict_warning_not_collapsed_for_missing_model_id(self, caplog): - import logging - - check = ModelRateLimitingCheck(dual_cache=DualCache()) - deployment = { - "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100, "tpm": 1000}, - "model_info": {}, - } - with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): - check._warn_io_token_and_tpm_rpm_coexist_once(deployment) - check._warn_io_token_and_tpm_rpm_coexist_once(deployment) - - warnings = [r for r in caplog.records if "both limit types are enforced" in r.message] - # id-less deployments are not collapsed onto a single dedup key. - assert len(warnings) == 2 - - @pytest.mark.asyncio - async def test_missing_deployment_id_skips_io_reservation(self): - dual_cache = DualCache() - deployment = { - "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100}, - "model_info": {}, # no id -> cannot build a per-deployment cache key - "model_name": "opus", - } - request_kwargs = { - "messages": [{"role": "user", "content": "hello world"}], - "max_tokens": 5, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(request_kwargs) - - result = await async_io_token_pre_call_check(dual_cache, deployment) - - assert result is deployment - # No reservation is stashed, so nothing lands in a shared None:None bucket. - assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] - - @pytest.mark.asyncio - async def test_reconcile_uses_reservation_minute_key(self): - dual_cache = DualCache() - # Reservation was made on a fixed minute key; a call that finishes in a - # later minute must reconcile against that same key, never a key built - # from the response-time minute. - itpm_key = "global_router:io-min-id:bedrock_mantle/test:itpm:99-99" - await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) - - kwargs = {"metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}} - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=4, completion_tokens=0, total_tokens=4), - ) - await async_io_token_reconcile_success(dual_cache, kwargs, response) - - assert await dual_cache.async_get_cache(key=itpm_key) == 4 - - @pytest.mark.asyncio - async def test_reconcile_missing_usage_keeps_reservation(self): - dual_cache = DualCache() - itpm_key = "global_router:io-missing-usage:bedrock_mantle/test:itpm:00-00" - otpm_key = "global_router:io-missing-usage:bedrock_mantle/test:otpm:00-00" - await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) - await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) - - kwargs = { - "metadata": { - ITPM_RESERVED_KEY: 8, - OTPM_RESERVED_KEY: 5, - ITPM_CACHE_KEY: itpm_key, - OTPM_CACHE_KEY: otpm_key, - } - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], - ) - - await async_io_token_reconcile_success(dual_cache, kwargs, response) - - assert await dual_cache.async_get_cache(key=itpm_key) == 8 - assert await dual_cache.async_get_cache(key=otpm_key) == 5 - assert ITPM_RESERVED_KEY not in kwargs["metadata"] - - @pytest.mark.asyncio - async def test_reconcile_total_tokens_only_keeps_reservation(self): - """ - A response usage object with only total_tokens (no prompt/completion - breakdown) can't be split into input/output, so it must be treated the - same as missing usage: keep the reservation instead of resolving to - (0, 0) and refunding it in full. - """ - dual_cache = DualCache() - itpm_key = "global_router:io-total-only:bedrock_mantle/test:itpm:00-00" - otpm_key = "global_router:io-total-only:bedrock_mantle/test:otpm:00-00" - await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) - await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) - - kwargs = { - "metadata": { - ITPM_RESERVED_KEY: 8, - OTPM_RESERVED_KEY: 5, - ITPM_CACHE_KEY: itpm_key, - OTPM_CACHE_KEY: otpm_key, - } - } - response = {"type": "message", "usage": {"total_tokens": 13}} - - await async_io_token_reconcile_success(dual_cache, kwargs, response) - - assert await dual_cache.async_get_cache(key=itpm_key) == 8 - assert await dual_cache.async_get_cache(key=otpm_key) == 5 - - def test_reconcile_standard_logging_total_tokens_only_keeps_reservation(self): - dual_cache = DualCache() - itpm_key = "global_router:io-slo-total-only:bedrock_mantle/test:itpm:00-00" - dual_cache.set_cache(key=itpm_key, value=10, ttl=60) - - kwargs = { - "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, - "standard_logging_object": {"total_tokens": 4}, - } - response = {"type": "message", "role": "assistant", "content": []} - - io_token_reconcile_success(dual_cache, kwargs, response) - - assert dual_cache.get_cache(key=itpm_key) == 10 - - @pytest.mark.asyncio - async def test_reconcile_falls_back_to_standard_logging_object(self): - dual_cache = DualCache() - itpm_key = "global_router:io-slo-fallback:bedrock_mantle/test:itpm:00-00" - await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) - - kwargs = { - "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, - "standard_logging_object": { - "prompt_tokens": 4, - "completion_tokens": 0, - "total_tokens": 4, - }, - } - response = {"type": "message", "role": "assistant", "content": []} - - await async_io_token_reconcile_success(dual_cache, kwargs, response) - - assert await dual_cache.async_get_cache(key=itpm_key) == 4 - - def test_sync_reconcile_anthropic_dict_usage(self): - dual_cache = DualCache() - itpm_key = "global_router:io-anthropic:bedrock_mantle/test:itpm:00-00" - otpm_key = "global_router:io-anthropic:bedrock_mantle/test:otpm:00-00" - dual_cache.set_cache(key=itpm_key, value=6, ttl=60) - dual_cache.set_cache(key=otpm_key, value=4, ttl=60) - - kwargs = { - "metadata": { - ITPM_RESERVED_KEY: 6, - OTPM_RESERVED_KEY: 4, - ITPM_CACHE_KEY: itpm_key, - OTPM_CACHE_KEY: otpm_key, - } - } - response = { - "type": "message", - "usage": {"input_tokens": 3, "output_tokens": 2, "cache_read_input_tokens": 1}, - } - - io_token_reconcile_success(dual_cache, kwargs, response) - - assert dual_cache.get_cache(key=itpm_key) == 2 - assert dual_cache.get_cache(key=otpm_key) == 2 - - @pytest.mark.asyncio - async def test_io_and_tpm_rpm_limits_both_enforced_with_warning(self, caplog): - import logging - - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - model_id = "io-mixed-id" - deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" - deployment = { - "litellm_params": { - "model": deployment_name, - "itpm": 100, - "rpm": 1, - }, - "model_info": {"id": model_id}, - "model_name": "opus", - } - - minute = get_utc_datetime().strftime("%H-%M") - rpm_key = f"{model_id}:{deployment_name}:rpm:{minute}" - itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" - await dual_cache.async_increment_cache(key=rpm_key, value=5, ttl=60) - - request_kwargs = { - "messages": [{"role": "user", "content": "hi"}], - "max_tokens": 5, - "metadata": {}, - } - set_io_token_rate_limit_request_kwargs(request_kwargs) - - with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): - with pytest.raises(litellm.RateLimitError): - await check.async_pre_call_check(deployment) - - assert await dual_cache.async_get_cache(key=rpm_key) == 6 - assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 - assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] - assert any("both limit types are enforced" in record.message for record in caplog.records) - - @pytest.mark.asyncio - async def test_io_success_still_tracks_tpm_for_mixed_deployment(self): - """ - A deployment with itpm/otpm AND tpm/rpm must have BOTH counters updated on - success, otherwise the tpm_key the pre-call check reads is never written - and the tpm_limit can never be enforced. - """ - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - model_id = "io-tpm-mixed-id" - deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" - tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=5, ttl=60) - - kwargs = { - "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, - "standard_logging_object": { - "model_id": model_id, - "total_tokens": 7, - "hidden_params": {"litellm_model_name": deployment_name}, - }, - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), - ) - - await check.async_log_success_event(kwargs, response, None, None) - - # ITPM reconciled down from the 5-token reservation to actual usage (3). - assert await dual_cache.async_get_cache(key=itpm_key) == 3 - # TPM tracking must still run so the tpm/rpm pre-call path can enforce it. - assert await dual_cache.async_get_cache(key=tpm_key) == 7 - - def test_io_success_still_tracks_tpm_for_mixed_deployment_sync(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - model_id = "io-tpm-mixed-sync-id" - deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" - tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" - dual_cache.set_cache(key=itpm_key, value=5, ttl=60) - - kwargs = { - "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, - "standard_logging_object": { - "model_id": model_id, - "total_tokens": 7, - "hidden_params": {"litellm_model_name": deployment_name}, - }, - } - response = ModelResponse( - choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], - usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), - ) - - check.log_success_event(kwargs, response, None, None) - - assert dual_cache.get_cache(key=itpm_key) == 3 - assert dual_cache.get_cache(key=tpm_key) == 7 - - @pytest.mark.asyncio - async def test_failure_refunds_itpm_reservation(self): - from litellm.utils import get_utc_datetime - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-refund-id:bedrock_mantle/test:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) - - reservation = {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key} - kwargs = { - "standard_logging_object": { - "model_id": "io-refund-id", - "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, - "metadata": dict(reservation), - }, - "metadata": dict(reservation), - } - await check.async_log_failure_event(kwargs, None, None, None) - - current = await dual_cache.async_get_cache(key=itpm_key) - assert current == 0 - - -class TestRouterIOTokenIntegration: - @pytest.mark.asyncio - async def test_model_group_info_aggregates_io_limits(self): - router = Router( - model_list=[ - { - "model_name": "opus", - "litellm_params": { - "model": "bedrock_mantle/anthropic.claude-opus-4-7", - "itpm": 100, - "otpm": 20, - }, - } - ], - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - info = router.get_model_group_info("opus") - assert info is not None - assert info.itpm == 100 - assert info.otpm == 20 - - -class TestContextSlotRetention: - def test_setter_stores_kwargs_only_for_io_limited_deployments(self): - """ - The context slot pins the entire request kwargs (messages included) - for the lifetime of the surrounding asyncio context, and pooled - resources created mid-request (e.g. redis connections) capture that - context, extending the pin far past the request. Only ITPM/OTPM - pre-call checks read the slot, so the setter must store None for - deployments without io token limits and still clear reservation - sentinels from kwargs either way. - """ - kwargs = { - "messages": [{"role": "user", "content": "x" * 1000}], - "metadata": {ITPM_RESERVED_KEY: 999, ITPM_CACHE_KEY: "forged"}, - } - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) - assert get_io_token_rate_limit_request_kwargs() is None - assert ITPM_RESERVED_KEY not in kwargs["metadata"] - assert ITPM_CACHE_KEY not in kwargs["metadata"] - - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=True) - assert get_io_token_rate_limit_request_kwargs() is kwargs - - set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) - assert get_io_token_rate_limit_request_kwargs() is None - - @pytest.mark.asyncio - async def test_router_does_not_pin_kwargs_without_io_limits(self): - router = Router( - model_list=[ - { - "model_name": "plain", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, - } - ] - ) - set_io_token_rate_limit_request_kwargs(None) - kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - deployment = router.get_deployment_by_model_group_name("plain") - assert deployment is not None - router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) - assert get_io_token_rate_limit_request_kwargs() is None - - @pytest.mark.asyncio - async def test_router_pins_kwargs_for_io_limited_deployment(self): - router = Router( - model_list=[ - { - "model_name": "limited", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "itpm": 100}, - } - ], - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - set_io_token_rate_limit_request_kwargs(None) - kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - deployment = router.get_deployment_by_model_group_name("limited") - assert deployment is not None - router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) - assert get_io_token_rate_limit_request_kwargs() is kwargs - - -@pytest.mark.asyncio -async def test_the_deployment_itpm_reservation_counts_the_request_off_the_event_loop(): - from litellm.utils import get_utc_datetime - from tests.large_text import text - from tests.test_litellm.litellm_core_utils.event_loop_lag import ( - assert_loop_stayed_free, - timed_with_loop_lags, - warm_tokenizer, - ) - - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - warm_tokenizer("anthropic/claude-fable-5") - deployment = { - "litellm_params": {"model": "anthropic/claude-fable-5", "itpm": 10_000_000}, - "model_info": {"id": "io-loop-id"}, - "model_name": "claude", - } - set_io_token_rate_limit_request_kwargs({"messages": [{"role": "user", "content": text * 100}], "metadata": {}}) - - _, took, lags = await timed_with_loop_lags(lambda: check.async_pre_call_check(deployment)) - - minute = get_utc_datetime().strftime("%H-%M") - reserved = await dual_cache.async_get_cache(key=f"global_router:io-loop-id:anthropic/claude-fable-5:itpm:{minute}") - assert reserved > 100_000 - assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/types/llms/test_types_llms_bedrock.py b/tests/test_litellm/types/llms/test_types_llms_bedrock.py deleted file mode 100644 index a5ad882e775..00000000000 --- a/tests/test_litellm/types/llms/test_types_llms_bedrock.py +++ /dev/null @@ -1,46 +0,0 @@ -import pytest -from pydantic import ValidationError - -from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams - - -def test_model_validate_keeps_auth_params_and_ignores_request_params(): - auth_params = AwsAuthParams.model_validate( - { - "aws_role_name": "arn:aws:iam::999999999999:role/litellm-role", - "aws_session_name": "litellm-session", - "aws_external_id": "litellm-external-id", - "aws_region_name": "us-west-2", - "aws_bedrock_runtime_endpoint": "https://bedrock.example.com", - "model": "anthropic.claude-haiku-4-5-20251001-v1:0", - "temperature": 0.1, - "messages": [{"role": "user", "content": "hi"}], - } - ) - - assert auth_params.aws_role_name == "arn:aws:iam::999999999999:role/litellm-role" - assert auth_params.aws_session_name == "litellm-session" - assert auth_params.aws_external_id == "litellm-external-id" - assert auth_params.aws_access_key_id is None - assert set(auth_params.model_dump()) == set(AWS_AUTH_PARAM_KEYS) - assert not set(AWS_AUTH_PARAM_KEYS) & {"aws_region_name", "aws_bedrock_runtime_endpoint", "model", "temperature"} - - -@pytest.mark.parametrize( - ("field", "value"), - [ - ("aws_role_name", 1234), - ("aws_session_name", ["litellm-session"]), - ("aws_external_id", {"id": "x"}), - ], -) -def test_model_validate_rejects_non_string_credentials(field, value): - with pytest.raises(ValidationError): - AwsAuthParams.model_validate({field: value}) - - -def test_frozen_struct_rejects_field_assignment(): - auth_params = AwsAuthParams(aws_role_name="arn:aws:iam::999999999999:role/litellm-role") - - with pytest.raises(ValidationError): - auth_params.aws_role_name = "arn:aws:iam::999999999999:role/other-role" diff --git a/tests/test_litellm/types/llms/test_types_llms_openai.py b/tests/test_litellm/types/llms/test_types_llms_openai.py deleted file mode 100644 index 64ec09838e8..00000000000 --- a/tests/test_litellm/types/llms/test_types_llms_openai.py +++ /dev/null @@ -1,591 +0,0 @@ -import asyncio -from typing import Optional -from unittest.mock import AsyncMock, patch - -import pytest - -import json - -import litellm -from litellm.types.llms.openai import HttpxBinaryResponseContent - - -@pytest.mark.parametrize("stream", (False, True)) -def test_completion_response_reasoning_summary_round_trip(stream: bool) -> None: - from typing import Final - - from litellm.types.llms.openai import ( - ChatCompletionReasoningItem, - ChatCompletionReasoningSummaryTextBlock, - ) - from litellm.types.utils import ( - Choices, - Delta, - Message, - ModelResponse, - ModelResponseStream, - StreamingChoices, - ) - - reasoning_item: Final = ChatCompletionReasoningItem( - type="reasoning", - id="rs_123", - encrypted_content="encrypted", - summary=[ChatCompletionReasoningSummaryTextBlock(type="summary_text", text="Reasoning summary")], - ) - response: Final = ( - ModelResponseStream(choices=[StreamingChoices(delta=Delta(reasoning_items=[reasoning_item]))]) - if stream - else ModelResponse(choices=[Choices(message=Message(reasoning_items=[reasoning_item]))]) - ) - message_key: Final = "delta" if stream else "message" - assert response.model_dump()["choices"][0][message_key]["reasoning_items"] == [reasoning_item] - - restored: Final = type(response).model_validate_json(response.model_dump_json()) - assert restored.model_dump()["choices"][0][message_key]["reasoning_items"] == [reasoning_item] - - -def test_generic_event(): - from litellm.types.llms.openai import GenericEvent - - event = {"type": "test", "test": "test"} - event = GenericEvent(**event) - assert event.type == "test" - assert event.test == "test" - - -def test_output_item_added_event(): - from litellm.types.llms.openai import OutputItemAddedEvent - - event = { - "type": "response.output_item.added", - "sequence_number": 4, - "output_index": 1, - "item": None, - } - event = OutputItemAddedEvent(**event) - assert event.type == "response.output_item.added" - assert event.sequence_number == 4 - assert event.output_index == 1 - assert event.item is None - - -class TestResponsesAPIResponseOutputText: - """Tests for the output_text property on ResponsesAPIResponse""" - - def test_output_text_with_single_message(self): - """Test output_text with a single message containing text output""" - from litellm.types.llms.openai import ResponsesAPIResponse - - response = ResponsesAPIResponse( - id="resp_123", - created_at=1234567890, - output=[ - { - "type": "message", - "id": "msg_123", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": "Hello, world!", - } - ], - } - ], - ) - - assert response.output_text == "Hello, world!" - - def test_output_text_with_multiple_messages(self): - """Test output_text with multiple messages aggregates all text""" - from litellm.types.llms.openai import ResponsesAPIResponse - - response = ResponsesAPIResponse( - id="resp_123", - created_at=1234567890, - output=[ - { - "type": "message", - "id": "msg_1", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": "First part. ", - } - ], - }, - { - "type": "message", - "id": "msg_2", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": "Second part.", - } - ], - }, - ], - ) - - assert response.output_text == "First part. Second part." - - def test_output_text_with_no_text_content(self): - """Test output_text returns empty string when no output_text content exists""" - from litellm.types.llms.openai import ResponsesAPIResponse - - response = ResponsesAPIResponse( - id="resp_123", - created_at=1234567890, - output=[ - { - "type": "function_call", - "id": "call_123", - "status": "completed", - "name": "get_weather", - "arguments": "{}", - } - ], - ) - - assert response.output_text == "" - - def test_output_text_with_mixed_content(self): - """Test output_text only aggregates output_text type content""" - from litellm.types.llms.openai import ResponsesAPIResponse - - response = ResponsesAPIResponse( - id="resp_123", - created_at=1234567890, - output=[ - { - "type": "message", - "id": "msg_1", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": "The weather is sunny. ", - }, - { - "type": "refusal", - "refusal": "I cannot do that.", - }, - ], - }, - { - "type": "function_call", - "id": "call_123", - "status": "completed", - "name": "get_weather", - "arguments": "{}", - }, - ], - ) - - assert response.output_text == "The weather is sunny. " - - def test_output_text_with_empty_output(self): - """Test output_text returns empty string with empty output list""" - from litellm.types.llms.openai import ResponsesAPIResponse - - response = ResponsesAPIResponse( - id="resp_123", - created_at=1234567890, - output=[], - ) - - assert response.output_text == "" - - -class TestAssistantMessageImageUrlContent: - """ - Regression tests for image_url blocks in assistant message content. - - Bug: ChatCompletionAssistantMessage.content did not include - ChatCompletionImageObject in its union, so Pydantic v2 silently dropped - image_url blocks (content → []) when serialising via AllMessageValues. - This affects users who store conversation history as JSON (e.g. in a DB) - and read it back typed as list[AllMessageValues]. - """ - - ASSISTANT_MESSAGE_WITH_IMAGE = { - "role": "assistant", - "content": [ - {"type": "text", "text": "Here is the image you requested:"}, - { - "type": "image_url", - "image_url": { - "url": ( - "data:image/png;base64," - "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAA" - "DUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - ) - }, - }, - ], - } - - def test_assistant_message_image_url_preserved_single(self): - """ - TypeAdapter(ChatCompletionAssistantMessage): image_url block must survive - validate_python → dump_python without being dropped or raising an error. - """ - from typing import List - - from pydantic import TypeAdapter - - from litellm.types.llms.openai import ChatCompletionAssistantMessage - - adapter = TypeAdapter(ChatCompletionAssistantMessage) - validated = adapter.validate_python(self.ASSISTANT_MESSAGE_WITH_IMAGE) - dumped = adapter.dump_python(validated) - - raw_content = dumped.get("content") - # Pydantic may return a lazy SerializationIterator for Iterable fields; - # convert to list to consume it — this must not raise ValidationError. - content_blocks = list(raw_content) if raw_content is not None else [] - - assert ( - len(content_blocks) == 2 - ), f"Expected 2 content blocks (text + image_url), got {len(content_blocks)}: {content_blocks}" - types = [b.get("type") for b in content_blocks if isinstance(b, dict)] - assert ( - "image_url" in types - ), f"image_url block was silently dropped; blocks: {content_blocks}" - - def test_assistant_message_image_url_preserved_in_all_message_values(self): - """ - TypeAdapter(List[AllMessageValues]) DB round-trip: image_url blocks in an - assistant message must not be silently dropped during dump_python(mode='json'). - - This is the primary failing path: conversation history stored as JSON in a - database and read back typed as list[AllMessageValues]. - """ - from typing import List - - from pydantic import TypeAdapter - - from litellm.types.llms.openai import AllMessageValues - - conversation = [ - { - "role": "user", - "content": "Generate an image of a banana wearing a LiteLLM costume", - }, - self.ASSISTANT_MESSAGE_WITH_IMAGE, - ] - - adapter = TypeAdapter(List[AllMessageValues]) - validated = adapter.validate_python(conversation) - dumped = adapter.dump_python(validated, mode="json") - - assistant = next((m for m in dumped if m.get("role") == "assistant"), None) - assert assistant is not None, "Assistant message missing after serialisation" - - content = assistant.get("content", []) - assert isinstance( - content, list - ), f"content should be a list, got {type(content)}" - assert ( - len(content) == 2 - ), f"Expected 2 content blocks (text + image_url), got {len(content)}: {content}" - types = [b.get("type") for b in content if isinstance(b, dict)] - assert ( - "image_url" in types - ), f"image_url block was silently dropped during AllMessageValues serialisation; blocks: {content}" - - -class TestResponsesAPIReasoningNullFields: - """ - Tests for issue #16824: reasoning output items should not include null - status/content/encrypted_content fields. - - When a provider returns reasoning items without these fields, LiteLLM's - Pydantic parsing adds them as Optional defaults (None). Serializing them - as null breaks downstream SDKs (e.g., the OpenAI C# SDK crashes on - status=null). - - The fix uses a field_serializer on ResponsesAPIResponse.output that - mirrors the request-side filtering in - OpenAIResponsesAPIConfig._handle_reasoning_item(). - """ - - def _make_response(self, output): - from litellm.types.llms.openai import ResponsesAPIResponse - - return ResponsesAPIResponse( - id="resp_test", - created_at=1741476542, - model="gpt-5-mini", - object="response", - status="completed", - output=output, - ) - - def test_reasoning_item_null_fields_removed_model_dump(self): - """Null status/content/encrypted_content should be absent from model_dump.""" - response = self._make_response( - output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] - ) - dumped = response.model_dump() - reasoning = dumped["output"][0] - assert "status" not in reasoning - assert "content" not in reasoning - assert "encrypted_content" not in reasoning - - def test_reasoning_item_null_fields_removed_model_dump_json(self): - """Null fields should also be absent from model_dump_json.""" - response = self._make_response( - output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] - ) - parsed = json.loads(response.model_dump_json()) - reasoning = parsed["output"][0] - assert "status" not in reasoning - assert "content" not in reasoning - assert "encrypted_content" not in reasoning - - def test_reasoning_item_non_null_values_preserved(self): - """Non-null values on reasoning items should be kept.""" - response = self._make_response( - output=[ - { - "id": "rs_abc", - "type": "reasoning", - "summary": [], - "status": "completed", - "encrypted_content": "gAAAA...", - } - ] - ) - dumped = response.model_dump() - reasoning = dumped["output"][0] - assert reasoning["status"] == "completed" - assert reasoning["encrypted_content"] == "gAAAA..." - - def test_message_item_not_affected(self): - """Non-reasoning output items should keep all their fields.""" - response = self._make_response( - output=[ - { - "id": "msg_abc", - "type": "message", - "role": "assistant", - "status": "completed", - "content": [ - { - "type": "output_text", - "text": "Hello!", - "annotations": [], - } - ], - } - ] - ) - dumped = response.model_dump() - message = dumped["output"][0] - assert message["status"] == "completed" - assert message["type"] == "message" - assert len(message["content"]) == 1 - - def test_mixed_output_reasoning_and_message(self): - """Reasoning items cleaned, message items untouched in same response.""" - response = self._make_response( - output=[ - {"id": "rs_abc", "type": "reasoning", "summary": []}, - { - "id": "msg_abc", - "type": "message", - "role": "assistant", - "status": "completed", - "content": [ - { - "type": "output_text", - "text": "Answer", - "annotations": [], - } - ], - }, - ] - ) - dumped = response.model_dump() - reasoning = [ - o - for o in dumped["output"] - if isinstance(o, dict) and o.get("type") == "reasoning" - ][0] - message = [ - o - for o in dumped["output"] - if isinstance(o, dict) and o.get("type") == "message" - ][0] - assert "status" not in reasoning - assert "content" not in reasoning - assert message["status"] == "completed" - assert len(message["content"]) == 1 - - def test_reasoning_core_fields_preserved(self): - """id, type, summary should always be present on reasoning items.""" - response = self._make_response( - output=[{"id": "rs_abc", "type": "reasoning", "summary": ["thinking..."]}] - ) - dumped = response.model_dump() - reasoning = dumped["output"][0] - assert reasoning["id"] == "rs_abc" - assert reasoning["type"] == "reasoning" - assert reasoning["summary"] == ["thinking..."] - - def test_top_level_null_fields_unaffected(self): - """Top-level response fields with None should not be affected.""" - response = self._make_response( - output=[{"id": "rs_abc", "type": "reasoning", "summary": []}] - ) - dumped = response.model_dump() - assert "error" in dumped - assert dumped["error"] is None - assert "instructions" in dumped - assert dumped["instructions"] is None - - -def test_normalize_fine_tuning_job_dict_maps_azure_pending(): - from litellm.llms.openai.fine_tuning.handler import _normalize_fine_tuning_job_dict - - out = _normalize_fine_tuning_job_dict( - {"organization_id": None, "result_files": None, "status": "pending"}, - is_azure=True, - ) - assert out["organization_id"] == "" - assert out["result_files"] == [] - assert out["status"] == "queued" - - -def test_normalize_fine_tuning_job_dict_openai_unchanged(): - from litellm.llms.openai.fine_tuning.handler import _normalize_fine_tuning_job_dict - - data = {"organization_id": None, "result_files": None, "status": "pending"} - out = _normalize_fine_tuning_job_dict(data, is_azure=False) - assert out is data - - -def test_openai_file_object_accepts_pending_status(): - from litellm.types.llms.openai import OpenAIFileObject - - file_obj = OpenAIFileObject( - id="file-123", - bytes=1024, - created_at=1677610602, - filename="train.jsonl", - object="file", - purpose="fine-tune", - status="pending", - ) - assert file_obj.status == "pending" - - -class TestOpenAIFileObjectBatchGuardrailSerialization: - """The proxy-only `litellm_batch_guardrail` key must reach the wire only when something set it.""" - - @staticmethod - def _file_object(**overrides): - from litellm.types.llms.openai import OpenAIFileObject - - return OpenAIFileObject( - id="file-123", - object="file", - bytes=1024, - created_at=1677610602, - filename="batch.jsonl", - purpose="batch", - status="uploaded", - **overrides, - ) - - @staticmethod - def _report(): - from litellm.types.llms.openai import BatchGuardrailRecord, BatchGuardrailReport - - return BatchGuardrailReport( - submitted_records=3, - modified_records=(BatchGuardrailRecord(line=2, custom_id="dirty", action="redacted"),), - ) - - @pytest.mark.parametrize("mode", ["python", "json"]) - def test_key_absent_when_unset(self, mode): - assert "litellm_batch_guardrail" not in self._file_object().model_dump(mode=mode) - - @pytest.mark.parametrize("mode", ["python", "json"]) - def test_key_present_when_set(self, mode): - dumped = self._file_object(litellm_batch_guardrail=self._report()).model_dump(mode=mode) - assert dumped["litellm_batch_guardrail"]["submitted_records"] == 3 - - def test_nested_nulls_of_a_set_report_survive(self): - """`exclude_none=True` was rejected as the fix because it would strip these.""" - dumped = self._file_object(litellm_batch_guardrail=self._report()).model_dump(mode="json") - assert dumped["litellm_batch_guardrail"]["modified_records"] == [ - {"line": 2, "custom_id": "dirty", "action": "redacted", "guardrail": None} - ] - - def test_by_alias_dump_also_omits_the_key(self): - """Tripwire: the serializer filters a literal key name, which an added alias would bypass.""" - assert "litellm_batch_guardrail" not in self._file_object().model_dump(mode="json", by_alias=True) - - def test_other_optional_fields_still_serialize_as_null(self): - dumped = self._file_object().model_dump(mode="json") - assert dumped["expires_at"] is None - assert dumped["status_details"] is None - - def test_round_trip_of_a_set_report_is_lossless(self): - from litellm.types.llms.openai import OpenAIFileObject - - original = self._file_object(litellm_batch_guardrail=self._report()) - assert OpenAIFileObject(**original.model_dump()) == original - - def test_serialization_json_schema_still_describes_the_model(self): - """A return annotation on the wrap serializer would collapse this to a bare object.""" - from litellm.types.llms.openai import OpenAIFileObject - - schema = OpenAIFileObject.model_json_schema(mode="serialization") - assert "litellm_batch_guardrail" in schema["properties"] - - def test_key_omitted_inside_a_file_list_page(self): - from litellm.types.llms.openai import FileListPage - - page = FileListPage(object="list", data=[self._file_object()], has_more=False) - assert "litellm_batch_guardrail" not in page.model_dump(mode="json")["data"][0] - - -def _binary_content(payload: bytes) -> HttpxBinaryResponseContent: - import httpx - - return HttpxBinaryResponseContent(httpx.Response(200, content=payload)) - - -def test_httpx_binary_response_content_hidden_params_are_per_instance(): - first = _binary_content(b"first") - second = _binary_content(b"second") - - first._hidden_params["response_cost"] = 0.5 - - assert second._hidden_params == {} - - -def test_set_response_cost_none_leaves_hidden_params_empty(): - binary_response = _binary_content(b"audio") - - binary_response.set_response_cost(None) - - assert "response_cost" not in binary_response._hidden_params - - binary_response.set_response_cost(0.25) - - assert binary_response._hidden_params["response_cost"] == 0.25 - - binary_response.set_response_cost(None) - - assert "response_cost" not in binary_response._hidden_params diff --git a/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py b/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py deleted file mode 100644 index 2e5986d3ef8..00000000000 --- a/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py +++ /dev/null @@ -1,168 +0,0 @@ -""" -Tests for pipeline type definitions. -""" - -import pytest -from pydantic import ValidationError - -from litellm.types.proxy.policy_engine.pipeline_types import ( - GuardrailPipeline, - PipelineExecutionResult, - PipelineStep, - PipelineStepResult, -) -from litellm.types.proxy.policy_engine.policy_types import ( - Policy, - PolicyGuardrails, -) - - -def test_pipeline_step_defaults(): - step = PipelineStep(guardrail="my-guard") - assert step.on_fail == "block" - assert step.on_pass == "allow" - assert step.on_error is None - assert step.pass_data is False - assert step.modify_response_message is None - - -def test_pipeline_step_valid_actions(): - step = PipelineStep(guardrail="my-guard", on_fail="next", on_pass="next") - assert step.on_fail == "next" - assert step.on_pass == "next" - - -def test_pipeline_step_all_action_types(): - for action in ("allow", "block", "next", "modify_response"): - step = PipelineStep( - guardrail="g", on_fail=action, on_pass=action, on_error=action - ) - assert step.on_fail == action - assert step.on_pass == action - assert step.on_error == action - - -def test_pipeline_step_invalid_action_rejected(): - with pytest.raises(ValidationError): - PipelineStep(guardrail="my-guard", on_fail="invalid_action") - - -def test_pipeline_step_invalid_on_pass_rejected(): - with pytest.raises(ValidationError): - PipelineStep(guardrail="my-guard", on_pass="skip") - - -def test_pipeline_step_on_error_valid(): - step = PipelineStep( - guardrail="g", on_error="next", on_fail="block", on_pass="allow" - ) - assert step.on_error == "next" - - -def test_pipeline_step_invalid_on_error_rejected(): - with pytest.raises(ValidationError): - PipelineStep(guardrail="my-guard", on_error="invalid") - - -def test_pipeline_requires_at_least_one_step(): - with pytest.raises(ValidationError): - GuardrailPipeline(mode="pre_call", steps=[]) - - -def test_pipeline_invalid_mode_rejected(): - with pytest.raises(ValidationError): - GuardrailPipeline( - mode="during_call", - steps=[PipelineStep(guardrail="g")], - ) - - -def test_pipeline_valid_modes(): - for mode in ("pre_call", "post_call"): - pipeline = GuardrailPipeline( - mode=mode, - steps=[PipelineStep(guardrail="g")], - ) - assert pipeline.mode == mode - - -def test_pipeline_with_multiple_steps(): - pipeline = GuardrailPipeline( - mode="pre_call", - steps=[ - PipelineStep(guardrail="g1", on_fail="next", on_pass="allow"), - PipelineStep(guardrail="g2", on_fail="block", on_pass="allow"), - ], - ) - assert len(pipeline.steps) == 2 - assert pipeline.steps[0].guardrail == "g1" - assert pipeline.steps[1].guardrail == "g2" - - -def test_policy_with_pipeline_parses(): - policy = Policy( - guardrails=PolicyGuardrails(add=["g1", "g2"]), - pipeline=GuardrailPipeline( - mode="pre_call", - steps=[ - PipelineStep(guardrail="g1", on_fail="next"), - PipelineStep(guardrail="g2"), - ], - ), - ) - assert policy.pipeline is not None - assert len(policy.pipeline.steps) == 2 - - -def test_policy_without_pipeline(): - policy = Policy( - guardrails=PolicyGuardrails(add=["g1"]), - ) - assert policy.pipeline is None - - -def test_pipeline_step_result(): - result = PipelineStepResult( - guardrail_name="g1", - outcome="fail", - action_taken="next", - error_detail="Content policy violation", - duration_seconds=0.05, - ) - assert result.outcome == "fail" - assert result.action_taken == "next" - - -def test_pipeline_execution_result(): - result = PipelineExecutionResult( - terminal_action="block", - step_results=[ - PipelineStepResult( - guardrail_name="g1", - outcome="fail", - action_taken="next", - ), - PipelineStepResult( - guardrail_name="g2", - outcome="fail", - action_taken="block", - ), - ], - error_message="Content blocked", - ) - assert result.terminal_action == "block" - assert len(result.step_results) == 2 - - -def test_pipeline_step_extra_fields_rejected(): - with pytest.raises(ValidationError): - PipelineStep(guardrail="g", unknown_field="value") - - -def test_pipeline_extra_fields_rejected(): - with pytest.raises(ValidationError): - GuardrailPipeline( - mode="pre_call", - steps=[PipelineStep(guardrail="g")], - unknown="value", - ) diff --git a/tests/test_litellm/types/proxy/policy_engine/test_policy_types.py b/tests/test_litellm/types/proxy/policy_engine/test_policy_types.py deleted file mode 100644 index bcd6d39aa4d..00000000000 --- a/tests/test_litellm/types/proxy/policy_engine/test_policy_types.py +++ /dev/null @@ -1,15 +0,0 @@ -import pytest -from pydantic import ValidationError - -from litellm.types.proxy.policy_engine.policy_types import PolicyAttachment - - -@pytest.mark.parametrize("priority", [-2147483648, 2147483647]) -def test_policy_attachment_accepts_int32_priority(priority: int): - assert PolicyAttachment(policy="p", priority=priority).priority == priority - - -@pytest.mark.parametrize("priority", [-2147483649, 2147483648]) -def test_policy_attachment_rejects_priority_outside_int32(priority: int): - with pytest.raises(ValidationError): - PolicyAttachment(policy="p", priority=priority) diff --git a/tests/test_litellm/types/proxy/policy_engine/test_resolver_types.py b/tests/test_litellm/types/proxy/policy_engine/test_resolver_types.py deleted file mode 100644 index f31b9d7e873..00000000000 --- a/tests/test_litellm/types/proxy/policy_engine/test_resolver_types.py +++ /dev/null @@ -1,115 +0,0 @@ -""" -Tests for pipeline field on policy CRUD types (resolver_types.py). -""" - -import pytest -from pydantic import ValidationError - -from litellm.types.proxy.policy_engine.resolver_types import ( - PolicyAttachmentCreateRequest, - PolicyCreateRequest, - PolicyDBResponse, - PolicyUpdateRequest, -) - - -def test_policy_create_request_with_pipeline(): - pipeline_data = { - "mode": "pre_call", - "steps": [ - {"guardrail": "g1", "on_fail": "next", "on_pass": "allow"}, - {"guardrail": "g2", "on_fail": "block", "on_pass": "allow"}, - ], - } - req = PolicyCreateRequest( - policy_name="test-policy", - guardrails_add=["g1", "g2"], - pipeline=pipeline_data, - ) - assert req.pipeline is not None - assert req.pipeline["mode"] == "pre_call" - assert len(req.pipeline["steps"]) == 2 - - -def test_policy_create_request_without_pipeline(): - req = PolicyCreateRequest( - policy_name="test-policy", - guardrails_add=["g1"], - ) - assert req.pipeline is None - - -def test_policy_update_request_with_pipeline(): - pipeline_data = { - "mode": "pre_call", - "steps": [ - {"guardrail": "g1", "on_fail": "block", "on_pass": "allow"}, - ], - } - req = PolicyUpdateRequest(pipeline=pipeline_data) - assert req.pipeline is not None - assert req.pipeline["steps"][0]["guardrail"] == "g1" - - -def test_policy_db_response_with_pipeline(): - pipeline_data = { - "mode": "pre_call", - "steps": [ - {"guardrail": "g1", "on_fail": "next", "on_pass": "allow"}, - {"guardrail": "g2", "on_fail": "block", "on_pass": "allow"}, - ], - } - resp = PolicyDBResponse( - policy_id="test-id", - policy_name="test-policy", - guardrails_add=["g1", "g2"], - pipeline=pipeline_data, - ) - assert resp.pipeline is not None - assert resp.pipeline["mode"] == "pre_call" - dumped = resp.model_dump() - assert dumped["pipeline"]["steps"][0]["guardrail"] == "g1" - - -def test_policy_db_response_without_pipeline(): - resp = PolicyDBResponse( - policy_id="test-id", - policy_name="test-policy", - ) - assert resp.pipeline is None - dumped = resp.model_dump() - assert dumped["pipeline"] is None - - -def test_policy_create_request_roundtrip(): - pipeline_data = { - "mode": "post_call", - "steps": [ - { - "guardrail": "g1", - "on_fail": "modify_response", - "on_pass": "next", - "pass_data": True, - "modify_response_message": "custom msg", - }, - ], - } - req = PolicyCreateRequest( - policy_name="roundtrip-test", - guardrails_add=["g1"], - pipeline=pipeline_data, - ) - dumped = req.model_dump() - restored = PolicyCreateRequest(**dumped) - assert restored.pipeline == pipeline_data - - -@pytest.mark.parametrize("priority", [-2147483648, 2147483647]) -def test_policy_attachment_create_request_accepts_int32_priority(priority: int): - assert PolicyAttachmentCreateRequest(policy_name="p", priority=priority).priority == priority - - -@pytest.mark.parametrize("priority", [-2147483649, 2147483648]) -def test_policy_attachment_create_request_rejects_priority_outside_int32(priority: int): - with pytest.raises(ValidationError): - PolicyAttachmentCreateRequest(policy_name="p", priority=priority) diff --git a/tests/test_litellm/videos/test_main.py b/tests/test_litellm/videos/test_main.py deleted file mode 100644 index 22e1e5c05eb..00000000000 --- a/tests/test_litellm/videos/test_main.py +++ /dev/null @@ -1,455 +0,0 @@ -""" -Dispatch-contract tests for litellm/videos/main.py - -Each public video operation is a pair: a sync `video_*` worker (decorated with -@client) that resolves the provider, fetches the provider config, logs, and then -forwards to exactly one `base_llm_http_handler.video_*_handler`; and an async -`avideo_*` wrapper that delegates to the sync worker in an executor. - -This file locks the contract of that layer so a regression fails loudly: - - 1. DISPATCH - the one correct handler fired and every sibling video handler - asserted NOT called. A copy-paste that calls the wrong handler - (e.g. remix -> edit) flips this. - 2. RESULT - the handler's return value is propagated by identity. - 3. PROVIDER - custom_llm_provider is decoded from an encoded video id when not - passed (status/content/remix/edit/extension), or defaults to - "openai" (list/create_character/get_character). This is the exact - surface of the historical "content defaulted to openai" bug. - 4. PAYLOAD - the provider config object and the operation's identifying args - (video_id/prompt/name/...) reach the handler; _is_async is False - on the sync path. - 5. SHORT-CIRCUIT - mock_response returns a typed object without any handler call. - 6. UNSUPPORTED - a None provider config raises before any handler fires. - 7. DELEGATION - avideo_* returns the sync worker's result untouched, sets - async_call=True, and pre-resolves the provider where it must. - -Seams mocked: the http handler (network), the provider-config registry lookup, -get_llm_provider, and the video-generation optional-param builders. The id decode -helper runs for real against genuinely-encoded ids, so the provider assertions -reflect production. -""" - -from contextlib import ExitStack -from dataclasses import dataclass -from typing import Any, Dict -from unittest.mock import MagicMock, patch - -import pytest - - -import litellm -from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.types.videos.main import CharacterObject, VideoObject -from litellm.types.videos.utils import encode_video_id_with_provider -from litellm.videos import main as videos_main - -# A real model-encoded video id: decodes (for real) to provider "azure". Used to -# prove the sync workers derive custom_llm_provider from the id, not a hardcode. -AZURE_VIDEO_ID = encode_video_id_with_provider("video_raw", "azure", "deployment-1") - -# The nine sync handlers on base_llm_http_handler. Dispatch tests assert exactly -# one fired and the other eight did not. -SYNC_HANDLERS = ( - "video_generation_handler", - "video_content_handler", - "video_remix_handler", - "video_create_character_handler", - "video_get_character_handler", - "video_edit_handler", - "video_extension_handler", - "video_list_handler", - "video_status_handler", -) - -GEN_OPTIONAL_PARAMS = {"seconds": "8", "size": "720x1280"} - - -@dataclass -class Seams: - handler: MagicMock - get_config: MagicMock - config: MagicMock - - def kwargs_of(self, handler_name: str) -> Dict[str, Any]: - method = getattr(self.handler, handler_name) - assert method.call_count == 1 - return dict(method.call_args.kwargs) - - def assert_only(self, handler_name: str) -> None: - for name in SYNC_HANDLERS: - method = getattr(self.handler, name) - if name == handler_name: - method.assert_called_once() - else: - method.assert_not_called() - - -@pytest.fixture -def seams(): - handler = MagicMock(spec=BaseLLMHTTPHandler) - config = MagicMock(name="provider_video_config") - get_config = MagicMock(return_value=config) - - with ExitStack() as stack: - stack.enter_context(patch.object(videos_main, "base_llm_http_handler", handler)) - stack.enter_context( - patch.object( - videos_main.ProviderConfigManager, - "get_provider_video_config", - get_config, - ) - ) - # video_generation resolves model+provider through get_llm_provider and - # builds optional params; mock those so the dispatch payload is deterministic. - stack.enter_context( - patch.object( - videos_main, - "get_llm_provider", - MagicMock(return_value=("sora-2", "openai", None, None)), - ) - ) - stack.enter_context( - patch.object( - videos_main.VideoGenerationRequestUtils, - "get_requested_video_generation_optional_param", - MagicMock(return_value={"seconds": "8"}), - ) - ) - stack.enter_context( - patch.object( - videos_main.VideoGenerationRequestUtils, - "get_optional_params_video_generation", - MagicMock(return_value=dict(GEN_OPTIONAL_PARAMS)), - ) - ) - yield Seams(handler=handler, get_config=get_config, config=config) - - -# =========================================================================== # -# Dispatch contract - one rich test per sync worker. -# =========================================================================== # - - -def test_video_generation__dispatch(seams): - result = videos_main.video_generation(prompt="a sunset", model="sora-2") - - seams.assert_only("video_generation_handler") - assert result is seams.handler.video_generation_handler.return_value - kw = seams.kwargs_of("video_generation_handler") - assert kw["model"] == "sora-2" - assert kw["prompt"] == "a sunset" - assert kw["custom_llm_provider"] == "openai" - assert kw["video_generation_provider_config"] is seams.config - assert kw["video_generation_optional_request_params"] == GEN_OPTIONAL_PARAMS - assert kw["_is_async"] is False - - -def test_video_status__dispatch_and_provider_from_id(seams): - result = videos_main.video_status(video_id=AZURE_VIDEO_ID) - - seams.assert_only("video_status_handler") - assert result is seams.handler.video_status_handler.return_value - kw = seams.kwargs_of("video_status_handler") - assert kw["video_id"] == AZURE_VIDEO_ID - assert kw["custom_llm_provider"] == "azure" # decoded from the id, not openai - assert kw["video_status_provider_config"] is seams.config - assert kw["_is_async"] is False - # provider config requested for the decoded provider, not a hardcode. - assert seams.get_config.call_args.kwargs["provider"] == litellm.LlmProviders.AZURE - - -def test_video_content__dispatch_and_provider_from_id(seams): - result = videos_main.video_content(video_id=AZURE_VIDEO_ID, variant="thumbnail") - - seams.assert_only("video_content_handler") - assert result is seams.handler.video_content_handler.return_value - kw = seams.kwargs_of("video_content_handler") - assert kw["video_id"] == AZURE_VIDEO_ID - assert kw["custom_llm_provider"] == "azure" - assert kw["variant"] == "thumbnail" - assert kw["video_content_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_content__plain_id_defaults_to_openai(seams): - videos_main.video_content(video_id="video_plain") - - assert seams.kwargs_of("video_content_handler")["custom_llm_provider"] == "openai" - - -def test_video_remix__dispatch_and_provider_from_id(seams): - result = videos_main.video_remix(video_id=AZURE_VIDEO_ID, prompt="new colors") - - seams.assert_only("video_remix_handler") - assert result is seams.handler.video_remix_handler.return_value - kw = seams.kwargs_of("video_remix_handler") - assert kw["video_id"] == AZURE_VIDEO_ID - assert kw["prompt"] == "new colors" - assert kw["custom_llm_provider"] == "azure" - assert kw["video_remix_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_edit__dispatch_and_provider_from_id(seams): - result = videos_main.video_edit(video_id=AZURE_VIDEO_ID, prompt="brighter") - - seams.assert_only("video_edit_handler") - assert result is seams.handler.video_edit_handler.return_value - kw = seams.kwargs_of("video_edit_handler") - assert kw["video_id"] == AZURE_VIDEO_ID - assert kw["prompt"] == "brighter" - assert kw["custom_llm_provider"] == "azure" - assert kw["video_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_extension__dispatch_and_provider_from_id(seams): - result = videos_main.video_extension( - video_id=AZURE_VIDEO_ID, prompt="continue", seconds="5" - ) - - seams.assert_only("video_extension_handler") - assert result is seams.handler.video_extension_handler.return_value - kw = seams.kwargs_of("video_extension_handler") - assert kw["video_id"] == AZURE_VIDEO_ID - assert kw["prompt"] == "continue" - assert kw["seconds"] == "5" - assert kw["custom_llm_provider"] == "azure" - assert kw["video_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_list__dispatch_defaults_to_openai(seams): - result = videos_main.video_list(after="cur", limit=5, order="desc") - - seams.assert_only("video_list_handler") - assert result is seams.handler.video_list_handler.return_value - kw = seams.kwargs_of("video_list_handler") - assert kw["after"] == "cur" - assert kw["limit"] == 5 - assert kw["order"] == "desc" - assert kw["custom_llm_provider"] == "openai" - assert kw["video_list_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_create_character__dispatch_defaults_to_openai(seams): - video = MagicMock(name="video_upload") - result = videos_main.video_create_character(name="hero", video=video) - - seams.assert_only("video_create_character_handler") - assert result is seams.handler.video_create_character_handler.return_value - kw = seams.kwargs_of("video_create_character_handler") - assert kw["name"] == "hero" - assert kw["video"] is video - assert kw["custom_llm_provider"] == "openai" - assert kw["video_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_video_get_character__dispatch_defaults_to_openai(seams): - result = videos_main.video_get_character(character_id="char_1") - - seams.assert_only("video_get_character_handler") - assert result is seams.handler.video_get_character_handler.return_value - kw = seams.kwargs_of("video_get_character_handler") - assert kw["character_id"] == "char_1" - assert kw["custom_llm_provider"] == "openai" - assert kw["video_provider_config"] is seams.config - assert kw["_is_async"] is False - - -def test_explicit_provider_beats_decoded_id(seams): - """An explicit custom_llm_provider wins over the one encoded in the id.""" - videos_main.video_status(video_id=AZURE_VIDEO_ID, custom_llm_provider="vertex_ai") - - assert seams.kwargs_of("video_status_handler")["custom_llm_provider"] == "vertex_ai" - - -# =========================================================================== # -# mock_response short-circuit - returns a typed object, no handler call. -# =========================================================================== # - - -def test_generation__mock_response_short_circuits(seams): - resp = videos_main.video_generation( - prompt="x", - model="sora-2", - mock_response={"id": "v1", "object": "video", "status": "queued"}, - ) - - assert isinstance(resp, VideoObject) - assert resp.id == "v1" - seams.handler.video_generation_handler.assert_not_called() - - -def test_list__mock_response_short_circuits(seams): - resp = videos_main.video_list( - mock_response=[{"id": "v1", "object": "video", "status": "completed"}] - ) - - assert isinstance(resp, list) - assert resp[0].id == "v1" - seams.handler.video_list_handler.assert_not_called() - - -def test_get_character__mock_response_short_circuits(seams): - resp = videos_main.video_get_character( - character_id="char_1", - mock_response={ - "id": "char_1", - "object": "character", - "created_at": 1, - "name": "hero", - }, - ) - - assert isinstance(resp, CharacterObject) - assert resp.id == "char_1" - seams.handler.video_get_character_handler.assert_not_called() - - -# =========================================================================== # -# Unsupported provider - a None provider config raises before any dispatch. -# =========================================================================== # - - -def test_unsupported_provider_raises_without_dispatch(seams): - seams.get_config.return_value = None - - with pytest.raises(litellm.APIConnectionError): - videos_main.video_status(video_id=AZURE_VIDEO_ID) - - seams.handler.video_status_handler.assert_not_called() - - -# =========================================================================== # -# Async-wrapper delegation - representative coverage. -# =========================================================================== # - - -@pytest.mark.asyncio -async def test_avideo_generation__delegates_with_async_flag(): - sentinel = VideoObject(id="v-async", object="video", status="queued") - with ( - patch.object( - videos_main, "video_generation", MagicMock(return_value=sentinel) - ) as sync, - patch.object( - litellm, - "get_llm_provider", - MagicMock(return_value=("sora-2", "openai", None, None)), - ), - ): - result = await videos_main.avideo_generation(prompt="x", model="sora-2") - - assert result is sentinel - assert sync.call_args.kwargs["async_call"] is True - assert sync.call_args.kwargs["custom_llm_provider"] == "openai" - - -@pytest.mark.asyncio -async def test_avideo_status__delegates_untouched(): - sentinel = VideoObject(id="v-async", object="video", status="queued") - with patch.object( - videos_main, "video_status", MagicMock(return_value=sentinel) - ) as sync: - result = await videos_main.avideo_status(video_id="video_plain") - - assert result is sentinel - assert sync.call_args.kwargs["async_call"] is True - assert sync.call_args.kwargs["video_id"] == "video_plain" - - -@pytest.mark.asyncio -async def test_avideo_content__pre_decodes_provider_before_delegating(): - """avideo_content resolves the provider from the encoded id itself before - handing off, so the sync worker receives the decoded provider, not None.""" - sentinel = b"mp4-bytes" - with patch.object( - videos_main, "video_content", MagicMock(return_value=sentinel) - ) as sync: - result = await videos_main.avideo_content(video_id=AZURE_VIDEO_ID) - - assert result is sentinel - assert sync.call_args.kwargs["async_call"] is True - assert sync.call_args.kwargs["custom_llm_provider"] == "azure" - - -# =========================================================================== # -# Credential passthrough - DB/YAML model-config credentials the router injects -# via kwargs must reach the provider call for EVERY video handler, carried in -# litellm_params. Distinct per-field values catch a cross-wired field. -# =========================================================================== # - -DB_YAML_CREDS = { - "api_key": "sk-db-credential", - "api_base": "https://db-resource.test", - "api_version": "2024-12-31", - "vertex_project": "db-project-xyz", -} - -CREDENTIAL_OPERATIONS = [ - ( - "video_generation_handler", - lambda: videos_main.video_generation( - prompt="p", model="sora-2", **DB_YAML_CREDS - ), - ), - ( - "video_status_handler", - lambda: videos_main.video_status(video_id=AZURE_VIDEO_ID, **DB_YAML_CREDS), - ), - ( - "video_content_handler", - lambda: videos_main.video_content(video_id=AZURE_VIDEO_ID, **DB_YAML_CREDS), - ), - ( - "video_remix_handler", - lambda: videos_main.video_remix( - video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS - ), - ), - ( - "video_edit_handler", - lambda: videos_main.video_edit( - video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS - ), - ), - ( - "video_extension_handler", - lambda: videos_main.video_extension( - video_id=AZURE_VIDEO_ID, prompt="p", seconds="5", **DB_YAML_CREDS - ), - ), - ( - "video_list_handler", - lambda: videos_main.video_list(**DB_YAML_CREDS), - ), - ( - "video_create_character_handler", - lambda: videos_main.video_create_character( - name="hero", video=MagicMock(name="vid"), **DB_YAML_CREDS - ), - ), - ( - "video_get_character_handler", - lambda: videos_main.video_get_character(character_id="char_1", **DB_YAML_CREDS), - ), -] - - -@pytest.mark.parametrize( - "handler_name,invoke", - CREDENTIAL_OPERATIONS, - ids=[op[0] for op in CREDENTIAL_OPERATIONS], -) -def test_db_yaml_credentials_reach_every_handler(seams, handler_name, invoke): - invoke() - - litellm_params = seams.kwargs_of(handler_name)["litellm_params"] - assert litellm_params.get("api_key") == DB_YAML_CREDS["api_key"] - assert litellm_params.get("api_base") == DB_YAML_CREDS["api_base"] - assert litellm_params.get("api_version") == DB_YAML_CREDS["api_version"] - assert litellm_params.get("vertex_project") == DB_YAML_CREDS["vertex_project"] diff --git a/tests/test_litellm/videos/test_utils.py b/tests/test_litellm/videos/test_utils.py deleted file mode 100644 index 57fb549c23d..00000000000 --- a/tests/test_litellm/videos/test_utils.py +++ /dev/null @@ -1,193 +0,0 @@ -""" -Pure-logic contract tests for litellm/videos/main.py's request utils -(litellm/videos/utils.py: VideoGenerationRequestUtils). - -These lock the exact param-shaping behavior so a mutation that drops a filter, -flips a precedence, or stops removing a key fails loudly. The only seam is the -provider config's map_openai_params (a provider boundary); filter_out_litellm_params -runs for real, so the "litellm-internal params get stripped" assertions reflect -production. Every test asserts the exact resulting dict, never "ran without error". -""" - -from unittest.mock import MagicMock - - - -import litellm -from litellm.videos.utils import VideoGenerationRequestUtils - -get_requested = ( - VideoGenerationRequestUtils.get_requested_video_generation_optional_param -) -get_optional = VideoGenerationRequestUtils.get_optional_params_video_generation - - -# =========================================================================== # -# get_requested_video_generation_optional_param -# -# Receives the caller's full local_vars; must return only the API-bound optional -# params. filter_out_litellm_params strips known internal keys for real; the -# values used below were chosen against the live set: seconds/size/user/foo_param/ -# vertex_project/extra/a/b survive, api_key/metadata/litellm_* are stripped. -# =========================================================================== # - - -def test_requested__drops_none_and_excluded_keys(): - result = get_requested( - { - "seconds": "8", - "size": None, # None -> dropped - "prompt": "a sunset", # excluded - "model": "sora-2", # excluded - "user": "u1", - } - ) - assert result == {"seconds": "8", "user": "u1"} - - -def test_requested__strips_litellm_internal_params(): - result = get_requested( - { - "seconds": "8", - "api_key": "sk-secret", - "metadata": {"x": 1}, - "litellm_call_id": "id-123", - } - ) - assert result == {"seconds": "8"} - - -def test_requested__timeout_always_removed(): - # timeout is NOT a litellm-internal param, so only the explicit pop removes it. - result = get_requested({"seconds": "8", "timeout": 30}) - assert result == {"seconds": "8"} - - -def test_requested__nested_kwargs_merge_and_override_base(): - result = get_requested( - {"seconds": "8", "kwargs": {"size": "720x1280", "seconds": "override"}} - ) - # nested kwargs win over the top-level base params on collision. - assert result == {"seconds": "override", "size": "720x1280"} - - -def test_requested__non_dict_kwargs_treated_as_empty(): - result = get_requested({"seconds": "8", "kwargs": "not-a-dict"}) - assert result == {"seconds": "8"} - - -def test_requested__none_input_returns_empty(): - assert get_requested(None) == {} - - -def test_requested__top_level_extra_body_spread_and_preserved(): - result = get_requested( - {"seconds": "8", "extra_body": {"vertex_project": "proj", "foo_param": "bar"}} - ) - # extra_body keys are both spread at top level AND kept under "extra_body". - assert result == { - "seconds": "8", - "vertex_project": "proj", - "foo_param": "bar", - "extra_body": {"vertex_project": "proj", "foo_param": "bar"}, - } - - -def test_requested__extra_body_kwargs_overrides_top_level(): - result = get_requested( - { - "extra_body": {"a": "top", "b": "top_b"}, - "kwargs": {"extra_body": {"a": "kw"}}, - } - ) - # kwargs' extra_body wins over the top-level extra_body on collision; the - # non-colliding top-level key survives. - assert result == { - "a": "kw", - "b": "top_b", - "extra_body": {"a": "kw", "b": "top_b"}, - } - - -def test_requested__extra_body_strips_litellm_internal_params(): - result = get_requested({"extra_body": {"api_key": "sk", "foo_param": "bar"}}) - # api_key filtered out of extra_body; only foo_param remains (and is spread). - assert result == {"foo_param": "bar", "extra_body": {"foo_param": "bar"}} - - -def test_requested__empty_extra_body_not_added(): - result = get_requested({"seconds": "8", "extra_body": {}}) - assert result == {"seconds": "8"} - assert "extra_body" not in result - - -# =========================================================================== # -# get_optional_params_video_generation -# -# Delegates mapping to the provider config (the seam) then folds extra_body in. -# =========================================================================== # - - -def _config(map_return): - config = MagicMock() - config.map_openai_params.return_value = map_return - return config - - -def test_optional__delegates_to_map_openai_params_with_drop_params(): - config = _config({"seconds": "8"}) - optional_params = {"seconds": "8"} - - result = get_optional( - model="sora-2", - video_generation_provider_config=config, - video_generation_optional_params=optional_params, - ) - - assert result == {"seconds": "8"} - config.map_openai_params.assert_called_once_with( - video_create_optional_params=optional_params, - model="sora-2", - drop_params=litellm.drop_params, - ) - - -def test_optional__extra_body_overrides_mapped_and_is_removed(): - # mapped output carries a leftover extra_body that must be popped; the input - # extra_body overrides a colliding mapped key and is spread in. - config = _config({"seconds": "8", "size": "mapped", "extra_body": {"leftover": 1}}) - - result = get_optional( - model="sora-2", - video_generation_provider_config=config, - video_generation_optional_params={ - "extra_body": {"size": "override", "extra": "x"} - }, - ) - - assert result == {"seconds": "8", "size": "override", "extra": "x"} - assert "extra_body" not in result - - -def test_optional__no_extra_body_returns_mapped_unchanged(): - config = _config({"seconds": "8"}) - - result = get_optional( - model="sora-2", - video_generation_provider_config=config, - video_generation_optional_params={"seconds": "8"}, - ) - - assert result == {"seconds": "8"} - - -def test_optional__non_dict_extra_body_ignored(): - config = _config({"seconds": "8"}) - - result = get_optional( - model="sora-2", - video_generation_provider_config=config, - video_generation_optional_params={"seconds": "8", "extra_body": None}, - ) - - assert result == {"seconds": "8"} From 72abd11b4ce43fb3762b008d97e60c95bc76e313 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:01:32 +0000 Subject: [PATCH 28/47] test: assert the responses bridge forwards aws_region_name Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...t_responses_bridge_provider_propagation.py | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) create mode 100644 tests/unit/completion_extras/test_responses_bridge_provider_propagation.py diff --git a/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py b/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py new file mode 100644 index 00000000000..09ef1889818 --- /dev/null +++ b/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py @@ -0,0 +1,59 @@ +from datetime import datetime +from unittest.mock import patch + +import pytest + +from litellm.completion_extras.litellm_responses_transformation.handler import ( + ResponsesToCompletionBridgeHandler, +) +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.types.utils import ModelResponse + +MODEL = "openai.gpt-5.5" +REGION = "us-east-2" + + +def _bedrock_mantle_kwargs() -> dict: + messages = [{"role": "user", "content": "hi"}] + logging_obj = LiteLLMLogging( + litellm_call_id="test-call", + call_type="acompletion", + model=MODEL, + messages=messages, + function_id="fn-id", + stream=False, + start_time=datetime.now(), + ) + return { + "model": MODEL, + "custom_llm_provider": "bedrock_mantle", + "messages": messages, + "optional_params": {}, + "litellm_params": { + "aws_region_name": REGION, + "api_base": "https://bedrock-mantle.us-east-1.api.aws/v1", + "custom_llm_provider": "bedrock_mantle", + }, + "headers": {}, + "model_response": ModelResponse(), + "logging_obj": logging_obj, + } + + +@pytest.mark.asyncio +async def test_acompletion_forwards_aws_region_name_to_aresponses(): + bridge = ResponsesToCompletionBridgeHandler() + cached = ModelResponse(id="chatcmpl-cached", model=MODEL) + + async def _fake_aresponses(**kwargs): + _fake_aresponses.kwargs = kwargs + return cached + + _fake_aresponses.kwargs = {} + + with patch("litellm.aresponses", _fake_aresponses): + result = await bridge.acompletion(**_bedrock_mantle_kwargs()) + + assert result is cached + assert _fake_aresponses.kwargs["aws_region_name"] == REGION + assert _fake_aresponses.kwargs["custom_llm_provider"] == "bedrock_mantle" From a69ca90ea5380bde22ad36ac9276604bae28652f Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:01:40 +0000 Subject: [PATCH 29/47] test: migrate phase 14 wave 2 provider tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...sonx_audio_transcription_transformation.py | 337 ---------- .../test_litellm/llms/watsonx/test_watsonx.py | 577 ------------------ .../llms/xai/xai_responses/__init__.py | 1 - .../xai/xai_responses/test_transformation.py | 105 ---- .../count_tokens/__init__.py | 0 .../test_count_tokens_location.py | 2 + .../test_count_tokens_no_vertexai_sdk.py | 0 ...ai_partner_models_llama3_transformation.py | 0 ...i_partner_models_mistral_transformation.py | 15 + .../vertex_ai/vertex_gemma_models/__init__.py | 0 .../test_vertex_gemma_transformation.py | 0 tests/unit/llms/vertex_ai/videos/__init__.py | 0 .../test_vertex_video_transformation.py | 0 ...est_volcengine_responses_transformation.py | 24 - tests/unit/llms/voyage/rerank/__init__.py | 0 .../test_voyage_rerank_transformation.py | 0 .../test_voyage_contextual_embedding.py | 0 .../test_voyage_multimodal_embedding.py | 0 tests/unit/llms/watsonx/__init__.py | 0 .../watsonx/audio_transcription/__init__.py | 0 ...sonx_audio_transcription_transformation.py | 85 +++ .../test_watsonx_embedding_transformation.py | 0 ...test_watsonx_passthrough_transformation.py | 0 tests/unit/llms/watsonx/rerank/__init__.py | 0 .../watsonx/rerank/test_watsonx_rerank.py | 0 tests/unit/llms/watsonx/test_watsonx.py | 74 +++ .../llms/watsonx/test_watsonx_common_utils.py | 0 .../test_xai_responses_transformation.py | 0 tests/unit/llms/you_com/__init__.py | 0 .../llms/you_com/test_you_com_search.py | 0 .../llms/zai/test_zai_provider.py | 0 31 files changed, 176 insertions(+), 1044 deletions(-) delete mode 100644 tests/test_litellm/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py delete mode 100644 tests/test_litellm/llms/watsonx/test_watsonx.py delete mode 100644 tests/test_litellm/llms/xai/xai_responses/__init__.py delete mode 100644 tests/test_litellm/llms/xai/xai_responses/test_transformation.py create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py (98%) rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py (100%) rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py (60%) create mode 100644 tests/unit/llms/vertex_ai/vertex_gemma_models/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py (100%) create mode 100644 tests/unit/llms/vertex_ai/videos/__init__.py rename tests/{test_litellm => unit}/llms/vertex_ai/videos/test_vertex_video_transformation.py (100%) rename tests/{test_litellm => unit}/llms/volcengine/responses/test_volcengine_responses_transformation.py (94%) create mode 100644 tests/unit/llms/voyage/rerank/__init__.py rename tests/{test_litellm => unit}/llms/voyage/rerank/test_voyage_rerank_transformation.py (100%) rename tests/{test_litellm => unit}/llms/voyage/test_voyage_contextual_embedding.py (100%) rename tests/{test_litellm => unit}/llms/voyage/test_voyage_multimodal_embedding.py (100%) create mode 100644 tests/unit/llms/watsonx/__init__.py create mode 100644 tests/unit/llms/watsonx/audio_transcription/__init__.py create mode 100644 tests/unit/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py rename tests/{test_litellm => unit}/llms/watsonx/embed/test_watsonx_embedding_transformation.py (100%) rename tests/{test_litellm => unit}/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py (100%) create mode 100644 tests/unit/llms/watsonx/rerank/__init__.py rename tests/{test_litellm => unit}/llms/watsonx/rerank/test_watsonx_rerank.py (100%) create mode 100644 tests/unit/llms/watsonx/test_watsonx.py rename tests/{test_litellm => unit}/llms/watsonx/test_watsonx_common_utils.py (100%) rename tests/{test_litellm => unit}/llms/xai/responses/test_xai_responses_transformation.py (100%) create mode 100644 tests/unit/llms/you_com/__init__.py rename tests/{test_litellm => unit}/llms/you_com/test_you_com_search.py (100%) rename tests/{test_litellm => unit}/llms/zai/test_zai_provider.py (100%) diff --git a/tests/test_litellm/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py b/tests/test_litellm/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py deleted file mode 100644 index e269e782061..00000000000 --- a/tests/test_litellm/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py +++ /dev/null @@ -1,337 +0,0 @@ -""" -Tests for IBM WatsonX Audio Transcription. - -Validates that litellm.transcription transforms requests correctly for WatsonX. -""" - -import json -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - - -import litellm -from litellm.llms.watsonx.audio_transcription.transformation import ( - IBMWatsonXAudioTranscriptionConfig, -) -from litellm.types.utils import TranscriptionResponse - - -class TestWatsonXAudioTranscription: - """Tests for WatsonX audio transcription via litellm.transcription.""" - - @pytest.mark.asyncio - async def test_watsonx_transcription_url_and_headers(self): - """ - Test that litellm.transcription sends request to correct WatsonX URL with proper headers. - """ - captured_request = {} - - async def mock_post(*args, **kwargs): - captured_request["url"] = str(kwargs.get("url", args[0] if args else None)) - captured_request["headers"] = kwargs.get("headers", {}) - captured_request["data"] = kwargs.get("data", {}) - captured_request["files"] = kwargs.get("files", {}) - - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "test transcription", - "duration": 1.0, - } - mock_response.status_code = 200 - return mock_response - - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new=mock_post, - ): - try: - await litellm.atranscription( - model="watsonx/whisper-large-v3-turbo", - file=b"fake_audio_data", - api_base="https://us-south.ml.cloud.ibm.com", - api_key="test-api-key", - project_id="test-project-123", - token="test-bearer-token", - ) - except Exception: - pass # We just want to capture the request - - # Validate URL contains WatsonX audio transcription endpoint - assert "/ml/v1/audio/transcriptions" in captured_request["url"] - assert "version=" in captured_request["url"] - # project_id should NOT be in URL (it should be in form data instead) - assert "project_id=test-project-123" not in captured_request["url"] - - # Validate headers contain WatsonX auth - assert "Authorization" in captured_request["headers"] - assert ( - "Bearer test-bearer-token" in captured_request["headers"]["Authorization"] - ) - - # Validate Content-Type is NOT set (httpx sets multipart/form-data automatically) - assert "Content-Type" not in captured_request["headers"] - - # Validate project_id is in form data, not URL - assert captured_request["data"].get("project_id") == "test-project-123" - - # Validate file is in files dict - assert "file" in captured_request["files"] - - @pytest.mark.asyncio - async def test_watsonx_transcription_request_body(self): - """ - Test that litellm.transcription sends correct request body for WatsonX. - - Validates that: - - Request uses multipart/form-data (data + files) - - Model name has watsonx/ prefix removed - - project_id is in form data, not URL - - Audio file is in files dict - - OpenAI params are included in form data - """ - captured_request = {} - - async def mock_post(*args, **kwargs): - captured_request["data"] = kwargs.get("data", {}) - captured_request["files"] = kwargs.get("files", {}) - - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "test transcription", - "duration": 1.0, - } - mock_response.status_code = 200 - return mock_response - - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new=mock_post, - ): - try: - await litellm.atranscription( - model="watsonx/whisper-large-v3-turbo", - file=b"fake_audio_data", - api_base="https://us-south.ml.cloud.ibm.com", - api_key="test-api-key", - project_id="test-project-123", - token="test-bearer-token", - language="en", - temperature=0.5, - ) - except Exception: - pass # We just want to capture the request - - # Validate form data contains expected fields - data = captured_request.get("data", {}) - - print("JSON DUMPS captured_request:") - print(json.dumps(captured_request, indent=4, default=str)) - - # Model name should NOT have watsonx/ prefix - assert data.get("model") == "whisper-large-v3-turbo" - - # project_id should be in form data - assert data.get("project_id") == "test-project-123" - - # OpenAI params should be in form data - assert data.get("language") == "en" - assert data.get("temperature") == 0.5 - # response_format should NOT be set by default - only send what user specifies - assert "response_format" not in data - - # Validate file is in files dict (multipart/form-data) - files = captured_request.get("files", {}) - assert "file" in files - assert isinstance( - files["file"], tuple - ) # Should be (filename, content, content_type) - - @pytest.mark.asyncio - async def test_watsonx_transcription_only_user_params_sent_with_project_id(self): - """ - Test that only user-specified params are sent in request body to WatsonX. - - LiteLLM should NOT add extra params like response_format if user didn't specify them. - """ - captured_request = {} - - async def mock_post(*args, **kwargs): - captured_request["data"] = kwargs.get("data", {}) - captured_request["files"] = kwargs.get("files", {}) - - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "test transcription", - "duration": 1.0, - } - mock_response.status_code = 200 - return mock_response - - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new=mock_post, - ): - try: - # Minimal request - only required params - await litellm.atranscription( - model="watsonx/whisper-large-v3-turbo", - file=b"fake_audio_data", - api_base="https://us-south.ml.cloud.ibm.com", - api_key="test-api-key", - project_id="test-project-123", - token="test-bearer-token", - ) - except Exception: - pass # We just want to capture the request - - data = captured_request.get("data", {}) - - # These are the ONLY keys that should be in data - expected_keys = {"model", "project_id"} - actual_keys = set(data.keys()) - - assert actual_keys == expected_keys, ( - f"Request body should only contain {expected_keys}, " - f"but got {actual_keys}. " - f"Extra keys: {actual_keys - expected_keys}" - ) - - # Specifically verify response_format is NOT added - assert ( - "response_format" not in data - ), "response_format should NOT be added by default" - - # Verify file is sent separately - files = captured_request.get("files", {}) - assert "file" in files - - @pytest.mark.asyncio - async def test_watsonx_transcription_only_user_params_sent_with_space_id(self): - """ - Test that only user-specified params are sent in request body to WatsonX. - - LiteLLM should NOT add extra params like response_format if user didn't specify them. - """ - captured_request = {} - - async def mock_post(*args, **kwargs): - captured_request["data"] = kwargs.get("data", {}) - captured_request["files"] = kwargs.get("files", {}) - - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "test transcription", - "duration": 1.0, - } - mock_response.status_code = 200 - return mock_response - - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new=mock_post, - ): - try: - # Minimal request - only required params - await litellm.atranscription( - model="watsonx/whisper-large-v3-turbo", - file=b"fake_audio_data", - api_base="https://us-south.ml.cloud.ibm.com", - api_key="test-api-key", - space_id="test-space_id-123", - token="test-bearer-token", - ) - except Exception: - pass # We just want to capture the request - - data = captured_request.get("data", {}) - - # These are the ONLY keys that should be in data - expected_keys = {"model", "space_id"} - actual_keys = set(data.keys()) - - assert actual_keys == expected_keys, ( - f"Request body should only contain {expected_keys}, " - f"but got {actual_keys}. " - f"Extra keys: {actual_keys - expected_keys}" - ) - - # Specifically verify response_format is NOT added - assert ( - "response_format" not in data - ), "response_format should NOT be added by default" - - # Verify file is sent separately - files = captured_request.get("files", {}) - assert "file" in files - - def test_transform_audio_transcription_response_removes_model_field(self): - """ - Test that transform_audio_transcription_response removes the 'model' field - from WatsonX response before creating TranscriptionResponse. - - This test ensures that when WatsonX returns a response with a 'model' field, - it is removed before creating the TranscriptionResponse object, since - TranscriptionResponse doesn't accept a 'model' parameter. - """ - handler = IBMWatsonXAudioTranscriptionConfig() - - # Mock response with 'model' field (as WatsonX may return) - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "Hello, this is a test transcription.", - "model": "whisper-large-v3-turbo", # This field should be removed - "duration": 5.5, - } - mock_response.text = '{"text": "Hello, this is a test transcription.", "model": "whisper-large-v3-turbo", "duration": 5.5}' - - # This should not raise a TypeError - model field should be removed - result = handler.transform_audio_transcription_response(mock_response) - - # Verify the result is a TranscriptionResponse - assert isinstance(result, TranscriptionResponse) - - # Verify the text is correct - assert result.text == "Hello, this is a test transcription." - - # Verify duration is set via dictionary assignment - assert result["duration"] == 5.5 - - # Verify the model field is NOT in the serialized result - # Check via model_dump() or dict() to ensure it's not in the output - try: - result_dict = result.model_dump() - except AttributeError: - # Fallback for pydantic v1 - result_dict = result.dict() - - # The 'model' field should not be in the result - assert "model" not in result_dict, "Model field should be removed from response" - - def test_transform_audio_transcription_response_without_model_field(self): - """ - Test that transform_audio_transcription_response works correctly - when WatsonX response doesn't include a 'model' field. - """ - handler = IBMWatsonXAudioTranscriptionConfig() - - # Mock response without 'model' field - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "Hello, this is a test transcription.", - "duration": 5.5, - } - mock_response.text = ( - '{"text": "Hello, this is a test transcription.", "duration": 5.5}' - ) - - result = handler.transform_audio_transcription_response(mock_response) - - # Verify the result is a TranscriptionResponse - assert isinstance(result, TranscriptionResponse) - - # Verify the text is correct - assert result.text == "Hello, this is a test transcription." - - # Verify duration is set via dictionary assignment - assert result["duration"] == 5.5 diff --git a/tests/test_litellm/llms/watsonx/test_watsonx.py b/tests/test_litellm/llms/watsonx/test_watsonx.py deleted file mode 100644 index 285afffefc0..00000000000 --- a/tests/test_litellm/llms/watsonx/test_watsonx.py +++ /dev/null @@ -1,577 +0,0 @@ -import json - -from typing import Optional -from unittest.mock import Mock, patch - -import pytest - -import litellm -from litellm import completion -from litellm.llms.custom_httpx.http_handler import HTTPHandler - - -@pytest.fixture -def watsonx_chat_completion_call(): - def _call( - model="watsonx/my-test-model", - messages=None, - api_key="test_api_key", - space_id: Optional[str] = None, - headers=None, - client=None, - patch_token_call=True, - ): - if messages is None: - messages = [{"role": "user", "content": "Hello, how are you?"}] - if client is None: - client = HTTPHandler() - - if patch_token_call: - mock_response = Mock() - mock_response.json.return_value = { - "access_token": "mock_access_token", - "expires_in": 3600, - } - mock_response.raise_for_status = Mock() # No-op to simulate no exception - - with ( - patch.object(client, "post") as mock_post, - patch.object( - litellm.module_level_client, "post", return_value=mock_response - ) as mock_get, - ): - try: - completion( - model=model, - messages=messages, - api_key=api_key, - headers=headers or {}, - client=client, - space_id=space_id, - ) - except Exception as e: - print(e) - - return mock_post, mock_get - else: - with patch.object(client, "post") as mock_post: - try: - completion( - model=model, - messages=messages, - api_key=api_key, - headers=headers or {}, - client=client, - space_id=space_id, - ) - except Exception as e: - print(e) - return mock_post, None - - return _call - - -def test_watsonx_deployment_model_id_not_in_payload( - monkeypatch, watsonx_chat_completion_call -): - """Test that deployment models do not include 'model_id' in the request payload""" - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - model = "watsonx/deployment/test-deployment-id" - messages = [{"role": "user", "content": "Test message"}] - - mock_post, _ = watsonx_chat_completion_call(model=model, messages=messages) - - assert mock_post.call_count == 1 - json_data = json.loads(mock_post.call_args.kwargs["data"]) - # Ensure model_id is not in the payload for deployment models - assert "model_id" not in json_data or json_data["model_id"] is None - # Ensure project_id is also not in the payload for deployment models - assert "project_id" not in json_data or json_data["project_id"] is None - - -def test_watsonx_regular_model_includes_model_id( - monkeypatch, watsonx_chat_completion_call -): - """Test that regular models include 'model_id' in the request payload""" - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - model = "watsonx/regular-model" - messages = [{"role": "user", "content": "Test message"}] - - mock_post, _ = watsonx_chat_completion_call(model=model, messages=messages) - - assert mock_post.call_count == 1 - json_data = json.loads(mock_post.call_args.kwargs["data"]) - # Ensure model_id is included in the payload for regular models - assert "model_id" in json_data - assert json_data["model_id"] == "regular-model" # Provider prefix is stripped - # Ensure project_id is also included for regular models - assert "project_id" in json_data - - -@pytest.fixture -def watsonx_completion_call(): - def _call( - model="watsonx_text/my-test-model", - prompt="Hello, how are you?", - api_key="test_api_key", - space_id: Optional[str] = None, - headers=None, - client=None, - patch_token_call=True, - ): - if client is None: - client = HTTPHandler() - - if patch_token_call: - mock_response = Mock() - mock_response.json.return_value = { - "access_token": "mock_access_token", - "expires_in": 3600, - } - mock_response.raise_for_status = Mock() - - with ( - patch.object(client, "post") as mock_post, - patch.object( - litellm.module_level_client, "post", return_value=mock_response - ) as mock_get, - ): - try: - litellm.text_completion( - model=model, - prompt=prompt, - api_key=api_key, - headers=headers or {}, - client=client, - space_id=space_id, - ) - except Exception as e: - print(e) - - return mock_post, mock_get - else: - with patch.object(client, "post") as mock_post: - try: - litellm.text_completion( - model=model, - prompt=prompt, - api_key=api_key, - headers=headers or {}, - client=client, - space_id=space_id, - ) - except Exception as e: - print(e) - return mock_post, None - - return _call - - -def test_watsonx_completion_deployment_model_id_not_in_payload( - monkeypatch, watsonx_completion_call -): - """Test that deployment models do not include 'model_id' in completion request payload""" - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - model = "watsonx_text/deployment/test-deployment-id" - prompt = "Test prompt" - - mock_post, _ = watsonx_completion_call(model=model, prompt=prompt) - - assert mock_post.call_count == 1 - json_data = json.loads(mock_post.call_args.kwargs["data"]) - # Ensure model_id is not in the payload for deployment models - assert "model_id" not in json_data - # Ensure project_id is also not in the payload for deployment models - assert "project_id" not in json_data - - -def test_watsonx_completion_regular_model_includes_model_id( - monkeypatch, watsonx_completion_call -): - """Test that regular models include 'model_id' in completion request payload""" - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - model = "watsonx_text/regular-model" - prompt = "Test prompt" - - mock_post, _ = watsonx_completion_call(model=model, prompt=prompt) - - assert mock_post.call_count == 1 - json_data = json.loads(mock_post.call_args.kwargs["data"]) - # Ensure model_id is included in the payload for regular models - assert "model_id" in json_data - assert json_data["model_id"] == "regular-model" # Provider prefix is stripped - # Ensure project_id is also included for regular models - assert "project_id" in json_data - - -def test_watsonx_gpt_oss_prompt_transformation(monkeypatch): - """ - Test that gpt-oss-120b model transforms messages to proper format instead of simple concatenation. - - This test calls litellm.completion (sync) and verifies what gets sent in the final POST request body. - Input messages should be transformed using the HuggingFace chat template from openai/gpt-oss-120b, - not just concatenated as "You are chatgpt Hi there". - """ - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - - # Test with gpt-oss model using watsonx_text provider (text generation endpoint) - model = "watsonx_text/openai/gpt-oss-120b" - - # Input messages - messages = [ - {"role": "system", "content": "You are chatgpt"}, - {"role": "user", "content": "Hi there"}, - ] - - client = HTTPHandler() - - # Mock HuggingFace template fetch to make test deterministic and avoid network flakiness. - # The test verifies that prompt transformation occurs (not simple concatenation), not the exact - # HuggingFace template format. Using a mock template that produces the correct format is sufficient. - # - # Mock template that produces gpt-oss-120b-like format. - # Note: This is a simplified version of the actual template. The real template is more complex - # (adds metadata, handles tools, thinking messages, etc.), but this captures the key aspects: - # - Converts system role to developer (matching real template behavior) - # - Uses the same tag structure (<|start|>, <|message|>, <|end|>) - # - Preserves message content - mock_tokenizer_config = { - "status": "success", - "tokenizer": { - "chat_template": "{% for message in messages %}{% if message['role'] == 'system' %}<|start|>developer<|message|>{% else %}<|start|>{{ message['role'] }}<|message|>{% endif %}{{ message['content'] }}<|end|>{% endfor %}", - "bos_token": None, - "eos_token": None, - }, - } - - # Isolate known_tokenizer_config so parallel tests don't interfere. - # monkeypatch.setitem restores the original value on teardown. - hf_model = "openai/gpt-oss-120b" - monkeypatch.setitem(litellm.known_tokenizer_config, hf_model, mock_tokenizer_config) - - # Mock IAM token generation to avoid real HTTP calls. - mock_token_response = Mock() - mock_token_response.json.return_value = { - "access_token": "mock_access_token", - "expires_in": 3600, - } - mock_token_response.raise_for_status = Mock() - - with ( - patch.object(client, "post") as mock_post, - patch.object( - litellm.module_level_client, "post", return_value=mock_token_response - ), - ): - try: - completion( - model=model, - messages=messages, - api_key="test_api_key", - client=client, - ) - except Exception as e: - print(f"Caught expected exception: {e}") - - # Verify the POST was called - assert ( - mock_post.call_count == 1 - ), f"POST should have been called exactly once, got {mock_post.call_count}" - - # Get the request body - call_args = mock_post.call_args - assert "data" in call_args.kwargs, "call_args.kwargs should contain 'data'" - json_data = json.loads(call_args.kwargs["data"]) - - # Verify the transformed input is in the request - assert "input" in json_data, "Request should have 'input' field" - transformed_prompt = json_data["input"] - - # Verify it's NOT simple concatenation - simple_concat = "You are chatgpt Hi there" - assert transformed_prompt != simple_concat, ( - f"Prompt should not be simple concatenation.\n" - f"Expected: Chat template with <|start|> tags\n" - f"Got: {transformed_prompt}" - ) - - # Verify it contains proper chat template formatting - assert "<|start|>" in transformed_prompt, "Prompt should contain <|start|> tag" - assert "<|message|>" in transformed_prompt, "Prompt should contain <|message|> tag" - assert "<|end|>" in transformed_prompt, "Prompt should contain <|end|> tag" - assert ( - "You are chatgpt" in transformed_prompt - ), "Prompt should contain system message content" - assert ( - "Hi there" in transformed_prompt - ), "Prompt should contain user message content" - - -@pytest.mark.asyncio -@pytest.mark.xdist_group("watsonx_heavy") -async def test_watsonx_gpt_oss_uses_async_http_handler(): - """ - Test that verifies async HTTP client is used when fetching HuggingFace templates. - """ - from unittest.mock import AsyncMock, MagicMock, patch - - from litellm.litellm_core_utils.prompt_templates.huggingface_template_handler import ( - _aget_chat_template_file, - ) - - # Mock the async HTTP client - mock_async_client = MagicMock() - mock_get = AsyncMock() - mock_async_client.get = mock_get - - # Create mock response for chat template file - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.content = b"test template content" - mock_get.return_value = mock_response - - # Test the async function directly - with patch( - "litellm.litellm_core_utils.prompt_templates.huggingface_template_handler.get_async_httpx_client", - return_value=mock_async_client, - ): - result = await _aget_chat_template_file(hf_model_name="test/model") - - # Verify async HTTP client was called - assert mock_get.called, "Async HTTP client's get method should be called" - assert mock_get.await_count > 0, "Async HTTP client's get should be awaited" - - # Verify it was called with HuggingFace URL - call_args = mock_get.call_args - assert call_args is not None, "get should have been called with arguments" - called_url = call_args.kwargs.get("url", "") - assert ( - "huggingface.co/test/model" in called_url - ), f"Should call HuggingFace API for test/model, got: {called_url}" - assert result["status"] == "success", "Should return success status" - - -@pytest.mark.parametrize("tokenizer_config_cached", [False, True], ids=["tokenizer_config", "cached_config_jinja"]) -async def test_watsonx_text_gpt_oss_async_completion_fetches_hf_template_off_the_event_loop( - monkeypatch, tokenizer_config_cached -): - import httpx - - from litellm._uuid import uuid - from litellm.litellm_core_utils.prompt_templates import huggingface_template_handler - from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - - hf_model = f"openai/gpt-oss-{uuid.uuid4()}" - chat_template = "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}{% endfor %}" - if tokenizer_config_cached: - cached_config = {"status": "success", "tokenizer": {"bos_token": None, "eos_token": None}} - monkeypatch.setattr(litellm, "known_tokenizer_config", {hf_model: cached_config}) - expected_fetch = f"https://huggingface.co/{hf_model}/raw/main/chat_template.jinja" - else: - monkeypatch.setattr(litellm, "known_tokenizer_config", {}) - expected_fetch = f"https://huggingface.co/{hf_model}/raw/main/tokenizer_config.json" - hf_fetched = [] - captured = {} - - def forbid_sync_client(): - raise AssertionError("sync HuggingFace fetch ran on the request path") - - async def serve_hf_file(url, **kwargs): - hf_fetched.append(url) - if url.endswith(".jinja"): - return httpx.Response(200, content=chat_template.encode()) - return httpx.Response(200, json={"chat_template": chat_template, "bos_token": None, "eos_token": None}) - - monkeypatch.setattr(huggingface_template_handler, "_get_httpx_client", forbid_sync_client) - monkeypatch.setattr(huggingface_template_handler, "get_async_httpx_client", lambda **kwargs: Mock(get=serve_hf_file)) - - def handle(request): - captured["body"] = json.loads(request.content) - return httpx.Response( - 200, - json={ - "model_id": hf_model, - "results": [ - { - "generated_text": "Hi", - "generated_token_count": 1, - "input_token_count": 1, - "stop_reason": "eos_token", - } - ], - }, - ) - - client = AsyncHTTPHandler() - client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) - - response = await litellm.acompletion( - model=f"watsonx_text/{hf_model}", - messages=[{"role": "user", "content": "Hi there"}], - api_base="https://test-api.watsonx.ai", - project_id="test-project-id", - token="test-token", - client=client, - ) - - assert response.choices[0].message.content == "Hi" - assert hf_fetched == [expected_fetch] - assert captured["body"]["input"] == "<|user|>Hi there" - - -def test_watsonx_chat_completion_with_reasoning_effort(monkeypatch): - """ - Test that 'reasoning_effort' is correctly passed through to the WatsonX API payload. - """ - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - - model = "watsonx/openai/gpt-oss-120b" - messages = [{"role": "user", "content": "Test message"}] - - client = HTTPHandler() - - # Mock the token generation call - mock_token_response = Mock() - mock_token_response.json.return_value = { - "access_token": "mock_access_token", - "expires_in": 3600, - } - mock_token_response.raise_for_status = Mock() - - # Call litellm.completion with the new parameter - with ( - patch.object(client, "post") as mock_post, - patch.object( - litellm.module_level_client, "post", return_value=mock_token_response - ), - ): - try: - completion( - model=model, - messages=messages, - api_key="test_api_key", - client=client, - reasoning_effort="low", - ) - except Exception as e: - print(f"Caught expected exception: {e}") - - # Verify the parameter is in the final request payload - assert ( - mock_post.call_count == 1 - ), "The completion endpoint should have been called once." - - # Get the JSON data sent in the POST request - request_kwargs = mock_post.call_args.kwargs - json_data = json.loads(request_kwargs["data"]) - - print("\nRequest payload sent to WatsonX API:") - print(json.dumps(json_data, indent=2)) - - # Check for the parameter at the top level of the payload - assert ( - "reasoning_effort" in json_data - ), "'reasoning_effort' should be at the top level of the payload." - assert ( - json_data["reasoning_effort"] == "low" - ), "The value of 'reasoning_effort' should be 'low'." - - -def test_watsonx_zen_api_key_from_client(monkeypatch, watsonx_chat_completion_call): - """ - Test that zen_api_key can be passed from client code and is used in Authorization header. - """ - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - - model = "watsonx/ibm/granite-3-3-8b-instruct" - messages = [{"role": "user", "content": "What is your favorite color?"}] - - client = HTTPHandler() - - zen_api_key = "U1ZDLWQo=" - - # No need to patch token call since zen_api_key should skip token generation - with patch.object(client, "post") as mock_post: - try: - completion( - model=model, - messages=messages, - api_key="test_api_key", - client=client, - zen_api_key=zen_api_key, - ) - except Exception as e: - print(f"Caught expected exception: {e}") - - # Verify the request was made - assert ( - mock_post.call_count == 1 - ), "The completion endpoint should have been called once." - - # Get the headers sent in the POST request - request_kwargs = mock_post.call_args.kwargs - headers = request_kwargs["headers"] - - print("\nHeaders sent to WatsonX API:") - print(json.dumps(dict(headers), indent=2)) - - # Verify Authorization header uses ZenApiKey format - assert "Authorization" in headers, "Authorization header should be present." - assert headers["Authorization"] == f"ZenApiKey {zen_api_key}", ( - f"Authorization header should use ZenApiKey format. " - f"Expected: 'ZenApiKey {zen_api_key}', Got: '{headers['Authorization']}'" - ) - - -def test_watsonx_zen_api_key_from_env(monkeypatch, watsonx_chat_completion_call): - """ - Test that zen_api_key from environment variable is used in Authorization header. - """ - monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") - monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - - zen_api_key = "U1ZDLWxpdG--===" - monkeypatch.setenv("WATSONX_ZENAPIKEY", zen_api_key) - - model = "watsonx/ibm/granite-3-3-8b-instruct" - messages = [{"role": "user", "content": "What is your favorite color?"}] - - client = HTTPHandler() - - # No need to patch token call since zen_api_key should skip token generation - with patch.object(client, "post") as mock_post: - try: - completion( - model=model, - messages=messages, - api_key="test_api_key", - client=client, - ) - except Exception as e: - print(f"Caught expected exception: {e}") - - # Verify the request was made - assert ( - mock_post.call_count == 1 - ), "The completion endpoint should have been called once." - - # Get the headers sent in the POST request - request_kwargs = mock_post.call_args.kwargs - headers = request_kwargs["headers"] - - print("\nHeaders sent to WatsonX API:") - print(json.dumps(dict(headers), indent=2)) - - # Verify Authorization header uses ZenApiKey format - assert "Authorization" in headers, "Authorization header should be present." - assert headers["Authorization"] == f"ZenApiKey {zen_api_key}", ( - f"Authorization header should use ZenApiKey format. " - f"Expected: 'ZenApiKey {zen_api_key}', Got: '{headers['Authorization']}'" - ) diff --git a/tests/test_litellm/llms/xai/xai_responses/__init__.py b/tests/test_litellm/llms/xai/xai_responses/__init__.py deleted file mode 100644 index 330e9f5a560..00000000000 --- a/tests/test_litellm/llms/xai/xai_responses/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# XAI Responses API tests diff --git a/tests/test_litellm/llms/xai/xai_responses/test_transformation.py b/tests/test_litellm/llms/xai/xai_responses/test_transformation.py deleted file mode 100644 index 3ea3fe631bd..00000000000 --- a/tests/test_litellm/llms/xai/xai_responses/test_transformation.py +++ /dev/null @@ -1,105 +0,0 @@ -""" -Tests for XAI Responses API transformation - -Tests the XAIResponsesAPIConfig class that handles XAI-specific -transformations for the Responses API. - -Source: litellm/llms/xai/responses/transformation.py -""" - - - -import pytest -from litellm.types.utils import LlmProviders -from litellm.utils import ProviderConfigManager -from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig -from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams - - -class TestXAIResponsesAPITransformation: - """Test XAI Responses API configuration and transformations""" - - def test_xai_provider_config_registration(self): - """Test that XAI provider returns XAIResponsesAPIConfig""" - config = ProviderConfigManager.get_provider_responses_api_config( - model="xai/grok-4-fast", - provider=LlmProviders.XAI, - ) - - assert config is not None, "Config should not be None for XAI provider" - assert isinstance( - config, XAIResponsesAPIConfig - ), f"Expected XAIResponsesAPIConfig, got {type(config)}" - assert ( - config.custom_llm_provider == LlmProviders.XAI - ), "custom_llm_provider should be XAI" - - def test_code_interpreter_container_field_removed(self): - """Test that container field is removed from code_interpreter tools""" - config = XAIResponsesAPIConfig() - - params = ResponsesAPIOptionalRequestParams( - tools=[{"type": "code_interpreter", "container": {"type": "auto"}}] - ) - - result = config.map_openai_params( - response_api_optional_params=params, model="grok-4-fast", drop_params=False - ) - - assert "tools" in result - assert len(result["tools"]) == 1 - assert result["tools"][0]["type"] == "code_interpreter" - assert ( - "container" not in result["tools"][0] - ), "Container field should be removed" - - def test_instructions_parameter_forwarded(self): - """xAI supports 'instructions' on /v1/responses, so it must survive param mapping""" - config = XAIResponsesAPIConfig() - - params = ResponsesAPIOptionalRequestParams( - instructions="You are a helpful assistant.", temperature=0.7 - ) - - result = config.map_openai_params( - response_api_optional_params=params, model="grok-4-fast", drop_params=False - ) - - assert result.get("instructions") == "You are a helpful assistant." - assert result.get("temperature") == 0.7, "Other params should be preserved" - - def test_supported_params_includes_instructions(self): - """A system message bridged to 'instructions' must not be rejected for xAI""" - config = XAIResponsesAPIConfig() - supported = config.get_supported_openai_params("grok-4-fast") - - assert "instructions" in supported, "instructions should be supported" - assert "tools" in supported, "tools should be supported" - assert "temperature" in supported, "temperature should be supported" - assert "model" in supported, "model should be supported" - - def test_xai_responses_endpoint_url(self): - """Test that get_complete_url returns correct XAI endpoint""" - config = XAIResponsesAPIConfig() - - # Test with default XAI API base - url = config.get_complete_url(api_base=None, litellm_params={}) - assert ( - url == "https://api.x.ai/v1/responses" - ), f"Expected XAI responses endpoint, got {url}" - - # Test with custom api_base - custom_url = config.get_complete_url( - api_base="https://custom.x.ai/v1", litellm_params={} - ) - assert ( - custom_url == "https://custom.x.ai/v1/responses" - ), f"Expected custom endpoint, got {custom_url}" - - # Test with trailing slash - url_with_slash = config.get_complete_url( - api_base="https://api.x.ai/v1/", litellm_params={} - ) - assert ( - url_with_slash == "https://api.x.ai/v1/responses" - ), "Should handle trailing slash" diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py similarity index 98% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py index 4b710175a48..e2fb81bc240 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_location.py @@ -98,6 +98,8 @@ class TestCountTokensLocationResolution: self, counter, monkeypatch ): """Claude models without any location should default to us-east5.""" + monkeypatch.delenv("VERTEXAI_LOCATION", raising=False) + monkeypatch.delenv("VERTEX_LOCATION", raising=False) captured = {} async def fake_ensure_access_token( diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py similarity index 60% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py index f7df4507651..15df8e47af3 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py @@ -1,6 +1,21 @@ +import pytest + import litellm +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + def test_reasoning_effort_stays_unsupported_on_vertex_partner_models(local_model_cost_map): assert "reasoning_effort" in litellm.get_supported_openai_params( model="mistral-medium-3", custom_llm_provider="mistral" diff --git a/tests/unit/llms/vertex_ai/vertex_gemma_models/__init__.py b/tests/unit/llms/vertex_ai/vertex_gemma_models/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py rename to tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py diff --git a/tests/unit/llms/vertex_ai/videos/__init__.py b/tests/unit/llms/vertex_ai/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py rename to tests/unit/llms/vertex_ai/videos/test_vertex_video_transformation.py diff --git a/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py b/tests/unit/llms/volcengine/responses/test_volcengine_responses_transformation.py similarity index 94% rename from tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py rename to tests/unit/llms/volcengine/responses/test_volcengine_responses_transformation.py index d42bf7b7a1c..5c8d67ecc70 100644 --- a/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py +++ b/tests/unit/llms/volcengine/responses/test_volcengine_responses_transformation.py @@ -137,30 +137,6 @@ class TestVolcengineResponsesAPITransformation: with pytest.raises(ValueError, match='Volcengine API key is required\\. Set ARK_API_KEY /'): config.validate_environment(headers={}, model="volcengine/demo", litellm_params={}) - def test_unsupported_params_are_dropped_with_extra_body(self): - """Unknown fields (including extra_body) should be dropped before send.""" - config = VolcEngineResponsesAPIConfig() - - request = config.transform_responses_api_request( - model="volcengine/demo-model", - input="hi", - response_api_optional_request_params={ - "unsupported_custom_param": 0.1, - "temperature": 0.2, - "metadata": {"k": "v"}, - "extra_body": {"unsupported_custom_param": 1, "temperature": 0.3}, - }, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) - - assert "unsupported_custom_param" not in request - assert "metadata" not in request - assert request["temperature"] == 0.2 - assert "extra_body" in request - assert "unsupported_custom_param" not in request["extra_body"] - assert request["extra_body"]["temperature"] == 0.3 - def test_valid_thinking_caching_and_expire_at_pass(self): """Documented params should pass through without validation errors.""" config = VolcEngineResponsesAPIConfig() diff --git a/tests/unit/llms/voyage/rerank/__init__.py b/tests/unit/llms/voyage/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py b/tests/unit/llms/voyage/rerank/test_voyage_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py rename to tests/unit/llms/voyage/rerank/test_voyage_rerank_transformation.py diff --git a/tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py b/tests/unit/llms/voyage/test_voyage_contextual_embedding.py similarity index 100% rename from tests/test_litellm/llms/voyage/test_voyage_contextual_embedding.py rename to tests/unit/llms/voyage/test_voyage_contextual_embedding.py diff --git a/tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py b/tests/unit/llms/voyage/test_voyage_multimodal_embedding.py similarity index 100% rename from tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py rename to tests/unit/llms/voyage/test_voyage_multimodal_embedding.py diff --git a/tests/unit/llms/watsonx/__init__.py b/tests/unit/llms/watsonx/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/watsonx/audio_transcription/__init__.py b/tests/unit/llms/watsonx/audio_transcription/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py b/tests/unit/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py new file mode 100644 index 00000000000..efe592f515e --- /dev/null +++ b/tests/unit/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py @@ -0,0 +1,85 @@ +""" +Tests for IBM WatsonX Audio Transcription. + +Validates the WatsonX transcription response transformation. +""" + +from unittest.mock import MagicMock + +from litellm.llms.watsonx.audio_transcription.transformation import ( + IBMWatsonXAudioTranscriptionConfig, +) +from litellm.types.utils import TranscriptionResponse + + +class TestWatsonXAudioTranscription: + def test_transform_audio_transcription_response_removes_model_field(self): + """ + Test that transform_audio_transcription_response removes the 'model' field + from WatsonX response before creating TranscriptionResponse. + + This test ensures that when WatsonX returns a response with a 'model' field, + it is removed before creating the TranscriptionResponse object, since + TranscriptionResponse doesn't accept a 'model' parameter. + """ + handler = IBMWatsonXAudioTranscriptionConfig() + + # Mock response with 'model' field (as WatsonX may return) + mock_response = MagicMock() + mock_response.json.return_value = { + "text": "Hello, this is a test transcription.", + "model": "whisper-large-v3-turbo", # This field should be removed + "duration": 5.5, + } + mock_response.text = '{"text": "Hello, this is a test transcription.", "model": "whisper-large-v3-turbo", "duration": 5.5}' + + # This should not raise a TypeError - model field should be removed + result = handler.transform_audio_transcription_response(mock_response) + + # Verify the result is a TranscriptionResponse + assert isinstance(result, TranscriptionResponse) + + # Verify the text is correct + assert result.text == "Hello, this is a test transcription." + + # Verify duration is set via dictionary assignment + assert result["duration"] == 5.5 + + # Verify the model field is NOT in the serialized result + # Check via model_dump() or dict() to ensure it's not in the output + try: + result_dict = result.model_dump() + except AttributeError: + # Fallback for pydantic v1 + result_dict = result.dict() + + # The 'model' field should not be in the result + assert "model" not in result_dict, "Model field should be removed from response" + + def test_transform_audio_transcription_response_without_model_field(self): + """ + Test that transform_audio_transcription_response works correctly + when WatsonX response doesn't include a 'model' field. + """ + handler = IBMWatsonXAudioTranscriptionConfig() + + # Mock response without 'model' field + mock_response = MagicMock() + mock_response.json.return_value = { + "text": "Hello, this is a test transcription.", + "duration": 5.5, + } + mock_response.text = ( + '{"text": "Hello, this is a test transcription.", "duration": 5.5}' + ) + + result = handler.transform_audio_transcription_response(mock_response) + + # Verify the result is a TranscriptionResponse + assert isinstance(result, TranscriptionResponse) + + # Verify the text is correct + assert result.text == "Hello, this is a test transcription." + + # Verify duration is set via dictionary assignment + assert result["duration"] == 5.5 diff --git a/tests/test_litellm/llms/watsonx/embed/test_watsonx_embedding_transformation.py b/tests/unit/llms/watsonx/embed/test_watsonx_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/watsonx/embed/test_watsonx_embedding_transformation.py rename to tests/unit/llms/watsonx/embed/test_watsonx_embedding_transformation.py diff --git a/tests/test_litellm/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py b/tests/unit/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py similarity index 100% rename from tests/test_litellm/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py rename to tests/unit/llms/watsonx/passthrough/test_watsonx_passthrough_transformation.py diff --git a/tests/unit/llms/watsonx/rerank/__init__.py b/tests/unit/llms/watsonx/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py b/tests/unit/llms/watsonx/rerank/test_watsonx_rerank.py similarity index 100% rename from tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py rename to tests/unit/llms/watsonx/rerank/test_watsonx_rerank.py diff --git a/tests/unit/llms/watsonx/test_watsonx.py b/tests/unit/llms/watsonx/test_watsonx.py new file mode 100644 index 00000000000..077539c9acd --- /dev/null +++ b/tests/unit/llms/watsonx/test_watsonx.py @@ -0,0 +1,74 @@ +import json +from unittest.mock import Mock + +import pytest + +import litellm + + +@pytest.mark.parametrize("tokenizer_config_cached", [False, True], ids=["tokenizer_config", "cached_config_jinja"]) +async def test_watsonx_text_gpt_oss_async_completion_fetches_hf_template_off_the_event_loop( + monkeypatch, tokenizer_config_cached +): + import httpx + + from litellm._uuid import uuid + from litellm.litellm_core_utils.prompt_templates import huggingface_template_handler + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + hf_model = f"openai/gpt-oss-{uuid.uuid4()}" + chat_template = "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}{% endfor %}" + if tokenizer_config_cached: + cached_config = {"status": "success", "tokenizer": {"bos_token": None, "eos_token": None}} + monkeypatch.setattr(litellm, "known_tokenizer_config", {hf_model: cached_config}) + expected_fetch = f"https://huggingface.co/{hf_model}/raw/main/chat_template.jinja" + else: + monkeypatch.setattr(litellm, "known_tokenizer_config", {}) + expected_fetch = f"https://huggingface.co/{hf_model}/raw/main/tokenizer_config.json" + hf_fetched = [] + captured = {} + + def forbid_sync_client(): + raise AssertionError("sync HuggingFace fetch ran on the request path") + + async def serve_hf_file(url, **kwargs): + hf_fetched.append(url) + if url.endswith(".jinja"): + return httpx.Response(200, content=chat_template.encode()) + return httpx.Response(200, json={"chat_template": chat_template, "bos_token": None, "eos_token": None}) + + monkeypatch.setattr(huggingface_template_handler, "_get_httpx_client", forbid_sync_client) + monkeypatch.setattr(huggingface_template_handler, "get_async_httpx_client", lambda **kwargs: Mock(get=serve_hf_file)) + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "model_id": hf_model, + "results": [ + { + "generated_text": "Hi", + "generated_token_count": 1, + "input_token_count": 1, + "stop_reason": "eos_token", + } + ], + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model=f"watsonx_text/{hf_model}", + messages=[{"role": "user", "content": "Hi there"}], + api_base="https://test-api.watsonx.ai", + project_id="test-project-id", + token="test-token", + client=client, + ) + + assert response.choices[0].message.content == "Hi" + assert hf_fetched == [expected_fetch] + assert captured["body"]["input"] == "<|user|>Hi there" diff --git a/tests/test_litellm/llms/watsonx/test_watsonx_common_utils.py b/tests/unit/llms/watsonx/test_watsonx_common_utils.py similarity index 100% rename from tests/test_litellm/llms/watsonx/test_watsonx_common_utils.py rename to tests/unit/llms/watsonx/test_watsonx_common_utils.py diff --git a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py b/tests/unit/llms/xai/responses/test_xai_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py rename to tests/unit/llms/xai/responses/test_xai_responses_transformation.py diff --git a/tests/unit/llms/you_com/__init__.py b/tests/unit/llms/you_com/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/you_com/test_you_com_search.py b/tests/unit/llms/you_com/test_you_com_search.py similarity index 100% rename from tests/test_litellm/llms/you_com/test_you_com_search.py rename to tests/unit/llms/you_com/test_you_com_search.py diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/unit/llms/zai/test_zai_provider.py similarity index 100% rename from tests/test_litellm/llms/zai/test_zai_provider.py rename to tests/unit/llms/zai/test_zai_provider.py From ef6237b9f2bcafba3e7e535a283dd542d67adc70 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:02:53 +0000 Subject: [PATCH 30/47] test: migrate phase 12 legacy llm provider tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 14 ++++++++++++++ ...est_openrouter_responses_transformation.py | 0 .../parallel_ai/test_parallel_ai_search.py | 0 .../test_parallel_ai_search_gateway.py | 0 .../llms/parasail/test_parasail.py | 0 .../test_perplexity_chat_transformation.py | 0 ...est_perplexity_embedding_transformation.py | 0 ...est_perplexity_responses_transformation.py | 0 .../test_publicai_chat_transformation.py | 0 .../chat/test_ragflow_chat_transformation.py | 0 .../test_recraft_image_edit_transformation.py | 0 .../test_recraft_image_gen_transformation.py | 0 .../test_text_to_speech_transformation.py | 0 .../test_runway_video_transformation.py | 0 .../test_s3_vectors_transformation.py | 4 ---- .../llms/sap/test_sap_fetch_creds.py | 0 ...eway_audio_transcription_transformation.py | 0 .../test_snowflake_native_endpoints.py | 0 .../test_soniox_provider_registration.py | 0 .../test_stability_image_generation.py | 19 +------------------ .../chat/test_tencent_chat_transformation.py | 0 21 files changed, 15 insertions(+), 22 deletions(-) rename tests/{test_litellm => unit}/llms/openrouter/responses/test_openrouter_responses_transformation.py (100%) rename tests/{test_litellm => unit}/llms/parallel_ai/test_parallel_ai_search.py (100%) rename tests/{test_litellm => unit}/llms/parallel_ai/test_parallel_ai_search_gateway.py (100%) rename tests/{test_litellm => unit}/llms/parasail/test_parasail.py (100%) rename tests/{test_litellm => unit}/llms/perplexity/chat/test_perplexity_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/perplexity/embedding/test_perplexity_embedding_transformation.py (100%) rename tests/{test_litellm => unit}/llms/perplexity/responses/test_perplexity_responses_transformation.py (100%) rename tests/{test_litellm => unit}/llms/publicai/test_publicai_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/ragflow/chat/test_ragflow_chat_transformation.py (100%) rename tests/{test_litellm => unit}/llms/recraft/image_edit/test_recraft_image_edit_transformation.py (100%) rename tests/{test_litellm => unit}/llms/recraft/image_generation/test_recraft_image_gen_transformation.py (100%) rename tests/{test_litellm => unit}/llms/runwayml/test_text_to_speech_transformation.py (100%) rename tests/{test_litellm => unit}/llms/runwayml/videos/test_runway_video_transformation.py (100%) rename tests/{test_litellm => unit}/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py (99%) rename tests/{test_litellm => unit}/llms/sap/test_sap_fetch_creds.py (100%) rename tests/{test_litellm => unit}/llms/scaleway/test_scaleway_audio_transcription_transformation.py (100%) rename tests/{test_litellm => unit}/llms/snowflake/test_snowflake_native_endpoints.py (100%) rename tests/{test_litellm => unit}/llms/soniox/test_soniox_provider_registration.py (100%) rename tests/{test_litellm => unit}/llms/stability/image_generation/test_stability_image_generation.py (93%) rename tests/{test_litellm => unit}/llms/tencent/chat/test_tencent_chat_transformation.py (100%) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 017e63ed1b8..fa91437964f 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,6 +1,7 @@ import os from typing import Final +import litellm import pytest from pytest_socket import enable_socket, socket_allow_hosts @@ -9,6 +10,19 @@ os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + def _allow_loopback_only() -> None: socket_allow_hosts(LOOPBACK_HOSTS, allow_unix_socket=True) diff --git a/tests/test_litellm/llms/openrouter/responses/test_openrouter_responses_transformation.py b/tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/openrouter/responses/test_openrouter_responses_transformation.py rename to tests/unit/llms/openrouter/responses/test_openrouter_responses_transformation.py diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py b/tests/unit/llms/parallel_ai/test_parallel_ai_search.py similarity index 100% rename from tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py rename to tests/unit/llms/parallel_ai/test_parallel_ai_search.py diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py b/tests/unit/llms/parallel_ai/test_parallel_ai_search_gateway.py similarity index 100% rename from tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py rename to tests/unit/llms/parallel_ai/test_parallel_ai_search_gateway.py diff --git a/tests/test_litellm/llms/parasail/test_parasail.py b/tests/unit/llms/parasail/test_parasail.py similarity index 100% rename from tests/test_litellm/llms/parasail/test_parasail.py rename to tests/unit/llms/parasail/test_parasail.py diff --git a/tests/test_litellm/llms/perplexity/chat/test_perplexity_chat_transformation.py b/tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/perplexity/chat/test_perplexity_chat_transformation.py rename to tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py diff --git a/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py b/tests/unit/llms/perplexity/embedding/test_perplexity_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py rename to tests/unit/llms/perplexity/embedding/test_perplexity_embedding_transformation.py diff --git a/tests/test_litellm/llms/perplexity/responses/test_perplexity_responses_transformation.py b/tests/unit/llms/perplexity/responses/test_perplexity_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/perplexity/responses/test_perplexity_responses_transformation.py rename to tests/unit/llms/perplexity/responses/test_perplexity_responses_transformation.py diff --git a/tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py b/tests/unit/llms/publicai/test_publicai_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py rename to tests/unit/llms/publicai/test_publicai_chat_transformation.py diff --git a/tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py b/tests/unit/llms/ragflow/chat/test_ragflow_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py rename to tests/unit/llms/ragflow/chat/test_ragflow_chat_transformation.py diff --git a/tests/test_litellm/llms/recraft/image_edit/test_recraft_image_edit_transformation.py b/tests/unit/llms/recraft/image_edit/test_recraft_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/recraft/image_edit/test_recraft_image_edit_transformation.py rename to tests/unit/llms/recraft/image_edit/test_recraft_image_edit_transformation.py diff --git a/tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py b/tests/unit/llms/recraft/image_generation/test_recraft_image_gen_transformation.py similarity index 100% rename from tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py rename to tests/unit/llms/recraft/image_generation/test_recraft_image_gen_transformation.py diff --git a/tests/test_litellm/llms/runwayml/test_text_to_speech_transformation.py b/tests/unit/llms/runwayml/test_text_to_speech_transformation.py similarity index 100% rename from tests/test_litellm/llms/runwayml/test_text_to_speech_transformation.py rename to tests/unit/llms/runwayml/test_text_to_speech_transformation.py diff --git a/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py b/tests/unit/llms/runwayml/videos/test_runway_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py rename to tests/unit/llms/runwayml/videos/test_runway_video_transformation.py diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py b/tests/unit/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py similarity index 99% rename from tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py rename to tests/unit/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py index 781e92ea7d9..c39887d86ce 100644 --- a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py +++ b/tests/unit/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py @@ -55,10 +55,6 @@ def _search_kwargs(**overrides): class TestS3VectorsVectorStoreConfig: - def test_init(self): - config = S3VectorsVectorStoreConfig() - assert config is not None - def test_get_supported_openai_params(self): config = S3VectorsVectorStoreConfig() params = config.get_supported_openai_params("test-model") diff --git a/tests/test_litellm/llms/sap/test_sap_fetch_creds.py b/tests/unit/llms/sap/test_sap_fetch_creds.py similarity index 100% rename from tests/test_litellm/llms/sap/test_sap_fetch_creds.py rename to tests/unit/llms/sap/test_sap_fetch_creds.py diff --git a/tests/test_litellm/llms/scaleway/test_scaleway_audio_transcription_transformation.py b/tests/unit/llms/scaleway/test_scaleway_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/scaleway/test_scaleway_audio_transcription_transformation.py rename to tests/unit/llms/scaleway/test_scaleway_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py b/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py similarity index 100% rename from tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py rename to tests/unit/llms/snowflake/test_snowflake_native_endpoints.py diff --git a/tests/test_litellm/llms/soniox/test_soniox_provider_registration.py b/tests/unit/llms/soniox/test_soniox_provider_registration.py similarity index 100% rename from tests/test_litellm/llms/soniox/test_soniox_provider_registration.py rename to tests/unit/llms/soniox/test_soniox_provider_registration.py diff --git a/tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py b/tests/unit/llms/stability/image_generation/test_stability_image_generation.py similarity index 93% rename from tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py rename to tests/unit/llms/stability/image_generation/test_stability_image_generation.py index c5b3c8fbdc5..c5a78603f9c 100644 --- a/tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py +++ b/tests/unit/llms/stability/image_generation/test_stability_image_generation.py @@ -10,10 +10,7 @@ from unittest.mock import MagicMock import httpx import pytest -from litellm.llms.stability.image_generation import ( - StabilityImageGenerationConfig, - get_stability_image_generation_config, -) +from litellm.llms.stability.image_generation import StabilityImageGenerationConfig from litellm.types.llms.stability import ( OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO, STABILITY_GENERATION_MODELS, @@ -266,20 +263,6 @@ class TestStabilityImageGenerationConfig: assert "filtered" in str(exc_info.value).lower() -class TestFactoryFunction: - """Test the factory function""" - - def test_get_stability_image_generation_config(self): - """Test that factory returns correct config type""" - config = get_stability_image_generation_config("stability/sd3") - assert isinstance(config, StabilityImageGenerationConfig) - - def test_factory_returns_config_for_any_model(self): - """Test that factory works for any model name""" - config = get_stability_image_generation_config("stability/custom-model") - assert isinstance(config, StabilityImageGenerationConfig) - - class TestOpenAISizeMapping: """Test the size to aspect ratio mapping""" diff --git a/tests/test_litellm/llms/tencent/chat/test_tencent_chat_transformation.py b/tests/unit/llms/tencent/chat/test_tencent_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/tencent/chat/test_tencent_chat_transformation.py rename to tests/unit/llms/tencent/chat/test_tencent_chat_transformation.py From a3dd47ea11b420018b5082e431b1f0ea611a6fad Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:05:03 +0000 Subject: [PATCH 31/47] test(unit): clear ambient Azure credentials in entra token tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 14 ++++++++++++++ .../azure/realtime/test_azure_realtime_handler.py | 2 +- .../test_azure_ai_image_edit_transformation.py | 4 ++-- .../test_mai_image_edit_transformation.py | 2 +- .../test_azure_ai_passthrough_transformation.py | 4 ++-- .../rerank/test_azure_ai_rerank_transformation.py | 2 +- 6 files changed, 21 insertions(+), 7 deletions(-) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 6b6f7c43760..b3bb19a8b8a 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -12,6 +12,14 @@ import litellm.router as litellm_router_module # noqa: E402 # same import-time import litellm.utils as litellm_utils_module # noqa: E402 # same import-time dependency LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] +AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = ( + "AZURE_AD_TOKEN", + "AZURE_TENANT_ID", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_USERNAME", + "AZURE_PASSWORD", +) def _allow_loopback_only() -> None: @@ -53,5 +61,11 @@ def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: litellm.get_model_info.cache_clear() +@pytest.fixture +def no_ambient_azure_credentials(monkeypatch: pytest.MonkeyPatch) -> None: + for name in AMBIENT_AZURE_CREDENTIAL_ENV_VARS: + monkeypatch.delenv(name, raising=False) + + def pytest_sessionfinish() -> None: enable_socket() diff --git a/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py index c1ba286f8c0..73f43ec8d8a 100644 --- a/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py +++ b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py @@ -707,7 +707,7 @@ async def test_realtime_health_check_uses_bearer_token_when_no_api_key(monkeypat @pytest.mark.asyncio -async def test_arealtime_forwards_deployment_azure_ad_token(monkeypatch): +async def test_arealtime_forwards_deployment_azure_ad_token(monkeypatch, no_ambient_azure_credentials): """ The router binds a deployment's `azure_ad_token` to `_arealtime`'s named parameter rather than **kwargs, so it must still reach the handler. diff --git a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index 39001c1795b..51ba2c34cd7 100644 --- a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py @@ -41,7 +41,7 @@ def test_azure_ai_url_generation(): assert complete_url == expected_url -def test_azure_ai_validate_environment_with_entra_token(monkeypatch): +def test_azure_ai_validate_environment_with_entra_token(monkeypatch, no_ambient_azure_credentials): monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) monkeypatch.setattr(litellm, "api_key", None) config = AzureFoundryFluxImageEditConfig() @@ -55,7 +55,7 @@ def test_azure_ai_validate_environment_with_entra_token(monkeypatch): assert headers == {"Authorization": "Bearer entra-token"} -def test_flux2_validate_environment_with_entra_token(monkeypatch): +def test_flux2_validate_environment_with_entra_token(monkeypatch, no_ambient_azure_credentials): monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) monkeypatch.setattr(litellm, "api_key", None) config = AzureFoundryFlux2ImageEditConfig() diff --git a/tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py index 75e046825a3..2d6f0083194 100644 --- a/tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py +++ b/tests/unit/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py @@ -174,7 +174,7 @@ class TestAzureMAIImageEdit: assert image_response.usage.total_tokens == 1024 -def test_mai_validate_environment_with_entra_token(monkeypatch): +def test_mai_validate_environment_with_entra_token(monkeypatch, no_ambient_azure_credentials): monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) monkeypatch.setattr(litellm, "api_key", None) diff --git a/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index f00698a6624..f9fd9681db8 100644 --- a/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -256,13 +256,13 @@ def test_serverless_host_gets_a_bearer_token(): assert "api-key" not in headers -def test_entra_token_is_used_when_the_deployment_has_no_api_key(): +def test_entra_token_is_used_when_the_deployment_has_no_api_key(no_ambient_azure_credentials): headers = _auth_headers(api_key=None, api_base=FOUNDRY_BASE, litellm_params={"azure_ad_token": "entra-token"}) assert headers["Authorization"] == "Bearer entra-token" -def test_no_credentials_at_all_raises(): +def test_no_credentials_at_all_raises(no_ambient_azure_credentials): with pytest.raises(ValueError, match="Missing Azure AI credentials"): _auth_headers(api_key=None, api_base=FOUNDRY_BASE) diff --git a/tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py b/tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py index 91bf665f18d..3de27199e2e 100644 --- a/tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py +++ b/tests/unit/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py @@ -105,7 +105,7 @@ class TestAzureAIRerankConfigValidateEnvironment: assert headers["Authorization"] == "Bearer my-key" - def test_falls_back_to_entra_token(self, monkeypatch): + def test_falls_back_to_entra_token(self, monkeypatch, no_ambient_azure_credentials): monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) monkeypatch.setattr(litellm, "azure_key", None) From e0b92b2b257bd26650a5887343cfc0de3127129d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 11:07:38 +0000 Subject: [PATCH 32/47] refactor(types): keep a2a response_dict annotation as it was Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/agent_endpoints/a2a_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index faf3e98a3a7..834c16ba6dc 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -880,7 +880,7 @@ async def invoke_agent_a2a( logging_obj._enqueue_deferred_logging = None _enqueue_fn() - response_dict: Final[dict[str, object]] = ( + response_dict: Final[dict[str, Any]] = ( response.model_dump(mode="json", exclude_none=True) if hasattr(response, "model_dump") else response From 6ba4c2e3399bc16a9996c379c020d1b726f8247a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 11:40:11 +0000 Subject: [PATCH 33/47] refactor(types): keep guardrail metadata helper accepting dicts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/base_llm/guardrail_translation/base_translation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index ace5af8124f..89ad67f0485 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -81,7 +81,7 @@ class BaseTranslation(ABC): @staticmethod def transform_user_api_key_dict_to_metadata( - user_api_key_dict: Optional["UserAPIKeyAuth"], + user_api_key_dict: Any | None, ) -> dict[str, object]: """ Transform user_api_key_dict to a metadata dict with prefixed keys. From e4a58ef91acdebe2a25c66f7dbdf7d4bdc338fd7 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:50:59 +0000 Subject: [PATCH 34/47] test(unit): make every tests/unit directory a package so pytest collection is unique Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/__init__.py | 0 tests/unit/integrations/__init__.py | 0 tests/unit/integrations/levo/__init__.py | 0 tests/unit/integrations/litellm_agent/__init__.py | 0 tests/unit/integrations/mavvrik_focus/__init__.py | 0 tests/unit/integrations/opik/__init__.py | 0 tests/unit/integrations/pointfive/__init__.py | 0 .../vector_store_integrations/__init__.py | 0 tests/unit/litellm_core_utils/__init__.py | 0 .../unit/litellm_core_utils/audio_utils/__init__.py | 0 .../llm_response_utils/__init__.py | 0 tests/unit/llms/__init__.py | 0 tests/unit/llms/a2a/__init__.py | 0 tests/unit/llms/a2a/chat/__init__.py | 0 .../llms/a2a/chat/guardrail_translation/__init__.py | 0 tests/unit/llms/anthropic/__init__.py | 0 tests/unit/llms/anthropic/batches/__init__.py | 0 tests/unit/llms/base_llm/__init__.py | 0 tests/unit/llms/base_llm/batches/__init__.py | 0 tests/unit/llms/base_llm/realtime/__init__.py | 0 tests/unit/llms/baseten/__init__.py | 0 tests/unit/llms/baseten/chat/__init__.py | 0 tests/unit/llms/bedrock/__init__.py | 0 tests/unit/llms/bedrock/chat/__init__.py | 0 tests/unit/llms/bedrock/chat/agentcore/__init__.py | 0 .../bedrock/chat/invoke_transformations/__init__.py | 0 tests/unit/llms/bedrock/chat/mantle/__init__.py | 0 tests/unit/llms/bedrock/count_tokens/__init__.py | 0 tests/unit/llms/bedrock/files/__init__.py | 0 tests/unit/llms/bedrock/image/__init__.py | 0 tests/unit/llms/bedrock/image_edit/__init__.py | 0 tests/unit/llms/bedrock/invoke_agent/__init__.py | 0 tests/unit/llms/bedrock/passthrough/__init__.py | 0 .../passthrough/guardrail_translation/__init__.py | 0 tests/unit/llms/bedrock/realtime/__init__.py | 0 tests/unit/llms/bedrock/rerank/__init__.py | 0 tests/unit/llms/bedrock/vector_stores/__init__.py | 0 tests/unit/llms/bedrock_mantle/__init__.py | 0 .../unit/llms/bedrock_mantle/passthrough/__init__.py | 0 tests/unit/llms/black_forest_labs/__init__.py | 0 .../llms/black_forest_labs/image_edit/__init__.py | 0 .../black_forest_labs/image_generation/__init__.py | 0 tests/unit/llms/bytez/__init__.py | 0 tests/unit/llms/bytez/chat/__init__.py | 0 tests/unit/llms/cerebras/__init__.py | 0 tests/unit/llms/chat/__init__.py | 0 tests/unit/llms/chatgpt/__init__.py | 0 tests/unit/llms/chatgpt/chat/__init__.py | 0 tests/unit/llms/chatgpt/responses/__init__.py | 0 tests/unit/llms/cloudflare/__init__.py | 0 tests/unit/llms/cohere/__init__.py | 0 tests/unit/llms/cohere/chat/__init__.py | 0 tests/unit/llms/cohere/embed/__init__.py | 0 tests/unit/llms/cohere/ocr/__init__.py | 0 tests/unit/llms/cohere/rerank/__init__.py | 0 tests/unit/llms/crusoe/__init__.py | 0 tests/unit/llms/databricks/__init__.py | 0 tests/unit/llms/databricks/chat/__init__.py | 0 tests/unit/llms/databricks/responses/__init__.py | 0 tests/unit/llms/datarobot/__init__.py | 0 tests/unit/llms/datarobot/chat/__init__.py | 0 tests/unit/llms/deepseek/__init__.py | 0 tests/unit/llms/deepseek/chat/__init__.py | 0 tests/unit/llms/deepseek/messages/__init__.py | 0 tests/unit/llms/docker_model_runner/__init__.py | 0 tests/unit/llms/elevenlabs/__init__.py | 0 tests/unit/llms/fastcrw/__init__.py | 0 tests/unit/llms/fastcrw/search/__init__.py | 0 tests/unit/llms/fireworks_ai/__init__.py | 0 tests/unit/llms/fireworks_ai/chat/__init__.py | 0 tests/unit/llms/fireworks_ai/rerank/__init__.py | 0 tests/unit/llms/fireworks_ai/responses/__init__.py | 0 tests/unit/llms/gemini/__init__.py | 0 .../unit/llms/gemini/audio_transcription/__init__.py | 0 tests/unit/llms/gemini/files/__init__.py | 0 tests/unit/llms/gemini/google_genai/__init__.py | 0 .../google_genai/guardrail_translation/__init__.py | 0 tests/unit/llms/gemini/image_edit/__init__.py | 0 tests/unit/llms/gemini/realtime/__init__.py | 0 tests/unit/llms/gemini/videos/__init__.py | 0 tests/unit/llms/gigachat/__init__.py | 0 tests/unit/llms/gigachat/chat/__init__.py | 0 tests/unit/llms/gigachat/embedding/__init__.py | 0 tests/unit/llms/gigachat/passthrough/__init__.py | 0 tests/unit/llms/github_copilot/__init__.py | 0 tests/unit/llms/github_copilot/embedding/__init__.py | 0 tests/unit/llms/github_copilot/messages/__init__.py | 0 tests/unit/llms/github_copilot/responses/__init__.py | 0 tests/unit/llms/gradient_ai/__init__.py | 0 tests/unit/llms/gradient_ai/chat/__init__.py | 0 tests/unit/llms/groq/__init__.py | 0 tests/unit/llms/groq/chat/__init__.py | 0 tests/unit/llms/hosted_vllm/__init__.py | 0 tests/unit/llms/hosted_vllm/chat/__init__.py | 0 tests/unit/llms/hosted_vllm/embedding/__init__.py | 0 tests/unit/llms/hosted_vllm/image_edit/__init__.py | 0 tests/unit/llms/hosted_vllm/responses/__init__.py | 0 tests/unit/llms/hosted_vllm/videos/__init__.py | 0 tests/unit/llms/huggingface/__init__.py | 0 tests/unit/llms/huggingface/rerank/__init__.py | 0 tests/unit/llms/inception/__init__.py | 0 tests/unit/llms/jina_ai/__init__.py | 0 tests/unit/llms/jina_ai/embedding/__init__.py | 0 tests/unit/llms/langflow/__init__.py | 0 tests/unit/llms/langflow/chat/__init__.py | 0 tests/unit/llms/litellm_proxy/__init__.py | 0 tests/unit/llms/litellm_proxy/chat/__init__.py | 0 tests/unit/llms/litellm_proxy/skills/__init__.py | 0 tests/unit/llms/llamafile/__init__.py | 0 tests/unit/llms/llamafile/chat/__init__.py | 0 tests/unit/llms/meta/__init__.py | 0 tests/unit/llms/meta/realtime/__init__.py | 0 tests/unit/llms/meta_llama/__init__.py | 0 tests/unit/llms/mistral/audio_speech/__init__.py | 0 tests/unit/llms/modelscope/__init__.py | 0 .../llms/modelscope/image_generation/__init__.py | 0 tests/unit/llms/mongodb/__init__.py | 0 tests/unit/llms/mongodb/vector_stores/__init__.py | 0 tests/unit/llms/moonshot/__init__.py | 0 tests/unit/llms/neosantara/__init__.py | 0 tests/unit/llms/nimble/__init__.py | 0 tests/unit/llms/nimble/search/__init__.py | 0 tests/unit/llms/novita/__init__.py | 0 tests/unit/llms/novita/chat/__init__.py | 0 tests/unit/llms/nscale/__init__.py | 0 tests/unit/llms/nscale/chat/__init__.py | 0 tests/unit/llms/nvidia_nim/__init__.py | 0 tests/unit/llms/nvidia_nim/passthrough/__init__.py | 0 tests/unit/llms/nvidia_nim/rerank/__init__.py | 0 tests/unit/llms/nvidia_riva/__init__.py | 0 .../llms/nvidia_riva/audio_transcription/__init__.py | 0 tests/unit/llms/oci/__init__.py | 0 tests/unit/llms/oci/chat/__init__.py | 0 tests/unit/llms/oci/embed/__init__.py | 0 tests/unit/llms/ocr/__init__.py | 0 .../unit/llms/ocr/guardrail_translation/__init__.py | 0 tests/unit/llms/oobabooga/__init__.py | 0 tests/unit/llms/oobabooga/chat/__init__.py | 0 tests/unit/llms/openai/__init__.py | 0 tests/unit/llms/openai/chat/__init__.py | 0 .../openai/chat/guardrail_translation/__init__.py | 0 tests/unit/llms/openai/completion/__init__.py | 0 tests/unit/test_package_layout.py | 12 ++++++++++++ 143 files changed, 12 insertions(+) create mode 100644 tests/unit/__init__.py create mode 100644 tests/unit/integrations/__init__.py create mode 100644 tests/unit/integrations/levo/__init__.py create mode 100644 tests/unit/integrations/litellm_agent/__init__.py create mode 100644 tests/unit/integrations/mavvrik_focus/__init__.py create mode 100644 tests/unit/integrations/opik/__init__.py create mode 100644 tests/unit/integrations/pointfive/__init__.py create mode 100644 tests/unit/integrations/vector_store_integrations/__init__.py create mode 100644 tests/unit/litellm_core_utils/__init__.py create mode 100644 tests/unit/litellm_core_utils/audio_utils/__init__.py create mode 100644 tests/unit/litellm_core_utils/llm_response_utils/__init__.py create mode 100644 tests/unit/llms/__init__.py create mode 100644 tests/unit/llms/a2a/__init__.py create mode 100644 tests/unit/llms/a2a/chat/__init__.py create mode 100644 tests/unit/llms/a2a/chat/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/anthropic/__init__.py create mode 100644 tests/unit/llms/anthropic/batches/__init__.py create mode 100644 tests/unit/llms/base_llm/__init__.py create mode 100644 tests/unit/llms/base_llm/batches/__init__.py create mode 100644 tests/unit/llms/base_llm/realtime/__init__.py create mode 100644 tests/unit/llms/baseten/__init__.py create mode 100644 tests/unit/llms/baseten/chat/__init__.py create mode 100644 tests/unit/llms/bedrock/__init__.py create mode 100644 tests/unit/llms/bedrock/chat/__init__.py create mode 100644 tests/unit/llms/bedrock/chat/agentcore/__init__.py create mode 100644 tests/unit/llms/bedrock/chat/invoke_transformations/__init__.py create mode 100644 tests/unit/llms/bedrock/chat/mantle/__init__.py create mode 100644 tests/unit/llms/bedrock/count_tokens/__init__.py create mode 100644 tests/unit/llms/bedrock/files/__init__.py create mode 100644 tests/unit/llms/bedrock/image/__init__.py create mode 100644 tests/unit/llms/bedrock/image_edit/__init__.py create mode 100644 tests/unit/llms/bedrock/invoke_agent/__init__.py create mode 100644 tests/unit/llms/bedrock/passthrough/__init__.py create mode 100644 tests/unit/llms/bedrock/passthrough/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/bedrock/realtime/__init__.py create mode 100644 tests/unit/llms/bedrock/rerank/__init__.py create mode 100644 tests/unit/llms/bedrock/vector_stores/__init__.py create mode 100644 tests/unit/llms/bedrock_mantle/__init__.py create mode 100644 tests/unit/llms/bedrock_mantle/passthrough/__init__.py create mode 100644 tests/unit/llms/black_forest_labs/__init__.py create mode 100644 tests/unit/llms/black_forest_labs/image_edit/__init__.py create mode 100644 tests/unit/llms/black_forest_labs/image_generation/__init__.py create mode 100644 tests/unit/llms/bytez/__init__.py create mode 100644 tests/unit/llms/bytez/chat/__init__.py create mode 100644 tests/unit/llms/cerebras/__init__.py create mode 100644 tests/unit/llms/chat/__init__.py create mode 100644 tests/unit/llms/chatgpt/__init__.py create mode 100644 tests/unit/llms/chatgpt/chat/__init__.py create mode 100644 tests/unit/llms/chatgpt/responses/__init__.py create mode 100644 tests/unit/llms/cloudflare/__init__.py create mode 100644 tests/unit/llms/cohere/__init__.py create mode 100644 tests/unit/llms/cohere/chat/__init__.py create mode 100644 tests/unit/llms/cohere/embed/__init__.py create mode 100644 tests/unit/llms/cohere/ocr/__init__.py create mode 100644 tests/unit/llms/cohere/rerank/__init__.py create mode 100644 tests/unit/llms/crusoe/__init__.py create mode 100644 tests/unit/llms/databricks/__init__.py create mode 100644 tests/unit/llms/databricks/chat/__init__.py create mode 100644 tests/unit/llms/databricks/responses/__init__.py create mode 100644 tests/unit/llms/datarobot/__init__.py create mode 100644 tests/unit/llms/datarobot/chat/__init__.py create mode 100644 tests/unit/llms/deepseek/__init__.py create mode 100644 tests/unit/llms/deepseek/chat/__init__.py create mode 100644 tests/unit/llms/deepseek/messages/__init__.py create mode 100644 tests/unit/llms/docker_model_runner/__init__.py create mode 100644 tests/unit/llms/elevenlabs/__init__.py create mode 100644 tests/unit/llms/fastcrw/__init__.py create mode 100644 tests/unit/llms/fastcrw/search/__init__.py create mode 100644 tests/unit/llms/fireworks_ai/__init__.py create mode 100644 tests/unit/llms/fireworks_ai/chat/__init__.py create mode 100644 tests/unit/llms/fireworks_ai/rerank/__init__.py create mode 100644 tests/unit/llms/fireworks_ai/responses/__init__.py create mode 100644 tests/unit/llms/gemini/__init__.py create mode 100644 tests/unit/llms/gemini/audio_transcription/__init__.py create mode 100644 tests/unit/llms/gemini/files/__init__.py create mode 100644 tests/unit/llms/gemini/google_genai/__init__.py create mode 100644 tests/unit/llms/gemini/google_genai/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/gemini/image_edit/__init__.py create mode 100644 tests/unit/llms/gemini/realtime/__init__.py create mode 100644 tests/unit/llms/gemini/videos/__init__.py create mode 100644 tests/unit/llms/gigachat/__init__.py create mode 100644 tests/unit/llms/gigachat/chat/__init__.py create mode 100644 tests/unit/llms/gigachat/embedding/__init__.py create mode 100644 tests/unit/llms/gigachat/passthrough/__init__.py create mode 100644 tests/unit/llms/github_copilot/__init__.py create mode 100644 tests/unit/llms/github_copilot/embedding/__init__.py create mode 100644 tests/unit/llms/github_copilot/messages/__init__.py create mode 100644 tests/unit/llms/github_copilot/responses/__init__.py create mode 100644 tests/unit/llms/gradient_ai/__init__.py create mode 100644 tests/unit/llms/gradient_ai/chat/__init__.py create mode 100644 tests/unit/llms/groq/__init__.py create mode 100644 tests/unit/llms/groq/chat/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/chat/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/embedding/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/image_edit/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/responses/__init__.py create mode 100644 tests/unit/llms/hosted_vllm/videos/__init__.py create mode 100644 tests/unit/llms/huggingface/__init__.py create mode 100644 tests/unit/llms/huggingface/rerank/__init__.py create mode 100644 tests/unit/llms/inception/__init__.py create mode 100644 tests/unit/llms/jina_ai/__init__.py create mode 100644 tests/unit/llms/jina_ai/embedding/__init__.py create mode 100644 tests/unit/llms/langflow/__init__.py create mode 100644 tests/unit/llms/langflow/chat/__init__.py create mode 100644 tests/unit/llms/litellm_proxy/__init__.py create mode 100644 tests/unit/llms/litellm_proxy/chat/__init__.py create mode 100644 tests/unit/llms/litellm_proxy/skills/__init__.py create mode 100644 tests/unit/llms/llamafile/__init__.py create mode 100644 tests/unit/llms/llamafile/chat/__init__.py create mode 100644 tests/unit/llms/meta/__init__.py create mode 100644 tests/unit/llms/meta/realtime/__init__.py create mode 100644 tests/unit/llms/meta_llama/__init__.py create mode 100644 tests/unit/llms/mistral/audio_speech/__init__.py create mode 100644 tests/unit/llms/modelscope/__init__.py create mode 100644 tests/unit/llms/modelscope/image_generation/__init__.py create mode 100644 tests/unit/llms/mongodb/__init__.py create mode 100644 tests/unit/llms/mongodb/vector_stores/__init__.py create mode 100644 tests/unit/llms/moonshot/__init__.py create mode 100644 tests/unit/llms/neosantara/__init__.py create mode 100644 tests/unit/llms/nimble/__init__.py create mode 100644 tests/unit/llms/nimble/search/__init__.py create mode 100644 tests/unit/llms/novita/__init__.py create mode 100644 tests/unit/llms/novita/chat/__init__.py create mode 100644 tests/unit/llms/nscale/__init__.py create mode 100644 tests/unit/llms/nscale/chat/__init__.py create mode 100644 tests/unit/llms/nvidia_nim/__init__.py create mode 100644 tests/unit/llms/nvidia_nim/passthrough/__init__.py create mode 100644 tests/unit/llms/nvidia_nim/rerank/__init__.py create mode 100644 tests/unit/llms/nvidia_riva/__init__.py create mode 100644 tests/unit/llms/nvidia_riva/audio_transcription/__init__.py create mode 100644 tests/unit/llms/oci/__init__.py create mode 100644 tests/unit/llms/oci/chat/__init__.py create mode 100644 tests/unit/llms/oci/embed/__init__.py create mode 100644 tests/unit/llms/ocr/__init__.py create mode 100644 tests/unit/llms/ocr/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/oobabooga/__init__.py create mode 100644 tests/unit/llms/oobabooga/chat/__init__.py create mode 100644 tests/unit/llms/openai/__init__.py create mode 100644 tests/unit/llms/openai/chat/__init__.py create mode 100644 tests/unit/llms/openai/chat/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/openai/completion/__init__.py create mode 100644 tests/unit/test_package_layout.py diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/__init__.py b/tests/unit/integrations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/levo/__init__.py b/tests/unit/integrations/levo/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/litellm_agent/__init__.py b/tests/unit/integrations/litellm_agent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/mavvrik_focus/__init__.py b/tests/unit/integrations/mavvrik_focus/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/opik/__init__.py b/tests/unit/integrations/opik/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/pointfive/__init__.py b/tests/unit/integrations/pointfive/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/vector_store_integrations/__init__.py b/tests/unit/integrations/vector_store_integrations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/litellm_core_utils/__init__.py b/tests/unit/litellm_core_utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/litellm_core_utils/audio_utils/__init__.py b/tests/unit/litellm_core_utils/audio_utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/litellm_core_utils/llm_response_utils/__init__.py b/tests/unit/litellm_core_utils/llm_response_utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/__init__.py b/tests/unit/llms/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/a2a/__init__.py b/tests/unit/llms/a2a/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/a2a/chat/__init__.py b/tests/unit/llms/a2a/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/a2a/chat/guardrail_translation/__init__.py b/tests/unit/llms/a2a/chat/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/__init__.py b/tests/unit/llms/anthropic/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/batches/__init__.py b/tests/unit/llms/anthropic/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/base_llm/__init__.py b/tests/unit/llms/base_llm/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/base_llm/batches/__init__.py b/tests/unit/llms/base_llm/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/base_llm/realtime/__init__.py b/tests/unit/llms/base_llm/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/baseten/__init__.py b/tests/unit/llms/baseten/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/baseten/chat/__init__.py b/tests/unit/llms/baseten/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/__init__.py b/tests/unit/llms/bedrock/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/chat/__init__.py b/tests/unit/llms/bedrock/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/chat/agentcore/__init__.py b/tests/unit/llms/bedrock/chat/agentcore/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/__init__.py b/tests/unit/llms/bedrock/chat/invoke_transformations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/chat/mantle/__init__.py b/tests/unit/llms/bedrock/chat/mantle/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/count_tokens/__init__.py b/tests/unit/llms/bedrock/count_tokens/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/files/__init__.py b/tests/unit/llms/bedrock/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/image/__init__.py b/tests/unit/llms/bedrock/image/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/image_edit/__init__.py b/tests/unit/llms/bedrock/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/invoke_agent/__init__.py b/tests/unit/llms/bedrock/invoke_agent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/passthrough/__init__.py b/tests/unit/llms/bedrock/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/__init__.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/realtime/__init__.py b/tests/unit/llms/bedrock/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/rerank/__init__.py b/tests/unit/llms/bedrock/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/vector_stores/__init__.py b/tests/unit/llms/bedrock/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock_mantle/__init__.py b/tests/unit/llms/bedrock_mantle/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock_mantle/passthrough/__init__.py b/tests/unit/llms/bedrock_mantle/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/black_forest_labs/__init__.py b/tests/unit/llms/black_forest_labs/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/black_forest_labs/image_edit/__init__.py b/tests/unit/llms/black_forest_labs/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/black_forest_labs/image_generation/__init__.py b/tests/unit/llms/black_forest_labs/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bytez/__init__.py b/tests/unit/llms/bytez/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bytez/chat/__init__.py b/tests/unit/llms/bytez/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cerebras/__init__.py b/tests/unit/llms/cerebras/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/chat/__init__.py b/tests/unit/llms/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/chatgpt/__init__.py b/tests/unit/llms/chatgpt/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/chatgpt/chat/__init__.py b/tests/unit/llms/chatgpt/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/chatgpt/responses/__init__.py b/tests/unit/llms/chatgpt/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cloudflare/__init__.py b/tests/unit/llms/cloudflare/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cohere/__init__.py b/tests/unit/llms/cohere/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cohere/chat/__init__.py b/tests/unit/llms/cohere/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cohere/embed/__init__.py b/tests/unit/llms/cohere/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cohere/ocr/__init__.py b/tests/unit/llms/cohere/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/cohere/rerank/__init__.py b/tests/unit/llms/cohere/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/crusoe/__init__.py b/tests/unit/llms/crusoe/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/databricks/__init__.py b/tests/unit/llms/databricks/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/databricks/chat/__init__.py b/tests/unit/llms/databricks/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/databricks/responses/__init__.py b/tests/unit/llms/databricks/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/datarobot/__init__.py b/tests/unit/llms/datarobot/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/datarobot/chat/__init__.py b/tests/unit/llms/datarobot/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/deepseek/__init__.py b/tests/unit/llms/deepseek/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/deepseek/chat/__init__.py b/tests/unit/llms/deepseek/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/deepseek/messages/__init__.py b/tests/unit/llms/deepseek/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/docker_model_runner/__init__.py b/tests/unit/llms/docker_model_runner/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/elevenlabs/__init__.py b/tests/unit/llms/elevenlabs/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fastcrw/__init__.py b/tests/unit/llms/fastcrw/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fastcrw/search/__init__.py b/tests/unit/llms/fastcrw/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fireworks_ai/__init__.py b/tests/unit/llms/fireworks_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fireworks_ai/chat/__init__.py b/tests/unit/llms/fireworks_ai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fireworks_ai/rerank/__init__.py b/tests/unit/llms/fireworks_ai/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/fireworks_ai/responses/__init__.py b/tests/unit/llms/fireworks_ai/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/__init__.py b/tests/unit/llms/gemini/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/audio_transcription/__init__.py b/tests/unit/llms/gemini/audio_transcription/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/files/__init__.py b/tests/unit/llms/gemini/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/google_genai/__init__.py b/tests/unit/llms/gemini/google_genai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/google_genai/guardrail_translation/__init__.py b/tests/unit/llms/gemini/google_genai/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/image_edit/__init__.py b/tests/unit/llms/gemini/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/realtime/__init__.py b/tests/unit/llms/gemini/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/videos/__init__.py b/tests/unit/llms/gemini/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gigachat/__init__.py b/tests/unit/llms/gigachat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gigachat/chat/__init__.py b/tests/unit/llms/gigachat/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gigachat/embedding/__init__.py b/tests/unit/llms/gigachat/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gigachat/passthrough/__init__.py b/tests/unit/llms/gigachat/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/github_copilot/__init__.py b/tests/unit/llms/github_copilot/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/github_copilot/embedding/__init__.py b/tests/unit/llms/github_copilot/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/github_copilot/messages/__init__.py b/tests/unit/llms/github_copilot/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/github_copilot/responses/__init__.py b/tests/unit/llms/github_copilot/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gradient_ai/__init__.py b/tests/unit/llms/gradient_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gradient_ai/chat/__init__.py b/tests/unit/llms/gradient_ai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/groq/__init__.py b/tests/unit/llms/groq/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/groq/chat/__init__.py b/tests/unit/llms/groq/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/__init__.py b/tests/unit/llms/hosted_vllm/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/chat/__init__.py b/tests/unit/llms/hosted_vllm/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/embedding/__init__.py b/tests/unit/llms/hosted_vllm/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/image_edit/__init__.py b/tests/unit/llms/hosted_vllm/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/responses/__init__.py b/tests/unit/llms/hosted_vllm/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/hosted_vllm/videos/__init__.py b/tests/unit/llms/hosted_vllm/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/huggingface/__init__.py b/tests/unit/llms/huggingface/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/huggingface/rerank/__init__.py b/tests/unit/llms/huggingface/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/inception/__init__.py b/tests/unit/llms/inception/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/jina_ai/__init__.py b/tests/unit/llms/jina_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/jina_ai/embedding/__init__.py b/tests/unit/llms/jina_ai/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/langflow/__init__.py b/tests/unit/llms/langflow/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/langflow/chat/__init__.py b/tests/unit/llms/langflow/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/litellm_proxy/__init__.py b/tests/unit/llms/litellm_proxy/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/litellm_proxy/chat/__init__.py b/tests/unit/llms/litellm_proxy/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/litellm_proxy/skills/__init__.py b/tests/unit/llms/litellm_proxy/skills/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/llamafile/__init__.py b/tests/unit/llms/llamafile/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/llamafile/chat/__init__.py b/tests/unit/llms/llamafile/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/meta/__init__.py b/tests/unit/llms/meta/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/meta/realtime/__init__.py b/tests/unit/llms/meta/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/meta_llama/__init__.py b/tests/unit/llms/meta_llama/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/mistral/audio_speech/__init__.py b/tests/unit/llms/mistral/audio_speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/modelscope/__init__.py b/tests/unit/llms/modelscope/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/modelscope/image_generation/__init__.py b/tests/unit/llms/modelscope/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/mongodb/__init__.py b/tests/unit/llms/mongodb/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/mongodb/vector_stores/__init__.py b/tests/unit/llms/mongodb/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/moonshot/__init__.py b/tests/unit/llms/moonshot/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/neosantara/__init__.py b/tests/unit/llms/neosantara/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nimble/__init__.py b/tests/unit/llms/nimble/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nimble/search/__init__.py b/tests/unit/llms/nimble/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/novita/__init__.py b/tests/unit/llms/novita/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/novita/chat/__init__.py b/tests/unit/llms/novita/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nscale/__init__.py b/tests/unit/llms/nscale/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nscale/chat/__init__.py b/tests/unit/llms/nscale/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nvidia_nim/__init__.py b/tests/unit/llms/nvidia_nim/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nvidia_nim/passthrough/__init__.py b/tests/unit/llms/nvidia_nim/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nvidia_nim/rerank/__init__.py b/tests/unit/llms/nvidia_nim/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nvidia_riva/__init__.py b/tests/unit/llms/nvidia_riva/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/nvidia_riva/audio_transcription/__init__.py b/tests/unit/llms/nvidia_riva/audio_transcription/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/oci/__init__.py b/tests/unit/llms/oci/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/oci/chat/__init__.py b/tests/unit/llms/oci/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/oci/embed/__init__.py b/tests/unit/llms/oci/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/ocr/__init__.py b/tests/unit/llms/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/ocr/guardrail_translation/__init__.py b/tests/unit/llms/ocr/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/oobabooga/__init__.py b/tests/unit/llms/oobabooga/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/oobabooga/chat/__init__.py b/tests/unit/llms/oobabooga/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/__init__.py b/tests/unit/llms/openai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/chat/__init__.py b/tests/unit/llms/openai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/chat/guardrail_translation/__init__.py b/tests/unit/llms/openai/chat/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/completion/__init__.py b/tests/unit/llms/openai/completion/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/test_package_layout.py b/tests/unit/test_package_layout.py new file mode 100644 index 00000000000..4ea68fc06ba --- /dev/null +++ b/tests/unit/test_package_layout.py @@ -0,0 +1,12 @@ +import os + +TESTS_UNIT_DIR = os.path.dirname(os.path.abspath(__file__)) + + +def test_every_directory_under_tests_unit_is_a_package(): + missing = [] + for root, dirs, _files in os.walk(TESTS_UNIT_DIR): + dirs[:] = [d for d in dirs if d != "__pycache__"] + if not os.path.isfile(os.path.join(root, "__init__.py")): + missing.append(os.path.relpath(root, TESTS_UNIT_DIR)) + assert missing == [] From 5599c59923c5362dccc5c17012ba0c02a2d7254f Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:52:07 +0000 Subject: [PATCH 35/47] test: add package initializers to migrated unit test directories Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/llms/openai/embeddings/__init__.py | 0 .../unit/llms/openai/embeddings/guardrail_translation/__init__.py | 0 tests/unit/llms/openai/evals/__init__.py | 0 tests/unit/llms/openai/image_generation/__init__.py | 0 tests/unit/llms/openai/speech/__init__.py | 0 tests/unit/llms/openai/transcriptions/__init__.py | 0 tests/unit/llms/openai/vector_store_files/__init__.py | 0 tests/unit/llms/openai/vector_stores/__init__.py | 0 tests/unit/llms/openai/videos/__init__.py | 0 tests/unit/llms/openai_like/__init__.py | 0 tests/unit/llms/openai_like/chat/__init__.py | 0 tests/unit/llms/openai_like/embedding/__init__.py | 0 tests/unit/llms/openai_like/messages/__init__.py | 0 tests/unit/llms/openrouter/__init__.py | 0 tests/unit/llms/openrouter/chat/__init__.py | 0 tests/unit/llms/openrouter/image_edit/__init__.py | 0 tests/unit/llms/openrouter/image_generation/__init__.py | 0 17 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/llms/openai/embeddings/__init__.py create mode 100644 tests/unit/llms/openai/embeddings/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/openai/evals/__init__.py create mode 100644 tests/unit/llms/openai/image_generation/__init__.py create mode 100644 tests/unit/llms/openai/speech/__init__.py create mode 100644 tests/unit/llms/openai/transcriptions/__init__.py create mode 100644 tests/unit/llms/openai/vector_store_files/__init__.py create mode 100644 tests/unit/llms/openai/vector_stores/__init__.py create mode 100644 tests/unit/llms/openai/videos/__init__.py create mode 100644 tests/unit/llms/openai_like/__init__.py create mode 100644 tests/unit/llms/openai_like/chat/__init__.py create mode 100644 tests/unit/llms/openai_like/embedding/__init__.py create mode 100644 tests/unit/llms/openai_like/messages/__init__.py create mode 100644 tests/unit/llms/openrouter/__init__.py create mode 100644 tests/unit/llms/openrouter/chat/__init__.py create mode 100644 tests/unit/llms/openrouter/image_edit/__init__.py create mode 100644 tests/unit/llms/openrouter/image_generation/__init__.py diff --git a/tests/unit/llms/openai/embeddings/__init__.py b/tests/unit/llms/openai/embeddings/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/embeddings/guardrail_translation/__init__.py b/tests/unit/llms/openai/embeddings/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/evals/__init__.py b/tests/unit/llms/openai/evals/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/image_generation/__init__.py b/tests/unit/llms/openai/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/speech/__init__.py b/tests/unit/llms/openai/speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/transcriptions/__init__.py b/tests/unit/llms/openai/transcriptions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/vector_store_files/__init__.py b/tests/unit/llms/openai/vector_store_files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/vector_stores/__init__.py b/tests/unit/llms/openai/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/videos/__init__.py b/tests/unit/llms/openai/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai_like/__init__.py b/tests/unit/llms/openai_like/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai_like/chat/__init__.py b/tests/unit/llms/openai_like/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai_like/embedding/__init__.py b/tests/unit/llms/openai_like/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai_like/messages/__init__.py b/tests/unit/llms/openai_like/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openrouter/__init__.py b/tests/unit/llms/openrouter/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openrouter/chat/__init__.py b/tests/unit/llms/openrouter/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openrouter/image_edit/__init__.py b/tests/unit/llms/openrouter/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openrouter/image_generation/__init__.py b/tests/unit/llms/openrouter/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 49cd32affe1cfd597c223cdb62049d52ed55e5b5 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:52:09 +0000 Subject: [PATCH 36/47] test: add __init__.py to every tests/unit directory this migration touches Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/llms/vertex_ai/__init__.py | 0 tests/unit/llms/vertex_ai/vertex_ai_partner_models/__init__.py | 0 .../llms/vertex_ai/vertex_ai_partner_models/llama3/__init__.py | 0 .../llms/vertex_ai/vertex_ai_partner_models/mistral/__init__.py | 0 tests/unit/llms/volcengine/__init__.py | 0 tests/unit/llms/volcengine/responses/__init__.py | 0 tests/unit/llms/voyage/__init__.py | 0 tests/unit/llms/watsonx/embed/__init__.py | 0 tests/unit/llms/watsonx/passthrough/__init__.py | 0 tests/unit/llms/xai/__init__.py | 0 tests/unit/llms/xai/responses/__init__.py | 0 tests/unit/llms/zai/__init__.py | 0 12 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/llms/vertex_ai/__init__.py create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/__init__.py create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/__init__.py create mode 100644 tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/__init__.py create mode 100644 tests/unit/llms/volcengine/__init__.py create mode 100644 tests/unit/llms/volcengine/responses/__init__.py create mode 100644 tests/unit/llms/voyage/__init__.py create mode 100644 tests/unit/llms/watsonx/embed/__init__.py create mode 100644 tests/unit/llms/watsonx/passthrough/__init__.py create mode 100644 tests/unit/llms/xai/__init__.py create mode 100644 tests/unit/llms/xai/responses/__init__.py create mode 100644 tests/unit/llms/zai/__init__.py diff --git a/tests/unit/llms/vertex_ai/__init__.py b/tests/unit/llms/vertex_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/mistral/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/volcengine/__init__.py b/tests/unit/llms/volcengine/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/volcengine/responses/__init__.py b/tests/unit/llms/volcengine/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/voyage/__init__.py b/tests/unit/llms/voyage/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/watsonx/embed/__init__.py b/tests/unit/llms/watsonx/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/watsonx/passthrough/__init__.py b/tests/unit/llms/watsonx/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/xai/__init__.py b/tests/unit/llms/xai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/xai/responses/__init__.py b/tests/unit/llms/xai/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/zai/__init__.py b/tests/unit/llms/zai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 434af084e7572d2b2f23101f56d7c722189aa33d Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:52:38 +0000 Subject: [PATCH 37/47] test: add __init__.py to every tests/unit directory phase 16 touches Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/__init__.py | 0 tests/unit/sandbox/__init__.py | 0 tests/unit/skills/__init__.py | 0 tests/unit/test_router/__init__.py | 0 tests/unit/types/__init__.py | 0 tests/unit/types/llms/__init__.py | 0 tests/unit/types/proxy/__init__.py | 0 7 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/__init__.py create mode 100644 tests/unit/sandbox/__init__.py create mode 100644 tests/unit/skills/__init__.py create mode 100644 tests/unit/test_router/__init__.py create mode 100644 tests/unit/types/__init__.py create mode 100644 tests/unit/types/llms/__init__.py create mode 100644 tests/unit/types/proxy/__init__.py diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/sandbox/__init__.py b/tests/unit/sandbox/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/skills/__init__.py b/tests/unit/skills/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/test_router/__init__.py b/tests/unit/test_router/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/types/__init__.py b/tests/unit/types/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/types/llms/__init__.py b/tests/unit/types/llms/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/types/proxy/__init__.py b/tests/unit/types/proxy/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From baf40ea5e896674dd8e75d6ac8b1937fae07e475 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:53:14 +0000 Subject: [PATCH 38/47] test(unit): add package markers to migrated unit test directories Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/__init__.py | 0 tests/unit/llms/__init__.py | 0 tests/unit/llms/anthropic/__init__.py | 0 tests/unit/llms/anthropic/experimental_pass_through/__init__.py | 0 tests/unit/llms/anthropic/files/__init__.py | 0 tests/unit/llms/anthropic/messages/__init__.py | 0 tests/unit/llms/apiserpent/__init__.py | 0 tests/unit/llms/azure/__init__.py | 0 tests/unit/llms/azure/image_edit/__init__.py | 0 tests/unit/llms/azure/image_generation/__init__.py | 0 tests/unit/llms/azure/passthrough/__init__.py | 0 tests/unit/llms/azure/realtime/__init__.py | 0 tests/unit/llms/azure/response/__init__.py | 0 tests/unit/llms/azure/search/__init__.py | 0 tests/unit/llms/azure/text_to_speech/__init__.py | 0 tests/unit/llms/azure/vector_stores/__init__.py | 0 tests/unit/llms/azure_ai/__init__.py | 0 tests/unit/llms/azure_ai/chat/__init__.py | 0 tests/unit/llms/azure_ai/embed/__init__.py | 0 tests/unit/llms/azure_ai/image_edit/__init__.py | 0 tests/unit/llms/azure_ai/ocr/__init__.py | 0 tests/unit/llms/azure_ai/passthrough/__init__.py | 0 tests/unit/llms/azure_ai/rerank/__init__.py | 0 tests/unit/llms/azure_ai/responses/__init__.py | 0 24 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/__init__.py create mode 100644 tests/unit/llms/__init__.py create mode 100644 tests/unit/llms/anthropic/__init__.py create mode 100644 tests/unit/llms/anthropic/experimental_pass_through/__init__.py create mode 100644 tests/unit/llms/anthropic/files/__init__.py create mode 100644 tests/unit/llms/anthropic/messages/__init__.py create mode 100644 tests/unit/llms/apiserpent/__init__.py create mode 100644 tests/unit/llms/azure/__init__.py create mode 100644 tests/unit/llms/azure/image_edit/__init__.py create mode 100644 tests/unit/llms/azure/image_generation/__init__.py create mode 100644 tests/unit/llms/azure/passthrough/__init__.py create mode 100644 tests/unit/llms/azure/realtime/__init__.py create mode 100644 tests/unit/llms/azure/response/__init__.py create mode 100644 tests/unit/llms/azure/search/__init__.py create mode 100644 tests/unit/llms/azure/text_to_speech/__init__.py create mode 100644 tests/unit/llms/azure/vector_stores/__init__.py create mode 100644 tests/unit/llms/azure_ai/__init__.py create mode 100644 tests/unit/llms/azure_ai/chat/__init__.py create mode 100644 tests/unit/llms/azure_ai/embed/__init__.py create mode 100644 tests/unit/llms/azure_ai/image_edit/__init__.py create mode 100644 tests/unit/llms/azure_ai/ocr/__init__.py create mode 100644 tests/unit/llms/azure_ai/passthrough/__init__.py create mode 100644 tests/unit/llms/azure_ai/rerank/__init__.py create mode 100644 tests/unit/llms/azure_ai/responses/__init__.py diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/__init__.py b/tests/unit/llms/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/__init__.py b/tests/unit/llms/anthropic/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/experimental_pass_through/__init__.py b/tests/unit/llms/anthropic/experimental_pass_through/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/files/__init__.py b/tests/unit/llms/anthropic/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/anthropic/messages/__init__.py b/tests/unit/llms/anthropic/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/apiserpent/__init__.py b/tests/unit/llms/apiserpent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/__init__.py b/tests/unit/llms/azure/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/image_edit/__init__.py b/tests/unit/llms/azure/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/image_generation/__init__.py b/tests/unit/llms/azure/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/passthrough/__init__.py b/tests/unit/llms/azure/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/realtime/__init__.py b/tests/unit/llms/azure/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/response/__init__.py b/tests/unit/llms/azure/response/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/search/__init__.py b/tests/unit/llms/azure/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/text_to_speech/__init__.py b/tests/unit/llms/azure/text_to_speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure/vector_stores/__init__.py b/tests/unit/llms/azure/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/__init__.py b/tests/unit/llms/azure_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/chat/__init__.py b/tests/unit/llms/azure_ai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/embed/__init__.py b/tests/unit/llms/azure_ai/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/image_edit/__init__.py b/tests/unit/llms/azure_ai/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/ocr/__init__.py b/tests/unit/llms/azure_ai/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/passthrough/__init__.py b/tests/unit/llms/azure_ai/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/rerank/__init__.py b/tests/unit/llms/azure_ai/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/responses/__init__.py b/tests/unit/llms/azure_ai/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 30a422087c9bf28cc3fc94a6d68b241b703325a1 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:53:49 +0000 Subject: [PATCH 39/47] test: add __init__.py to migrated tests/unit packages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/a2a_protocol/__init__.py | 0 tests/unit/a2a_protocol/providers/__init__.py | 0 tests/unit/a2a_protocol/providers/pydantic_ai_agents/__init__.py | 0 tests/unit/a2a_protocol/providers/watsonx_orchestrate/__init__.py | 0 tests/unit/anthropic_interface/__init__.py | 0 tests/unit/anthropic_interface/exceptions/__init__.py | 0 tests/unit/batches/__init__.py | 0 tests/unit/chat_completions/__init__.py | 0 tests/unit/completion_extras/__init__.py | 0 tests/unit/compression/__init__.py | 0 tests/unit/endpoints/__init__.py | 0 tests/unit/endpoints/speech/__init__.py | 0 .../unit/endpoints/speech/speech_to_completion_bridge/__init__.py | 0 tests/unit/enterprise/__init__.py | 0 tests/unit/enterprise/enterprise_callbacks/__init__.py | 0 tests/unit/integrations/__init__.py | 0 tests/unit/integrations/compression_interception/__init__.py | 0 tests/unit/integrations/gcs_bucket/__init__.py | 0 tests/unit/integrations/gcs_pubsub/__init__.py | 0 tests/unit/integrations/helicone/__init__.py | 0 20 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/a2a_protocol/__init__.py create mode 100644 tests/unit/a2a_protocol/providers/__init__.py create mode 100644 tests/unit/a2a_protocol/providers/pydantic_ai_agents/__init__.py create mode 100644 tests/unit/a2a_protocol/providers/watsonx_orchestrate/__init__.py create mode 100644 tests/unit/anthropic_interface/__init__.py create mode 100644 tests/unit/anthropic_interface/exceptions/__init__.py create mode 100644 tests/unit/batches/__init__.py create mode 100644 tests/unit/chat_completions/__init__.py create mode 100644 tests/unit/completion_extras/__init__.py create mode 100644 tests/unit/compression/__init__.py create mode 100644 tests/unit/endpoints/__init__.py create mode 100644 tests/unit/endpoints/speech/__init__.py create mode 100644 tests/unit/endpoints/speech/speech_to_completion_bridge/__init__.py create mode 100644 tests/unit/enterprise/__init__.py create mode 100644 tests/unit/enterprise/enterprise_callbacks/__init__.py create mode 100644 tests/unit/integrations/__init__.py create mode 100644 tests/unit/integrations/compression_interception/__init__.py create mode 100644 tests/unit/integrations/gcs_bucket/__init__.py create mode 100644 tests/unit/integrations/gcs_pubsub/__init__.py create mode 100644 tests/unit/integrations/helicone/__init__.py diff --git a/tests/unit/a2a_protocol/__init__.py b/tests/unit/a2a_protocol/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/a2a_protocol/providers/__init__.py b/tests/unit/a2a_protocol/providers/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/a2a_protocol/providers/pydantic_ai_agents/__init__.py b/tests/unit/a2a_protocol/providers/pydantic_ai_agents/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/a2a_protocol/providers/watsonx_orchestrate/__init__.py b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/anthropic_interface/__init__.py b/tests/unit/anthropic_interface/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/anthropic_interface/exceptions/__init__.py b/tests/unit/anthropic_interface/exceptions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/batches/__init__.py b/tests/unit/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/chat_completions/__init__.py b/tests/unit/chat_completions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/completion_extras/__init__.py b/tests/unit/completion_extras/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/compression/__init__.py b/tests/unit/compression/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/endpoints/__init__.py b/tests/unit/endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/endpoints/speech/__init__.py b/tests/unit/endpoints/speech/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/endpoints/speech/speech_to_completion_bridge/__init__.py b/tests/unit/endpoints/speech/speech_to_completion_bridge/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/enterprise/__init__.py b/tests/unit/enterprise/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/enterprise/enterprise_callbacks/__init__.py b/tests/unit/enterprise/enterprise_callbacks/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/__init__.py b/tests/unit/integrations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/compression_interception/__init__.py b/tests/unit/integrations/compression_interception/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/gcs_bucket/__init__.py b/tests/unit/integrations/gcs_bucket/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/gcs_pubsub/__init__.py b/tests/unit/integrations/gcs_pubsub/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/helicone/__init__.py b/tests/unit/integrations/helicone/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 78a751c04972d0079ef01bb87fef58a6e472be65 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:55:43 +0000 Subject: [PATCH 40/47] test: migrate phase 15 legacy tests to tests/unit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/messages/__init__.py | 0 .../messages/test_dispatch.py | 11 +- .../models/test_models.py | 46 ++--- tests/unit/ocr/__init__.py | 0 .../ocr/test_dispatch.py | 0 tests/{test_litellm => unit}/ocr/test_main.py | 0 .../ocr/test_ocr_file_input.py | 48 ++--- tests/unit/passthrough/__init__.py | 0 .../test_async_streaming_error_propagation.py | 27 ++- .../passthrough/test_passthrough_main.py | 77 ++------ ...test_streaming_interrupt_spend_tracking.py | 31 +--- tests/unit/rag/ingestion/__init__.py | 0 .../ingestion/test_s3_vectors_ingestion.py | 12 +- .../realtime_api/test_main.py | 11 +- .../repositories/test_repositories.py | 48 ++--- .../repositories/test_unit_of_work.py | 0 .../complexity_router/test_jev_classifier.py | 0 .../test_deployment_affinity_check.py | 70 ++++---- .../test_encrypted_content_affinity_check.py | 169 ++++++++++-------- .../test_prompt_caching_deployment_check.py | 47 +++-- .../test_responses_api_deployment_check.py | 26 +-- .../test_session_id_affinity.py | 70 ++++---- tests/unit/rust_bridge/__init__.py | 0 .../rust_bridge/chat_completions/__init__.py | 0 .../chat_completions/test_route_host.py | 0 tests/unit/rust_bridge/messages/__init__.py | 0 .../rust_bridge/messages/test_route_host.py | 0 27 files changed, 290 insertions(+), 403 deletions(-) create mode 100644 tests/unit/messages/__init__.py rename tests/{test_litellm => unit}/messages/test_dispatch.py (97%) rename tests/{test_litellm => unit}/models/test_models.py (93%) create mode 100644 tests/unit/ocr/__init__.py rename tests/{test_litellm => unit}/ocr/test_dispatch.py (100%) rename tests/{test_litellm => unit}/ocr/test_main.py (100%) rename tests/{test_litellm => unit}/ocr/test_ocr_file_input.py (92%) create mode 100644 tests/unit/passthrough/__init__.py rename tests/{test_litellm => unit}/passthrough/test_async_streaming_error_propagation.py (92%) rename tests/{test_litellm => unit}/passthrough/test_passthrough_main.py (94%) rename tests/{test_litellm => unit}/passthrough/test_streaming_interrupt_spend_tracking.py (91%) create mode 100644 tests/unit/rag/ingestion/__init__.py rename tests/{test_litellm => unit}/rag/ingestion/test_s3_vectors_ingestion.py (94%) rename tests/{test_litellm => unit}/realtime_api/test_main.py (98%) rename tests/{test_litellm => unit}/repositories/test_repositories.py (98%) rename tests/{test_litellm => unit}/repositories/test_unit_of_work.py (100%) rename tests/{test_litellm => unit}/router_strategy/complexity_router/test_jev_classifier.py (100%) rename tests/{test_litellm => unit}/router_utils/pre_call_checks/test_deployment_affinity_check.py (95%) rename tests/{test_litellm => unit}/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py (96%) rename tests/{test_litellm => unit}/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py (94%) rename tests/{test_litellm => unit}/router_utils/pre_call_checks/test_responses_api_deployment_check.py (95%) rename tests/{test_litellm => unit}/router_utils/pre_call_checks/test_session_id_affinity.py (95%) create mode 100644 tests/unit/rust_bridge/__init__.py create mode 100644 tests/unit/rust_bridge/chat_completions/__init__.py rename tests/{test_litellm => unit}/rust_bridge/chat_completions/test_route_host.py (100%) create mode 100644 tests/unit/rust_bridge/messages/__init__.py rename tests/{test_litellm => unit}/rust_bridge/messages/test_route_host.py (100%) diff --git a/tests/unit/messages/__init__.py b/tests/unit/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/messages/test_dispatch.py b/tests/unit/messages/test_dispatch.py similarity index 97% rename from tests/test_litellm/messages/test_dispatch.py rename to tests/unit/messages/test_dispatch.py index 2eaf4cd9a50..586b77d9a25 100644 --- a/tests/test_litellm/messages/test_dispatch.py +++ b/tests/unit/messages/test_dispatch.py @@ -29,9 +29,7 @@ RUST_RULES: Final[Rules] = (Rule(Route.MESSAGES, Rollout.RUST_REQUIRED),) def messages_binding(native: NativeMessages | None) -> NativeBinding[NativeMessages]: - binding: Final[NativeBinding[NativeMessages]] = NativeBinding( - "anthropic_messages_handler", validate=lambda _: None - ) + binding: Final[NativeBinding[NativeMessages]] = NativeBinding("anthropic_messages_handler", validate=lambda _: None) binding.override(native) return binding @@ -99,7 +97,8 @@ async def test_async_python_route_forwards_original_call_shape() -> None: expected: Final = response() async def python( - *call_args: object, **call_kwargs: object # kwargs-ok: records call shape + *call_args: object, + **call_kwargs: object, # kwargs-ok: records call shape ) -> AnthropicMessagesResponse: captured.append((call_args, call_kwargs)) return expected @@ -217,7 +216,9 @@ def test_binding_errors_delegate_to_python(args: tuple[object, ...], kwargs: Map captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = [] expected: Final = response() - def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: records invalid call + def python( + *call_args: object, **call_kwargs: object + ) -> AnthropicMessagesResponse: # kwargs-ok: records invalid call captured.append((call_args, call_kwargs)) return expected diff --git a/tests/test_litellm/models/test_models.py b/tests/unit/models/test_models.py similarity index 93% rename from tests/test_litellm/models/test_models.py rename to tests/unit/models/test_models.py index 777b4a265ac..b8bf55f1b4a 100644 --- a/tests/test_litellm/models/test_models.py +++ b/tests/unit/models/test_models.py @@ -5,7 +5,7 @@ Tests for backend domain models. from datetime import datetime, timezone import pytest -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, TypeAdapter, ValidationError from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.models.autorouter_session import LiteLLM_AutoRouterSession @@ -19,7 +19,6 @@ from litellm.models.credentials import CreateCredentialItem, CredentialItem from litellm.models.end_user import LiteLLM_EndUserTable from litellm.models.managed_files import ( LiteLLM_ManagedFileTable, - LiteLLM_ManagedObjectTable, LiteLLM_ManagedVectorStoresTable, ) from litellm.models.mcp_server import LiteLLM_MCPServerTable @@ -41,7 +40,6 @@ from litellm.models.verification_token import ( LiteLLM_DeletedVerificationToken, LiteLLM_VerificationToken, ) -from pydantic import ValidationError class TestBudget: @@ -121,9 +119,7 @@ class TestCredentials: assert item.credential_values is None def test_create_credential_item_requires_values_or_model_id(self): - with pytest.raises( - ValueError, match="Either credential_values or model_id must be set" - ): + with pytest.raises(ValueError, match="Either credential_values or model_id must be set"): CreateCredentialItem(credential_name="bad", credential_info={}) @@ -141,12 +137,8 @@ class TestModel: assert model.team_public_model_name == "my-gpt4" def test_is_blocked(self): - model_blocked = LiteLLM_ProxyModelTable( - model_id="m1", model_name="test", litellm_params={}, blocked=True - ) - model_unblocked = LiteLLM_ProxyModelTable( - model_id="m2", model_name="test", litellm_params={}, blocked=False - ) + model_blocked = LiteLLM_ProxyModelTable(model_id="m1", model_name="test", litellm_params={}, blocked=True) + model_unblocked = LiteLLM_ProxyModelTable(model_id="m2", model_name="test", litellm_params={}, blocked=False) assert model_blocked.is_blocked assert not model_unblocked.is_blocked @@ -188,9 +180,7 @@ class TestModel: assert model.blocked is True def test_team_helpers_none_when_no_model_info(self): - model = LiteLLM_ProxyModelTable( - model_id="m1", model_name="gpt-4", litellm_params={}, model_info=None - ) + model = LiteLLM_ProxyModelTable(model_id="m1", model_name="gpt-4", litellm_params={}, model_info=None) assert model.team_id is None assert model.team_public_model_name is None @@ -292,9 +282,7 @@ class TestTeam: assert team.model_max_budget == {"gpt-4": 5.0} def test_cached_team(self): - cached = LiteLLM_TeamTableCachedObj( - team_id="t1", last_refreshed_at=1234567890.0 - ) + cached = LiteLLM_TeamTableCachedObj(team_id="t1", last_refreshed_at=1234567890.0) assert cached.last_refreshed_at == 1234567890.0 def test_deleted_team(self): @@ -345,9 +333,7 @@ class TestUser: assert "password" not in user.model_dump() assert "password" not in user.model_dump_json() - with_keys = LiteLLM_UserTableWithKeyCount( - user_id="u1", user_email="a@b.c", password=secret, key_count=2 - ) + with_keys = LiteLLM_UserTableWithKeyCount(user_id="u1", user_email="a@b.c", password=secret, key_count=2) assert with_keys.password == secret assert "password" not in with_keys.model_dump() assert "password" not in with_keys.model_dump_json() @@ -479,9 +465,7 @@ class TestEndUserTable: class TestBudgetTableFull: def test_full_adds_server_managed_fields(self): now = datetime.now() - budget = LiteLLM_BudgetTableFull( - budget_id="b1", max_budget=10.0, created_at=now, budget_reset_at=now - ) + budget = LiteLLM_BudgetTableFull(budget_id="b1", max_budget=10.0, created_at=now, budget_reset_at=now) assert budget.created_at == now assert budget.budget_reset_at == now assert budget.max_budget == 10.0 @@ -493,9 +477,7 @@ class TestBudgetTableFull: class TestTeamMemberTable: def test_tracks_user_within_team(self): - member = LiteLLM_TeamMemberTable( - user_id="u1", team_id="t1", spend=3.0, budget_id="b1", max_budget=5.0 - ) + member = LiteLLM_TeamMemberTable(user_id="u1", team_id="t1", spend=3.0, budget_id="b1", max_budget=5.0) assert member.user_id == "u1" assert member.team_id == "t1" assert member.spend == 3.0 @@ -585,9 +567,7 @@ class TestSpendLogs: assert log.updated_at == updated_at def test_error_logs_creation(self): - log = LiteLLM_ErrorLogs( - request_id="r1", startTime=None, endTime=None, status_code="500" - ) + log = LiteLLM_ErrorLogs(request_id="r1", startTime=None, endTime=None, status_code="500") assert log.request_id == "r1" assert log.status_code == "500" @@ -603,12 +583,6 @@ class TestManagedTables: assert table.model_mappings == {"gpt-4": "file-abc"} assert table.flat_model_file_ids == ["file-abc"] - def test_managed_object_table_requires_purpose(self): - with pytest.raises(ValidationError): - LiteLLM_ManagedObjectTable( - unified_object_id="o1", model_object_id="m1", file_object={} - ) - def test_managed_vector_stores_table(self): table = LiteLLM_ManagedVectorStoresTable( vector_store_id="vs1", diff --git a/tests/unit/ocr/__init__.py b/tests/unit/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/ocr/test_dispatch.py b/tests/unit/ocr/test_dispatch.py similarity index 100% rename from tests/test_litellm/ocr/test_dispatch.py rename to tests/unit/ocr/test_dispatch.py diff --git a/tests/test_litellm/ocr/test_main.py b/tests/unit/ocr/test_main.py similarity index 100% rename from tests/test_litellm/ocr/test_main.py rename to tests/unit/ocr/test_main.py diff --git a/tests/test_litellm/ocr/test_ocr_file_input.py b/tests/unit/ocr/test_ocr_file_input.py similarity index 92% rename from tests/test_litellm/ocr/test_ocr_file_input.py rename to tests/unit/ocr/test_ocr_file_input.py index 4ac27d286e1..d67f5280195 100644 --- a/tests/test_litellm/ocr/test_ocr_file_input.py +++ b/tests/unit/ocr/test_ocr_file_input.py @@ -73,9 +73,7 @@ class TestConvertFileDocumentToUrlDocument: tmp_path = Path(f.name) try: - result = convert_file_document_to_url_document( - {"type": "file", "file": tmp_path} - ) + result = convert_file_document_to_url_document({"type": "file", "file": tmp_path}) assert result["type"] == "document_url" assert result["document_url"].startswith("data:application/pdf;base64,") @@ -95,9 +93,7 @@ class TestConvertFileDocumentToUrlDocument: tmp_path = Path(f.name) try: - result = convert_file_document_to_url_document( - {"type": "file", "file": tmp_path} - ) + result = convert_file_document_to_url_document({"type": "file", "file": tmp_path}) assert result["type"] == "image_url" assert result["image_url"].startswith("data:image/png;base64,") @@ -112,9 +108,7 @@ class TestConvertFileDocumentToUrlDocument: request handler the value is attacker-controlled, and opening it as a path is an arbitrary local file read on the proxy host.""" with pytest.raises(ValueError, match="does not accept bare str values"): - convert_file_document_to_url_document( - {"type": "file", "file": "/etc/passwd"} - ) + convert_file_document_to_url_document({"type": "file", "file": "/etc/passwd"}) def test_should_convert_pathlib_path(self): """pathlib.Path objects should work the same as string paths.""" @@ -126,9 +120,7 @@ class TestConvertFileDocumentToUrlDocument: tmp_path = Path(f.name) try: - result = convert_file_document_to_url_document( - {"type": "file", "file": tmp_path} - ) + result = convert_file_document_to_url_document({"type": "file", "file": tmp_path}) assert result["type"] == "document_url" assert result["document_url"].startswith("data:application/pdf;base64,") @@ -139,9 +131,7 @@ class TestConvertFileDocumentToUrlDocument: """Raw bytes should be converted using a fallback MIME type.""" content = b"raw bytes content" - result = convert_file_document_to_url_document( - {"type": "file", "file": content} - ) + result = convert_file_document_to_url_document({"type": "file", "file": content}) assert result["type"] == "document_url" assert "base64," in result["document_url"] @@ -164,9 +154,7 @@ class TestConvertFileDocumentToUrlDocument: """Raw bytes with an image MIME type should produce type=image_url.""" content = b"raw image content" - result = convert_file_document_to_url_document( - {"type": "file", "file": content, "mime_type": "image/jpeg"} - ) + result = convert_file_document_to_url_document({"type": "file", "file": content, "mime_type": "image/jpeg"}) assert result["type"] == "image_url" assert result["image_url"].startswith("data:image/jpeg;base64,") @@ -176,9 +164,7 @@ class TestConvertFileDocumentToUrlDocument: content = b"file-like content" file_obj = BytesIO(content) - result = convert_file_document_to_url_document( - {"type": "file", "file": file_obj} - ) + result = convert_file_document_to_url_document({"type": "file", "file": file_obj}) assert result["type"] == "document_url" assert "base64," in result["document_url"] @@ -189,9 +175,7 @@ class TestConvertFileDocumentToUrlDocument: file_obj = BytesIO(content) file_obj.name = "test_image.png" - result = convert_file_document_to_url_document( - {"type": "file", "file": file_obj} - ) + result = convert_file_document_to_url_document({"type": "file", "file": file_obj}) assert result["type"] == "image_url" assert result["image_url"].startswith("data:image/png;base64,") @@ -204,9 +188,7 @@ class TestConvertFileDocumentToUrlDocument: def test_should_raise_error_for_nonexistent_pathlib_path(self): """Non-existent pathlib.Path should raise FileNotFoundError.""" with pytest.raises(FileNotFoundError, match="File not found"): - convert_file_document_to_url_document( - {"type": "file", "file": Path("/nonexistent/path/to/file.pdf")} - ) + convert_file_document_to_url_document({"type": "file", "file": Path("/nonexistent/path/to/file.pdf")}) def test_should_raise_error_for_empty_file(self): """Empty file should raise ValueError.""" @@ -215,9 +197,7 @@ class TestConvertFileDocumentToUrlDocument: try: with pytest.raises(ValueError, match="File is empty"): - convert_file_document_to_url_document( - {"type": "file", "file": tmp_path} - ) + convert_file_document_to_url_document({"type": "file", "file": tmp_path}) finally: os.unlink(str(tmp_path)) @@ -248,9 +228,7 @@ class TestConvertFileDocumentToUrlDocument: tmp_path = Path(f.name) try: - result = convert_file_document_to_url_document( - {"type": "file", "file": tmp_path, "mime_type": "image/png"} - ) + result = convert_file_document_to_url_document({"type": "file", "file": tmp_path, "mime_type": "image/png"}) assert result["type"] == "image_url" assert result["image_url"].startswith("data:image/png;base64,") @@ -477,9 +455,7 @@ class TestProxySecurityGuard: result = await self._parse_multipart(mock_request) assert result["document"]["type"] == "document_url" - assert result["document"]["document_url"].startswith( - "data:application/pdf;base64," - ) + assert result["document"]["document_url"].startswith("data:application/pdf;base64,") assert result["model"] == "mistral/mistral-ocr-latest" diff --git a/tests/unit/passthrough/__init__.py b/tests/unit/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py b/tests/unit/passthrough/test_async_streaming_error_propagation.py similarity index 92% rename from tests/test_litellm/passthrough/test_async_streaming_error_propagation.py rename to tests/unit/passthrough/test_async_streaming_error_propagation.py index 9f2b436d2d8..cb93183957c 100644 --- a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py +++ b/tests/unit/passthrough/test_async_streaming_error_propagation.py @@ -21,9 +21,7 @@ def _make_mock_response(status_code: int, body: bytes, headers: dict = None): # def _raise_for_status(): if status_code >= 400: - request = httpx.Request( - "POST", "https://azure.example.com/openai/responses" - ) + request = httpx.Request("POST", "https://azure.example.com/openai/responses") real_response = httpx.Response( status_code=status_code, content=body, @@ -55,16 +53,15 @@ def _make_mock_logging_obj(): async def test_async_streaming_429_raises(): """429 from upstream should raise HTTPStatusError, not yield error bytes.""" from litellm.passthrough.main import AsyncPassthroughStreamingResponse - - error_body = json.dumps( - {"error": {"code": "429", "message": "Rate limit exceeded."}} - ).encode() + + error_body = json.dumps({"error": {"code": "429", "message": "Rate limit exceeded."}}).encode() mock_response = _make_mock_response(429, error_body) - + async def response_coro(): return mock_response - + chunks = [] + async def _drain(): async for chunk in AsyncPassthroughStreamingResponse( response=response_coro(), @@ -84,15 +81,13 @@ async def test_async_streaming_429_raises(): async def test_async_streaming_500_raises(): """500 from upstream should also raise, not yield error bytes.""" from litellm.passthrough.main import AsyncPassthroughStreamingResponse - - error_body = json.dumps( - {"error": {"code": "500", "message": "Internal server error"}} - ).encode() + + error_body = json.dumps({"error": {"code": "500", "message": "Internal server error"}}).encode() mock_response = _make_mock_response(500, error_body) - + async def response_coro(): return mock_response - + with pytest.raises(httpx.HTTPStatusError) as exc_info: async for _ in AsyncPassthroughStreamingResponse( response=response_coro(), @@ -100,7 +95,7 @@ async def test_async_streaming_500_raises(): provider_config=MagicMock(), ): pass - + assert exc_info.value.response.status_code == 500 diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/unit/passthrough/test_passthrough_main.py similarity index 94% rename from tests/test_litellm/passthrough/test_passthrough_main.py rename to tests/unit/passthrough/test_passthrough_main.py index 3f2c434cc00..82825ec2802 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/unit/passthrough/test_passthrough_main.py @@ -3,14 +3,9 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest -from fastapi.testclient import TestClient - -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler - - - import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.passthrough.main import allm_passthrough_route, llm_passthrough_route @@ -37,10 +32,7 @@ def test_llm_passthrough_route(): client=client, ) - assert ( - mock_post.call_args.kwargs["request"].url - == "http://localhost:8090/v1/chat/completions" - ) + assert mock_post.call_args.kwargs["request"].url == "http://localhost:8090/v1/chat/completions" assert response.status_code == 200 assert response.json == {"message": "Hello, world!"} @@ -74,12 +66,9 @@ def test_bedrock_application_inference_profile_url_encoding(): "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("test-model", "bedrock", "test-key", "test-base"), ), - patch.object( - client.client, "send", return_value=MagicMock(status_code=200) - ) as mock_send, + patch.object(client.client, "send", return_value=MagicMock(status_code=200)), patch.object(client.client, "build_request") as mock_build_request, ): - # Mock logging object mock_logging_obj = MagicMock() mock_logging_obj.update_environment_variables = MagicMock() @@ -132,12 +121,9 @@ def test_bedrock_non_application_inference_profile_no_encoding(): "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("test-model", "bedrock", "test-key", "test-base"), ), - patch.object( - client.client, "send", return_value=MagicMock(status_code=200) - ) as mock_send, + patch.object(client.client, "send", return_value=MagicMock(status_code=200)), patch.object(client.client, "build_request") as mock_build_request, ): - # Mock logging object mock_logging_obj = MagicMock() mock_logging_obj.update_environment_variables = MagicMock() @@ -202,7 +188,6 @@ def test_update_stream_param_based_on_request_body(): @pytest.fixture def mock_request(): """Create a mock request with headers""" - from typing import Optional class QueryParams: def __init__(self): @@ -215,9 +200,7 @@ def mock_request(): return self._dict.items() class MockRequest: - def __init__( - self, headers=None, method="POST", request_body: Optional[dict] = None - ): + def __init__(self, headers=None, method="POST", request_body: dict | None = None): self.headers = headers or {} self.query_params = QueryParams() self.method = method @@ -245,9 +228,7 @@ def mock_user_api_key_dict(): @pytest.mark.asyncio -async def test_pass_through_request_stream_param_override( - mock_request, mock_user_api_key_dict -): +async def test_pass_through_request_stream_param_override(mock_request, mock_user_api_key_dict): """ Test that when stream=None is passed as parameter but stream=True is in request body, the request body value takes precedence and @@ -346,9 +327,7 @@ async def test_pass_through_request_stream_param_override( @pytest.mark.asyncio -async def test_pass_through_request_stream_param_no_override( - mock_request, mock_user_api_key_dict -): +async def test_pass_through_request_stream_param_no_override(mock_request, mock_user_api_key_dict): """ Test that when stream=False is passed as parameter and no stream is in request body, the function parameter is used and @@ -448,15 +427,11 @@ def test_azure_with_custom_api_base_and_key(): # Mock the provider config and its methods mock_provider_config = MagicMock() mock_provider_config.get_complete_url.return_value = ( - httpx.URL( - "https://my-custom-base/openai/deployments/gpt-4.1/chat/completions?api-version=2024-02-01" - ), + httpx.URL("https://my-custom-base/openai/deployments/gpt-4.1/chat/completions?api-version=2024-02-01"), "https://my-custom-base", ) mock_provider_config.get_api_key.return_value = "my-custom-key" - mock_provider_config.validate_environment.return_value = { - "api-key": "my-custom-key" - } + mock_provider_config.validate_environment.return_value = {"api-key": "my-custom-key"} mock_provider_config.sign_request.return_value = ( {"api-key": "my-custom-key"}, None, @@ -484,13 +459,10 @@ def test_azure_with_custom_api_base_and_key(): patch.object( client.client, "send", - return_value=MagicMock( - status_code=200, json=lambda: {"id": "chatcmpl-123", "choices": []} - ), - ) as mock_send, + return_value=MagicMock(status_code=200, json=lambda: {"id": "chatcmpl-123", "choices": []}), + ), patch.object(client.client, "build_request") as mock_build_request, ): - # Mock logging object mock_logging_obj = MagicMock() mock_logging_obj.update_environment_variables = MagicMock() @@ -541,9 +513,7 @@ def test_content_param_forwarded_to_build_request(): mock_provider_config = MagicMock() mock_provider_config.get_complete_url.return_value = ( - httpx.URL( - "https://my-azure.openai.azure.com/openai/deployments/gpt-4/chat/completions" - ), + httpx.URL("https://my-azure.openai.azure.com/openai/deployments/gpt-4/chat/completions"), "https://my-azure.openai.azure.com", ) mock_provider_config.get_api_key.return_value = "test-key" @@ -575,7 +545,6 @@ def test_content_param_forwarded_to_build_request(): patch.object(client.client, "send", return_value=MagicMock(status_code=200)), patch.object(client.client, "build_request") as mock_build_request, ): - mock_logging_obj = MagicMock() mock_logging_obj.update_environment_variables = MagicMock() @@ -656,15 +625,11 @@ async def test_allm_passthrough_route_429_streaming_raises(): """ mock_provider_config = MagicMock() mock_provider_config.get_complete_url.return_value = ( - httpx.URL( - "https://my-azure.openai.azure.com/openai/deployments/gpt-4/responses" - ), + httpx.URL("https://my-azure.openai.azure.com/openai/deployments/gpt-4/responses"), "https://my-azure.openai.azure.com", ) mock_provider_config.get_api_key.return_value = "fake-azure-key" - mock_provider_config.validate_environment.return_value = { - "api-key": "fake-azure-key" - } + mock_provider_config.validate_environment.return_value = {"api-key": "fake-azure-key"} mock_provider_config.sign_request.return_value = ( {"api-key": "fake-azure-key"}, None, @@ -752,9 +717,7 @@ def test_llm_passthrough_route_sync_streaming_error_maps_upstream_status(): headers={"content-type": "application/json"}, ) - sync_client = HTTPHandler( - client=httpx.Client(transport=httpx.MockTransport(_handler)) - ) + sync_client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_handler))) mock_provider_config = MagicMock() mock_provider_config.get_complete_url.return_value = ( @@ -762,18 +725,14 @@ def test_llm_passthrough_route_sync_streaming_error_maps_upstream_status(): "https://gigachat.devices.sberbank.ru/api/v1", ) mock_provider_config.get_api_key.return_value = "fake-key" - mock_provider_config.validate_environment.return_value = { - "Authorization": "Bearer fake-key" - } + mock_provider_config.validate_environment.return_value = {"Authorization": "Bearer fake-key"} mock_provider_config.sign_request.return_value = ( {"Authorization": "Bearer fake-key"}, None, ) mock_provider_config.is_streaming_request.return_value = True - mock_provider_config.get_error_class.side_effect = ( - lambda error_message, status_code, headers: BaseLLMException( - status_code=status_code, message=error_message, headers=headers - ) + mock_provider_config.get_error_class.side_effect = lambda error_message, status_code, headers: BaseLLMException( + status_code=status_code, message=error_message, headers=headers ) mock_logging_obj = MagicMock() diff --git a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py b/tests/unit/passthrough/test_streaming_interrupt_spend_tracking.py similarity index 91% rename from tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py rename to tests/unit/passthrough/test_streaming_interrupt_spend_tracking.py index 5e13db9439b..922643f9834 100644 --- a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py +++ b/tests/unit/passthrough/test_streaming_interrupt_spend_tracking.py @@ -68,9 +68,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_normal_completion(): chunks = [b"chunk-1", b"chunk-2", b"chunk-3"] mock_response = _make_streaming_response(chunks) - mock_response.headers = httpx.Headers( - {"content-type": "application/octet-stream", "x-request-id": "req-123"} - ) + mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"}) async def response_coro(): return mock_response @@ -88,7 +86,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_normal_completion(): received.append(chunk) assert received == chunks - + assert received_response.headers["content-type"] == "application/octet-stream" assert received_response.headers["x-request-id"] == "req-123" @@ -107,9 +105,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_client_disconnect(): b'{"chunk": 3, "outputTokens": 8}', ] mock_response = _make_streaming_response(chunks) - mock_response.headers = httpx.Headers( - {"content-type": "application/octet-stream", "x-request-id": "req-123"} - ) + mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"}) async def response_coro(): return mock_response @@ -138,17 +134,13 @@ async def test_asyncpassthroughstreamingresponse_does_not_flush_on_4xx(): err_response = MagicMock(spec=httpx.Response) err_response.status_code = 429 - err_response.headers = httpx.Headers( - {"content-type": "application/octet-stream"} - ) + err_response.headers = httpx.Headers({"content-type": "application/octet-stream"}) def _raise(): raise httpx.HTTPStatusError( "429", request=httpx.Request("POST", "https://example.com"), - response=httpx.Response( - 429, request=httpx.Request("POST", "https://example.com") - ), + response=httpx.Response(429, request=httpx.Request("POST", "https://example.com")), ) err_response.raise_for_status = _raise @@ -180,9 +172,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_upstream_exception_w mock_response.status_code = 200 mock_response.raise_for_status = MagicMock(return_value=None) mock_response.aclose = AsyncMock() - mock_response.headers = httpx.Headers( - {"content-type": "application/octet-stream", "x-request-id": "req-123"} - ) + mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"}) async def _aiter_bytes_then_raise(): for c in partial_chunks: @@ -197,6 +187,7 @@ async def test_asyncpassthroughstreamingresponse_flushes_on_upstream_exception_w mock_logging_obj = _make_logging_obj() received = [] + async def _drain(): async for chunk in AsyncPassthroughStreamingResponse( response=response_coro(), @@ -222,9 +213,7 @@ def test_passthroughstreamingresponse_flushes_on_normal_completion(): mock_response = MagicMock(spec=httpx.Response) mock_response.status_code = 200 - mock_response.headers = httpx.Headers( - {"content-type": "application/octet-stream", "x-request-id": "req-123"} - ) + mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"}) def _iter_bytes(): yield from chunks @@ -258,9 +247,7 @@ def test_passthroughstreamingresponse_flushes_on_early_close(): mock_response = MagicMock(spec=httpx.Response) mock_response.status_code = 200 - mock_response.headers = httpx.Headers( - {"content-type": "application/octet-stream", "x-request-id": "req-123"} - ) + mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"}) def _iter_bytes(): yield from chunks diff --git a/tests/unit/rag/ingestion/__init__.py b/tests/unit/rag/ingestion/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py b/tests/unit/rag/ingestion/test_s3_vectors_ingestion.py similarity index 94% rename from tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py rename to tests/unit/rag/ingestion/test_s3_vectors_ingestion.py index 07fd2b765f3..1e5b62456b6 100644 --- a/tests/test_litellm/rag/ingestion/test_s3_vectors_ingestion.py +++ b/tests/unit/rag/ingestion/test_s3_vectors_ingestion.py @@ -21,10 +21,14 @@ class _RecordingRouter: def _ingestion(embedding=REQUEST_EMBEDDING, router=None, **vector_store): vector_store_options = {"custom_llm_provider": "s3_vectors", "aws_region_name": "us-west-2", **vector_store} - ingest_options = {"vector_store": vector_store_options} if embedding is None else { - "embedding": embedding, - "vector_store": vector_store_options, - } + ingest_options = ( + {"vector_store": vector_store_options} + if embedding is None + else { + "embedding": embedding, + "vector_store": vector_store_options, + } + ) return S3VectorsRAGIngestion(ingest_options=ingest_options, router=router) diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/unit/realtime_api/test_main.py similarity index 98% rename from tests/test_litellm/realtime_api/test_main.py rename to tests/unit/realtime_api/test_main.py index 86b25b2f9c8..5d3276dfae1 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/unit/realtime_api/test_main.py @@ -12,6 +12,15 @@ from litellm.realtime_api import main as realtime_main from litellm.realtime_api.main import _with_resolved_session_model +@pytest.fixture +def local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + class FakeLogging: def update_from_kwargs(self, **kwargs): pass @@ -502,8 +511,8 @@ async def test_arealtime_azure_env_beta_protocol_wins_over_a_ga_client(monkeypat async def _vertex_provider_config_for(monkeypatch, model: str, vertex_location: str | None): - from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import VertexChirpRealtimeConfig + from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig captured: dict[str, object] = {} diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/unit/repositories/test_repositories.py similarity index 98% rename from tests/test_litellm/repositories/test_repositories.py rename to tests/unit/repositories/test_repositories.py index 63fde9b2b8f..87cf2fc4268 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/unit/repositories/test_repositories.py @@ -78,17 +78,11 @@ class MockTable: record_data = dict(data) if self._pk_field and self._pk_field not in record_data: record_data[self._pk_field] = f"{self._pk_field}-{len(self._records)}" - key = ( - record_data.get(self._pk_field) - if self._pk_field - else record_data.get("id", str(len(self._records))) - ) + key = record_data.get(self._pk_field) if self._pk_field else record_data.get("id", str(len(self._records))) self._records[key] = record_data return MockRecord(record_data) - async def update( - self, where: Dict[str, Any], data: Dict[str, Any] - ) -> Optional[MockRecord]: + async def update(self, where: Dict[str, Any], data: Dict[str, Any]) -> Optional[MockRecord]: key_field = list(where.keys())[0] key_value = where[key_field] if key_value in self._records: @@ -140,9 +134,7 @@ class MockPrismaClient: self.db.litellm_config = MockTable() self.db.litellm_organizationtable = MockTable() self.db.litellm_projecttable = MockTable(pk_field="project_id") - self.db.litellm_objectpermissiontable = MockTable( - pk_field="object_permission_id" - ) + self.db.litellm_objectpermissiontable = MockTable(pk_field="object_permission_id") self.db.litellm_credentialstable = MockTable() @@ -200,9 +192,7 @@ class TestBaseRepository: prisma_client.db.litellm_budgettable._records = { "b1": {"budget_id": "b1", "max_budget": 100.0}, } - budgets = await repo.find_many( - where={"budget_id": "b1"}, skip=0, take=10, order={"budget_id": "asc"} - ) + budgets = await repo.find_many(where={"budget_id": "b1"}, skip=0, take=10, order={"budget_id": "asc"}) assert len(budgets) == 1 def test_record_to_dict_branches(self): @@ -1518,9 +1508,7 @@ class TestVerificationTokenRepositoryExtended: class MockTx: def __init__(self, client): - self.litellm_deletedverificationtoken = ( - client.db.litellm_deletedverificationtoken - ) + self.litellm_deletedverificationtoken = client.db.litellm_deletedverificationtoken self.litellm_verificationtoken = client.db.litellm_verificationtoken async def __aenter__(self): @@ -1563,9 +1551,7 @@ class TestVerificationTokenRepositoryExtended: class MockTx: def __init__(self, client): - self.litellm_deletedverificationtoken = ( - client.db.litellm_deletedverificationtoken - ) + self.litellm_deletedverificationtoken = client.db.litellm_deletedverificationtoken self.litellm_verificationtoken = client.db.litellm_verificationtoken async def __aenter__(self): @@ -1578,9 +1564,7 @@ class TestVerificationTokenRepositoryExtended: await repo.delete_token("sk-arch", deleted_by="admin") - archived = list( - repo._prisma_client.db.litellm_deletedverificationtoken._records.values() - )[0] + archived = list(repo._prisma_client.db.litellm_deletedverificationtoken._records.values())[0] assert isinstance(archived["aliases"], str) assert json.loads(archived["aliases"]) == {"a": "b"} @@ -1599,9 +1583,7 @@ class TestVerificationTokenRepositoryExtended: ): assert relation_field not in archived - assert ( - "sk-arch" not in repo._prisma_client.db.litellm_verificationtoken._records - ) + assert "sk-arch" not in repo._prisma_client.db.litellm_verificationtoken._records @pytest.mark.asyncio async def test_find_by_id_maps_org_and_budget_columns(self, repo): @@ -1977,9 +1959,7 @@ class TestDomainModelExtended: DomainModel.from_db_record(None) def test_from_db_record_dict(self): - model = _SampleDomainModel.from_db_record( - {"budget_id": "b1", "max_budget": 100.0} - ) + model = _SampleDomainModel.from_db_record({"budget_id": "b1", "max_budget": 100.0}) assert model.budget_id == "b1" def test_from_db_record_model_dump(self): @@ -2174,9 +2154,7 @@ class TestPrismaTableRepository: assert self.CONFIG_SYNCED_TABLE_NAMES <= seen -def _json_path_equals( - metadata: Optional[Dict[str, Any]], path: List[str], expected: Any -) -> bool: +def _json_path_equals(metadata: Optional[Dict[str, Any]], path: List[str], expected: Any) -> bool: """Reproduce Postgres jsonb path-equals semantics: a missing path yields SQL NULL, which never matches `equals`.""" value: Any = metadata @@ -2201,11 +2179,7 @@ class _ScimAwareUserTable: json_filter = where["metadata"] path = json_filter["path"] expected = getattr(json_filter["equals"], "data", json_filter["equals"]) - return sum( - 1 - for metadata in self._metadatas - if _json_path_equals(metadata, path, expected) - ) + return sum(1 for metadata in self._metadatas if _json_path_equals(metadata, path, expected)) class TestCountBillableUsers: diff --git a/tests/test_litellm/repositories/test_unit_of_work.py b/tests/unit/repositories/test_unit_of_work.py similarity index 100% rename from tests/test_litellm/repositories/test_unit_of_work.py rename to tests/unit/repositories/test_unit_of_work.py diff --git a/tests/test_litellm/router_strategy/complexity_router/test_jev_classifier.py b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py similarity index 100% rename from tests/test_litellm/router_strategy/complexity_router/test_jev_classifier.py rename to tests/unit/router_strategy/complexity_router/test_jev_classifier.py diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py similarity index 95% rename from tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py rename to tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index b5651062098..d8bc4c45ab8 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -14,6 +14,30 @@ from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( ) +@pytest.fixture(autouse=True) +def isolate_litellm_router_state(): + saved = { + name: getattr(litellm, name).copy() + if isinstance(getattr(litellm, name, None), list) + else getattr(litellm, name, None) + for name in ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", + "model_fallbacks", + "cache", + ) + if hasattr(litellm, name) + } + yield + for name, value in saved.items(): + setattr(litellm, name, value) + + class MockResponse: def __init__(self, json_data, status_code): self._json_data = json_data @@ -43,9 +67,7 @@ async def test_async_user_key_affinity_routes_to_same_deployment(): "id": "msg_123", "status": "completed", "role": "assistant", - "content": [ - {"type": "output_text", "text": "Hello there!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello there!", "annotations": []}], } ], "parallel_tool_calls": True, @@ -348,9 +370,7 @@ async def test_async_previous_response_id_priority_over_user_key_affinity(): model_group=model_group, user_key=user_api_key_hash, ) - await router.cache.async_set_cache( - affinity_cache_key, {"model_id": other_model_id}, ttl=3600 - ) + await router.cache.async_set_cache(affinity_cache_key, {"model_id": other_model_id}, ttl=3600) # Even though user-key affinity points elsewhere, previous_response_id should pin # to the deployment that created the original response. @@ -519,9 +539,7 @@ async def test_async_filter_deployments_uses_stable_model_map_key_for_affinity_s }, { "model_name": stable_model_map_key, - "litellm_params": { - "model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0" - }, + "litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"}, "model_info": {"id": "deployment-2"}, }, ] @@ -542,9 +560,7 @@ async def test_async_filter_deployments_uses_stable_model_map_key_for_affinity_s model="some-router-model-group", healthy_deployments=healthy_deployments, messages=None, - request_kwargs={ - "metadata": {"user_api_key_hash": user_key, "model_group": "alias-group"} - }, + request_kwargs={"metadata": {"user_api_key_hash": user_key, "model_group": "alias-group"}}, parent_otel_span=None, ) @@ -580,9 +596,7 @@ async def test_async_filter_deployments_falls_back_when_cached_deployment_is_unh }, { "model_name": stable_model_map_key, - "litellm_params": { - "model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0" - }, + "litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"}, "model_info": {"id": "deployment-2"}, }, ] @@ -618,9 +632,7 @@ async def test_async_filter_deployments_does_not_pin_when_target_order_is_set(): }, { "model_name": stable_model_map_key, - "litellm_params": { - "model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0" - }, + "litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"}, "model_info": {"id": "deployment-2"}, }, ] @@ -660,9 +672,7 @@ async def test_async_user_key_affinity_ttl_expiry_allows_reroute(): }, { "model_name": stable_model_map_key, - "litellm_params": { - "model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0" - }, + "litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"}, "model_info": {"id": "deployment-2"}, }, ] @@ -706,9 +716,7 @@ def test_cache_key_does_not_double_hash_user_api_key_hash(): The affinity cache key should not hash it again. """ - user_api_key_hash = ( - "b95b015b66dd02a1c14e1e0a8729211f8ee53ec962658764f4cf58546c2c68e1" - ) + user_api_key_hash = "b95b015b66dd02a1c14e1e0a8729211f8ee53ec962658764f4cf58546c2c68e1" key = DeploymentAffinityCheck.get_affinity_cache_key( model_group="any-model-group", user_key=user_api_key_hash, @@ -746,9 +754,7 @@ def test_get_effective_flags_returns_per_group_config(): assert session_id is True # unconfigured-model: falls back to global flags - user_key, responses_api, session_id = callback._get_effective_flags( - "unconfigured-model" - ) + user_key, responses_api, session_id = callback._get_effective_flags("unconfigured-model") assert user_key is True assert responses_api is True assert session_id is False @@ -980,12 +986,8 @@ async def test_model_group_affinity_config_overrides_global(): ] # Set up user-key affinity cache for claude-3 - cache_key = DeploymentAffinityCheck.get_affinity_cache_key( - model_group=stable_model_map_key, user_key=user_key - ) - await callback.cache.async_set_cache( - cache_key, {"model_id": "deployment-1"}, ttl=60 - ) + cache_key = DeploymentAffinityCheck.get_affinity_cache_key(model_group=stable_model_map_key, user_key=user_key) + await callback.cache.async_set_cache(cache_key, {"model_id": "deployment-1"}, ttl=60) # claude-3 has per-group config (session_affinity only), so user-key affinity # should NOT apply even though it's globally enabled @@ -1050,7 +1052,7 @@ async def test_async_jwt_user_affinity_routes_to_same_deployment(): return seq[0] return seq[1] if len(seq) > 1 else seq[0] - with patch( # test-quality-ok: simple-shuffle has no injectable RNG; forcing the other pick is what proves the pin overrides the strategy + with patch( "litellm.router_strategy.simple_shuffle.random.choice", side_effect=deterministic_choice, ): diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py similarity index 96% rename from tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py rename to tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index b93b8c1cdfc..aa34fbd6bf7 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -25,6 +25,31 @@ from litellm.models.credentials import CredentialItem from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIResponse + +@pytest.fixture(autouse=True) +def isolate_litellm_router_state(): + saved = { + name: getattr(litellm, name).copy() + if isinstance(getattr(litellm, name, None), list) + else getattr(litellm, name, None) + for name in ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", + "model_fallbacks", + "cache", + ) + if hasattr(litellm, name) + } + yield + for name, value in saved.items(): + setattr(litellm, name, value) + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -1088,21 +1113,19 @@ def test_boundary_key_resolves_missing_values_from_named_credential(): EncryptedContentAffinityCheck, ) - with ( - patch.object( # test-quality-ok: credential registry is the direct dependency under test - litellm, - "credential_list", - [ - CredentialItem( - credential_name="account-a", - credential_values={ - "api_base": "https://account-a.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ) - ], - ) + with patch.object( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://account-a.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ) + ], ): boundary = EncryptedContentAffinityCheck._encryption_boundary_key({"litellm_credential_name": "account-a"}) @@ -1114,21 +1137,19 @@ def test_boundary_key_matches_named_credential_precedence(): EncryptedContentAffinityCheck, ) - with ( - patch.object( # test-quality-ok: credential registry is the direct dependency under test - litellm, - "credential_list", - [ - CredentialItem( - credential_name="account-a", - credential_values={ - "api_base": "https://credential.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ) - ], - ) + with patch.object( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://credential.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ) + ], ): boundary = EncryptedContentAffinityCheck._encryption_boundary_key( { @@ -1146,21 +1167,19 @@ def test_boundary_key_resolves_credential_when_explicit_values_are_empty(): EncryptedContentAffinityCheck, ) - with ( - patch.object( # test-quality-ok: credential registry is the direct dependency under test - litellm, - "credential_list", - [ - CredentialItem( - credential_name="account-a", - credential_values={ - "api_base": "https://credential.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ) - ], - ) + with patch.object( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://credential.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ) + ], ): boundary = EncryptedContentAffinityCheck._encryption_boundary_key( { @@ -1178,37 +1197,35 @@ def test_boundary_fallback_matches_deployments_with_same_named_credential_values EncryptedContentAffinityCheck, ) - with ( - patch.object( # test-quality-ok: credential registry is the direct dependency under test - litellm, - "credential_list", - [ - CredentialItem( - credential_name="account-a", - credential_values={ - "api_base": "https://account-a.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ), - CredentialItem( - credential_name="account-a-peer", - credential_values={ - "api_base": "https://account-a.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ), - CredentialItem( - credential_name="account-b", - credential_values={ - "api_base": "https://account-b.example.com", - "api_key": "credential-key-b", - }, - credential_info={}, - ), - ], - ) + with patch.object( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://account-a.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ), + CredentialItem( + credential_name="account-a-peer", + credential_values={ + "api_base": "https://account-a.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ), + CredentialItem( + credential_name="account-b", + credential_values={ + "api_base": "https://account-b.example.com", + "api_key": "credential-key-b", + }, + credential_info={}, + ), + ], ): router = litellm.Router( model_list=[ diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py similarity index 94% rename from tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py rename to tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 333e7b2ff31..849edc8c537 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -1,10 +1,9 @@ import asyncio import copy -from typing import List, cast +from typing import cast import pytest - import litellm from litellm.caching.dual_cache import DualCache from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT @@ -22,6 +21,15 @@ MODEL_GROUP_ALIAS = "my-claude-group" OPUS_4_6_MIN_TOKENS = 4096 +@pytest.fixture +def local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + @pytest.fixture(autouse=True) def _local_model_cost_map_autouse(local_model_cost_map): """Every test here reads `prompt_cache_min_tokens`, which only the in-repo map @@ -30,8 +38,7 @@ def _local_model_cost_map_autouse(local_model_cost_map): yield - -def _deployments(*models: str) -> List[dict]: +def _deployments(*models: str) -> list[dict]: return [ { "model_name": MODEL_GROUP_ALIAS, @@ -42,9 +49,9 @@ def _deployments(*models: str) -> List[dict]: ] -def _messages(word_count: int) -> List[AllMessageValues]: +def _messages(word_count: int) -> list[AllMessageValues]: return cast( - List[AllMessageValues], + list[AllMessageValues], [ { "role": "user", @@ -84,7 +91,9 @@ def test_write_gate_is_what_prevents_a_pin_below_the_model_minimum(): """ messages = _messages(word_count=1400) - token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-5", use_default_image_token_count=True) + token_count = token_counter( + messages=messages, model="anthropic/claude-opus-4-5", use_default_image_token_count=True + ) assert 1024 < token_count < 4096 assert is_prompt_caching_valid_prompt(model="anthropic/claude-opus-4-5", messages=messages) is False @@ -110,7 +119,9 @@ async def test_async_filter_deployments_does_not_narrow_prompt_below_model_minim deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") messages = _messages(word_count=1400) - token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True) + token_count = token_counter( + messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True + ) assert DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT < token_count < OPUS_4_6_MIN_TOKENS await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None) @@ -136,7 +147,9 @@ async def test_async_filter_deployments_narrows_prompt_above_model_minimum(): deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") messages = _messages(word_count=5000) - token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True) + token_count = token_counter( + messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True + ) assert token_count > OPUS_4_6_MIN_TOKENS await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None) @@ -197,10 +210,10 @@ async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is AUTO_CACHING_MODEL = "anthropic/claude-sonnet-4-5" -def _auto_caching_messages() -> List[AllMessageValues]: +def _auto_caching_messages() -> list[AllMessageValues]: """A prompt over the model minimum that carries no client cache_control.""" return cast( - List[AllMessageValues], + list[AllMessageValues], [ {"role": "system", "content": "word " * 3000}, {"role": "user", "content": "hello"}, @@ -208,7 +221,7 @@ def _auto_caching_messages() -> List[AllMessageValues]: ) -def _affinity_messages(messages: List[AllMessageValues]) -> List[AllMessageValues]: +def _affinity_messages(messages: list[AllMessageValues]) -> list[AllMessageValues]: """The messages the check keys deployment affinity on, for a group of `AUTO_CACHING_MODEL`.""" return AnthropicCacheControlHook.messages_with_default_injections( messages=messages, @@ -218,7 +231,7 @@ def _affinity_messages(messages: List[AllMessageValues]) -> List[AllMessageValue class _SentMessagesCapture(CustomLogger): def __init__(self): - self.messages: List[AllMessageValues] | None = None + self.messages: list[AllMessageValues] | None = None async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): standard_logging_object = kwargs.get("standard_logging_object") @@ -338,7 +351,7 @@ async def test_claude_code_one_shot_subagent_does_not_reuse_an_auto_injected_aff cache = DualCache() check = PromptCachingDeploymentCheck(cache=cache) deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) - messages = cast(List[AllMessageValues], [{"role": "user", "content": "unique " * 3000}]) + messages = cast(list[AllMessageValues], [{"role": "user", "content": "unique " * 3000}]) request_kwargs = { "system": [ { @@ -441,7 +454,7 @@ def test_client_supplied_cache_control_keeps_its_own_prefix_boundary(monkeypatch """ monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) messages = cast( - List[AllMessageValues], + list[AllMessageValues], [ { "role": "system", @@ -491,7 +504,7 @@ async def test_async_filter_deployments_counts_the_prompt_off_the_event_loop(): warm_tokenizer("anthropic/claude-fable-5") check = PromptCachingDeploymentCheck(cache=DualCache()) deployments = _deployments("anthropic/claude-fable-5") - messages = cast(List[AllMessageValues], [{"role": "user", "content": text * 100}]) + messages = cast(list[AllMessageValues], [{"role": "user", "content": text * 100}]) result, took, lags = await timed_with_loop_lags( lambda: check.async_filter_deployments( @@ -516,7 +529,7 @@ async def test_async_log_success_event_counts_the_prompt_off_the_event_loop(): cache = DualCache() check = PromptCachingDeploymentCheck(cache=cache) messages = cast( - List[AllMessageValues], + list[AllMessageValues], [{"role": "user", "content": [{"type": "text", "text": text * 100, "cache_control": {"type": "ephemeral"}}]}], ) standard_logging_object = { diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_responses_api_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_responses_api_deployment_check.py similarity index 95% rename from tests/test_litellm/router_utils/pre_call_checks/test_responses_api_deployment_check.py rename to tests/unit/router_utils/pre_call_checks/test_responses_api_deployment_check.py index ee7fab7d19f..78cafbec70a 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_responses_api_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_responses_api_deployment_check.py @@ -1,21 +1,10 @@ import asyncio -from typing import Optional +import json from unittest.mock import AsyncMock, patch import pytest -import json - import litellm -from litellm.integrations.custom_logger import CustomLogger -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.types.llms.openai import ( - IncompleteDetails, - ResponseAPIUsage, - ResponseCompletedEvent, - ResponsesAPIResponse, -) -from litellm.types.utils import StandardLoggingPayload @pytest.mark.asyncio @@ -119,14 +108,11 @@ async def test_async_responses_api_routing_with_previous_response_id(): input="Hello, how are you?", truncation="auto", ) - print("RESPONSE", response) # Store the model_id from the response expected_model_id = response._hidden_params["model_id"] response_id = response.id - print("Response ID=", response_id, "came from model_id=", expected_model_id) - # Make 10 other requests with previous_response_id, assert that they are sent to the same model_id for i in range(10): # Reset the mock for the next call @@ -137,7 +123,7 @@ async def test_async_responses_api_routing_with_previous_response_id(): response = await router.aresponses( model=MODEL, - input=f"Follow-up question {i+1}", + input=f"Follow-up question {i + 1}", truncation="auto", previous_response_id=response_id, ) @@ -163,9 +149,7 @@ async def test_async_routing_without_previous_response_id(): "id": "msg_123", "status": "completed", "role": "assistant", - "content": [ - {"type": "output_text", "text": "Hello there!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello there!", "annotations": []}], } ], "parallel_tool_calls": True, @@ -266,9 +250,7 @@ async def test_async_routing_without_previous_response_id(): used_model_ids.add(response._hidden_params["model_id"]) # We should have used more than one model_id if load balancing is working - assert ( - len(used_model_ids) > 1 - ), "Load balancing isn't working, only one deployment was used" + assert len(used_model_ids) > 1, "Load balancing isn't working, only one deployment was used" @pytest.mark.asyncio diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py similarity index 95% rename from tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py rename to tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 780300bf9e1..9bbaed0ae1b 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -1,12 +1,10 @@ import asyncio +import json from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest - -import json - import litellm from litellm.caching.affinity_cache import claim_affinity_pin from litellm.caching.dual_cache import DualCache @@ -46,9 +44,7 @@ async def test_async_session_id_affinity_routes_to_same_deployment(): "id": "msg_123", "status": "completed", "role": "assistant", - "content": [ - {"type": "output_text", "text": "Hello there!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello there!", "annotations": []}], } ], "parallel_tool_calls": True, @@ -164,9 +160,7 @@ async def test_async_session_id_affinity_priority_over_user_key(): ) await callback.cache.async_set_cache( - DeploymentAffinityCheck.get_session_affinity_cache_key( - "model_group", "session1", user_key="user1" - ), + DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "session1", user_key="user1"), {"model_id": "deployment-2"}, ) @@ -175,9 +169,7 @@ async def test_async_session_id_affinity_priority_over_user_key(): model="model_group", healthy_deployments=healthy_deployments, messages=[], - request_kwargs={ - "metadata": {"user_api_key_hash": "user1", "session_id": "session1"} - }, + request_kwargs={"metadata": {"user_api_key_hash": "user1", "session_id": "session1"}}, ) assert len(filtered) == 1 @@ -575,16 +567,17 @@ async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down(): (None, {"model": "second"}), ], ) -async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl( - stored: object, expected: object -) -> None: +async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl(stored: object, expected: object) -> None: clock: Final = MagicMock(return_value=100.0) cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock)) cache.in_memory_cache.set_cache("tier-pin", stored, ttl=10) clock.return_value = 105.0 winner: Final = await claim_affinity_pin( - cache, "tier-pin", {"model": "second"}, 30, + cache, + "tier-pin", + {"model": "second"}, + 30, eligible_values=({"model": "first"}, {"model": "second"}), ) @@ -600,13 +593,18 @@ async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl( async def test_concurrent_eligible_claims_return_one_winner() -> None: cache: Final = DualCache() candidates: Final = ({"model": "first"}, {"model": "second"}) - winners: Final = await asyncio.gather(*( - claim_affinity_pin( - cache, "tier-pin", candidates[index % 2], 30, - eligible_values=candidates, + winners: Final = await asyncio.gather( + *( + claim_affinity_pin( + cache, + "tier-pin", + candidates[index % 2], + 30, + eligible_values=candidates, + ) + for index in range(20) ) - for index in range(20) - )) + ) assert winners == [{"model": "first"}] * 20 assert cache.in_memory_cache.get_cache("tier-pin") == {"model": "first"} @@ -628,23 +626,19 @@ async def test_legacy_deployment_claim_retains_decoder_and_keepalive( clock: Final = MagicMock(return_value=100.0) cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock)) callback: Final = DeploymentAffinityCheck( - cache=cache, ttl_seconds=30, - enable_user_key_affinity=False, enable_responses_api_affinity=False, + cache=cache, + ttl_seconds=30, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, ) cache.in_memory_cache.set_cache("deployment-pin", stored, ttl=10) clock.return_value = 105.0 - winner: Final = await callback._claim_pin( - "deployment-pin", {"model_id": "7"}, 30 - ) + winner: Final = await callback._claim_pin("deployment-pin", {"model_id": "7"}, 30) assert winner == expected - assert cache.in_memory_cache.ttl_dict["deployment-pin"] == ( - 135.0 if refresh else 110.0 - ) - assert cache.in_memory_cache.get_cache("deployment-pin") == ( - {"model_id": "7"} if refresh else stored - ) + assert cache.in_memory_cache.ttl_dict["deployment-pin"] == (135.0 if refresh else 110.0) + assert cache.in_memory_cache.get_cache("deployment-pin") == ({"model_id": "7"} if refresh else stored) @pytest.mark.asyncio @@ -668,13 +662,13 @@ async def test_redis_deployment_claim_preserves_legacy_result_decoding( redis.async_register_script.return_value = AsyncMock(return_value=raw) cache: Final = DualCache(redis_cache=redis) callback: Final = DeploymentAffinityCheck( - cache=cache, ttl_seconds=30, - enable_user_key_affinity=False, enable_responses_api_affinity=False, + cache=cache, + ttl_seconds=30, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, ) - winner: Final = await callback._claim_pin( - "deployment-pin", {"model_id": "candidate"}, 30 - ) + winner: Final = await callback._claim_pin("deployment-pin", {"model_id": "candidate"}, 30) assert winner == expected assert cache.in_memory_cache.get_cache("deployment-pin") == stored diff --git a/tests/unit/rust_bridge/__init__.py b/tests/unit/rust_bridge/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/rust_bridge/chat_completions/__init__.py b/tests/unit/rust_bridge/chat_completions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rust_bridge/chat_completions/test_route_host.py b/tests/unit/rust_bridge/chat_completions/test_route_host.py similarity index 100% rename from tests/test_litellm/rust_bridge/chat_completions/test_route_host.py rename to tests/unit/rust_bridge/chat_completions/test_route_host.py diff --git a/tests/unit/rust_bridge/messages/__init__.py b/tests/unit/rust_bridge/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rust_bridge/messages/test_route_host.py b/tests/unit/rust_bridge/messages/test_route_host.py similarity index 100% rename from tests/test_litellm/rust_bridge/messages/test_route_host.py rename to tests/unit/rust_bridge/messages/test_route_host.py From 73a35abeb303dff68f232719bac0893cddf69366 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 20 Sep 2026 11:57:53 +0000 Subject: [PATCH 41/47] refactor(types): keep object permission parsing as it was Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_helpers/object_permission_utils.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 61e432daa16..daab38d3662 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -8,7 +8,7 @@ from collections.abc import Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Optional +from typing import TYPE_CHECKING, Any, Final, Optional from fastapi import HTTPException, status from pydantic import TypeAdapter @@ -156,13 +156,12 @@ async def handle_update_object_permission_common( if prisma_client is None: raise ValueError("Prisma client not found") - raw_object_permission: Final[dict | str | None] = data_json.pop("object_permission", None) - if raw_object_permission is None: + new_object_permission: dict | str | None = data_json.pop("object_permission", None) + if new_object_permission is None: return None - new_object_permission: Final[object] = ( - json.loads(raw_object_permission) if isinstance(raw_object_permission, str) else raw_object_permission - ) + if isinstance(new_object_permission, str): + new_object_permission = json.loads(new_object_permission) upsert: Final = await prepare_object_permission_upsert( new_object_permission=new_object_permission if isinstance(new_object_permission, dict) else {}, @@ -231,7 +230,7 @@ def _dedupe_preserving_order(values: list[str]) -> list[str]: return result -def _mcp_server_identifier_matches(server: object, identifier: str) -> bool: +def _mcp_server_identifier_matches(server: Any, identifier: str) -> bool: return identifier in { getattr(server, "server_id", None), getattr(server, "alias", None), From eea00175655063cd398f7200e9663cc8be2aa6c5 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 12:36:13 +0000 Subject: [PATCH 42/47] test: add __init__.py to new tests/unit packages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/models/__init__.py | 0 tests/unit/realtime_api/__init__.py | 0 tests/unit/repositories/__init__.py | 0 tests/unit/router_strategy/complexity_router/__init__.py | 0 tests/unit/router_utils/pre_call_checks/__init__.py | 0 5 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/models/__init__.py create mode 100644 tests/unit/realtime_api/__init__.py create mode 100644 tests/unit/repositories/__init__.py create mode 100644 tests/unit/router_strategy/complexity_router/__init__.py create mode 100644 tests/unit/router_utils/pre_call_checks/__init__.py diff --git a/tests/unit/models/__init__.py b/tests/unit/models/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/realtime_api/__init__.py b/tests/unit/realtime_api/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/repositories/__init__.py b/tests/unit/repositories/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/router_strategy/complexity_router/__init__.py b/tests/unit/router_strategy/complexity_router/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/router_utils/pre_call_checks/__init__.py b/tests/unit/router_utils/pre_call_checks/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 42dd6a130016083166339c3ede9539a2c94a162c Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 12:36:15 +0000 Subject: [PATCH 43/47] test: add __init__.py to phase 12 unit test directories Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/llms/openrouter/responses/__init__.py | 0 tests/unit/llms/parallel_ai/__init__.py | 0 tests/unit/llms/parasail/__init__.py | 0 tests/unit/llms/perplexity/chat/__init__.py | 0 tests/unit/llms/perplexity/embedding/__init__.py | 0 tests/unit/llms/perplexity/responses/__init__.py | 0 tests/unit/llms/publicai/__init__.py | 0 tests/unit/llms/ragflow/chat/__init__.py | 0 tests/unit/llms/recraft/image_edit/__init__.py | 0 tests/unit/llms/recraft/image_generation/__init__.py | 0 tests/unit/llms/runwayml/__init__.py | 0 tests/unit/llms/runwayml/videos/__init__.py | 0 tests/unit/llms/s3_vectors/vector_stores/__init__.py | 0 tests/unit/llms/sap/__init__.py | 0 tests/unit/llms/scaleway/__init__.py | 0 tests/unit/llms/snowflake/__init__.py | 0 tests/unit/llms/soniox/__init__.py | 0 tests/unit/llms/stability/image_generation/__init__.py | 0 tests/unit/llms/tencent/chat/__init__.py | 0 19 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/llms/openrouter/responses/__init__.py create mode 100644 tests/unit/llms/parallel_ai/__init__.py create mode 100644 tests/unit/llms/parasail/__init__.py create mode 100644 tests/unit/llms/perplexity/chat/__init__.py create mode 100644 tests/unit/llms/perplexity/embedding/__init__.py create mode 100644 tests/unit/llms/perplexity/responses/__init__.py create mode 100644 tests/unit/llms/publicai/__init__.py create mode 100644 tests/unit/llms/ragflow/chat/__init__.py create mode 100644 tests/unit/llms/recraft/image_edit/__init__.py create mode 100644 tests/unit/llms/recraft/image_generation/__init__.py create mode 100644 tests/unit/llms/runwayml/__init__.py create mode 100644 tests/unit/llms/runwayml/videos/__init__.py create mode 100644 tests/unit/llms/s3_vectors/vector_stores/__init__.py create mode 100644 tests/unit/llms/sap/__init__.py create mode 100644 tests/unit/llms/scaleway/__init__.py create mode 100644 tests/unit/llms/snowflake/__init__.py create mode 100644 tests/unit/llms/soniox/__init__.py create mode 100644 tests/unit/llms/stability/image_generation/__init__.py create mode 100644 tests/unit/llms/tencent/chat/__init__.py diff --git a/tests/unit/llms/openrouter/responses/__init__.py b/tests/unit/llms/openrouter/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/parallel_ai/__init__.py b/tests/unit/llms/parallel_ai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/parasail/__init__.py b/tests/unit/llms/parasail/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/perplexity/chat/__init__.py b/tests/unit/llms/perplexity/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/perplexity/embedding/__init__.py b/tests/unit/llms/perplexity/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/perplexity/responses/__init__.py b/tests/unit/llms/perplexity/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/publicai/__init__.py b/tests/unit/llms/publicai/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/ragflow/chat/__init__.py b/tests/unit/llms/ragflow/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/recraft/image_edit/__init__.py b/tests/unit/llms/recraft/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/recraft/image_generation/__init__.py b/tests/unit/llms/recraft/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/runwayml/__init__.py b/tests/unit/llms/runwayml/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/runwayml/videos/__init__.py b/tests/unit/llms/runwayml/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/s3_vectors/vector_stores/__init__.py b/tests/unit/llms/s3_vectors/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/sap/__init__.py b/tests/unit/llms/sap/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/scaleway/__init__.py b/tests/unit/llms/scaleway/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/snowflake/__init__.py b/tests/unit/llms/snowflake/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/soniox/__init__.py b/tests/unit/llms/soniox/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/stability/image_generation/__init__.py b/tests/unit/llms/stability/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/tencent/chat/__init__.py b/tests/unit/llms/tencent/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From cd07acbdea3bc1d579e88e612c870b1b92b8ea28 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 12:37:21 +0000 Subject: [PATCH 44/47] test: add __init__.py to intermediate phase 12 unit test directories Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/llms/openrouter/__init__.py | 0 tests/unit/llms/perplexity/__init__.py | 0 tests/unit/llms/ragflow/__init__.py | 0 tests/unit/llms/recraft/__init__.py | 0 tests/unit/llms/s3_vectors/__init__.py | 0 tests/unit/llms/stability/__init__.py | 0 tests/unit/llms/tencent/__init__.py | 0 7 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/llms/openrouter/__init__.py create mode 100644 tests/unit/llms/perplexity/__init__.py create mode 100644 tests/unit/llms/ragflow/__init__.py create mode 100644 tests/unit/llms/recraft/__init__.py create mode 100644 tests/unit/llms/s3_vectors/__init__.py create mode 100644 tests/unit/llms/stability/__init__.py create mode 100644 tests/unit/llms/tencent/__init__.py diff --git a/tests/unit/llms/openrouter/__init__.py b/tests/unit/llms/openrouter/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/perplexity/__init__.py b/tests/unit/llms/perplexity/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/ragflow/__init__.py b/tests/unit/llms/ragflow/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/recraft/__init__.py b/tests/unit/llms/recraft/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/s3_vectors/__init__.py b/tests/unit/llms/s3_vectors/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/stability/__init__.py b/tests/unit/llms/stability/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/tencent/__init__.py b/tests/unit/llms/tencent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 4fb1a3e61043a8308de3159417f1fca3a329d52e Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 12:55:27 +0000 Subject: [PATCH 45/47] test: add __init__.py to intermediate tests/unit packages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/rag/__init__.py | 0 tests/unit/router_strategy/__init__.py | 0 tests/unit/router_utils/__init__.py | 0 3 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unit/rag/__init__.py create mode 100644 tests/unit/router_strategy/__init__.py create mode 100644 tests/unit/router_utils/__init__.py diff --git a/tests/unit/rag/__init__.py b/tests/unit/rag/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/router_strategy/__init__.py b/tests/unit/router_strategy/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/router_utils/__init__.py b/tests/unit/router_utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 30f3075010964a024480df032cc2bf818ac43607 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 13:44:30 +0000 Subject: [PATCH 46/47] test: use main's local_model_cost_map fixture in tests/unit/conftest.py Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/conftest.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 73a555e4455..b3bb19a8b8a 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -2,7 +2,6 @@ import os from collections.abc import Iterator from typing import Final -import litellm import pytest from pytest_socket import enable_socket, socket_allow_hosts @@ -23,19 +22,6 @@ AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = ( ) -@pytest.fixture -def local_model_cost_map(monkeypatch): - original_model_cost = litellm.model_cost - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - litellm.get_model_info.cache_clear() - try: - yield - finally: - litellm.model_cost = original_model_cost - litellm.get_model_info.cache_clear() - - def _allow_loopback_only() -> None: socket_allow_hosts(LOOPBACK_HOSTS, allow_unix_socket=True) From d086f8574dcaac12c494fa9d37d60cab3b5a5ab7 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 13:46:52 +0000 Subject: [PATCH 47/47] test: add sync and async custom_llm_provider bridge propagation tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...t_responses_bridge_provider_propagation.py | 52 +++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py b/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py index 09ef1889818..f2e36137a19 100644 --- a/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py +++ b/tests/unit/completion_extras/test_responses_bridge_provider_propagation.py @@ -40,6 +40,58 @@ def _bedrock_mantle_kwargs() -> dict: } +def _openai_kwargs() -> dict: + messages = [{"role": "user", "content": "hi"}] + logging_obj = LiteLLMLogging( + litellm_call_id="test-call", + call_type="completion", + model="gpt-5.5", + messages=messages, + function_id="fn-id", + stream=False, + start_time=datetime.now(), + ) + return { + "model": "gpt-5.5", + "custom_llm_provider": "openai", + "messages": messages, + "optional_params": {}, + "litellm_params": {}, + "headers": {}, + "model_response": ModelResponse(), + "logging_obj": logging_obj, + } + + +def test_completion_forwards_custom_llm_provider_to_responses(): + bridge = ResponsesToCompletionBridgeHandler() + cached = ModelResponse(id="chatcmpl-cached", model="gpt-5.5") + + with patch("litellm.responses", return_value=cached) as fake_responses: + result = bridge.completion(**_openai_kwargs()) + + assert result is cached + assert fake_responses.call_args.kwargs["custom_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_acompletion_forwards_custom_llm_provider_to_aresponses(): + bridge = ResponsesToCompletionBridgeHandler() + cached = ModelResponse(id="chatcmpl-cached", model="gpt-5.5") + + async def _fake_aresponses(**kwargs): + _fake_aresponses.kwargs = kwargs + return cached + + _fake_aresponses.kwargs = {} + + with patch("litellm.aresponses", _fake_aresponses): + result = await bridge.acompletion(**_openai_kwargs()) + + assert result is cached + assert _fake_aresponses.kwargs["custom_llm_provider"] == "openai" + + @pytest.mark.asyncio async def test_acompletion_forwards_aws_region_name_to_aresponses(): bridge = ResponsesToCompletionBridgeHandler()