From 7cf9a3035ce130fb4a07a02c7e139df6915fd378 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 10:59:02 +0000 Subject: [PATCH 1/3] 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 2/3] 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 434af084e7572d2b2f23101f56d7c722189aa33d Mon Sep 17 00:00:00 2001 From: yuneng Date: Sun, 20 Sep 2026 11:52:38 +0000 Subject: [PATCH 3/3] 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