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"}