diff --git a/.circleci/scripts/classify_changes.sh b/.circleci/scripts/classify_changes.sh index ad265a5e39f..8c2ac019b99 100755 --- a/.circleci/scripts/classify_changes.sh +++ b/.circleci/scripts/classify_changes.sh @@ -14,12 +14,12 @@ while IFS= read -r file || [ -n "$file" ]; do [ -n "$file" ] || continue case "$file" in *.md | *.mdx) : ;; - pyproject.toml | */pyproject.toml | uv.lock | uv.toml | .python-version | rust-toolchain.toml | litellm-rust/* | litellm/__init__.py | litellm/proxy/proxy_server.py | litellm/*mcp* | tests/*mcp* | litellm/integrations/arize/* | tests/base_sdk_tests/* | scripts/check_mcp_sdk_install.py | .github/workflows/test-mcp-dependency-resolution.yml | .github/actions/detect-changes/* | .github/actions/setup-uv-with-retries/* | .github/actions/cache-cargo-build/* | .github/scripts/detect_changes.sh | .github/scripts/uv_sync_with_retries.sh | .circleci/scripts/classify_changes.sh | tests/test_litellm/test_circleci_path_filter.py | tests/test_litellm/test_detect_changes.py) + pyproject.toml | */pyproject.toml | uv.lock | uv.toml | .python-version | rust-toolchain.toml | litellm-rust/* | litellm/__init__.py | litellm/proxy/proxy_server.py | litellm/*mcp* | tests/*mcp* | litellm/integrations/arize/* | tests/base_sdk_tests/* | scripts/check_mcp_sdk_install.py | .github/workflows/test-mcp-dependency-resolution.yml | .github/actions/detect-changes/* | .github/actions/setup-uv-with-retries/* | .github/actions/cache-cargo-build/* | .github/scripts/detect_changes.sh | .github/scripts/uv_sync_with_retries.sh | .circleci/scripts/classify_changes.sh | tests/unit/test_circleci_path_filter.py | tests/unit/test_detect_changes.py) has_mcp_dependencies=true ;; esac case "$file" in tests/e2e/*/*.py) : ;; - tests/e2e/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/test_litellm/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock) + tests/e2e/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/unit/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock) has_provider_harness=true ;; esac case "$file" in diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh index f2ee7550df3..5ce8b6c84ba 100755 --- a/.circleci/scripts/unit_selection.sh +++ b/.circleci/scripts/unit_selection.sh @@ -8,6 +8,7 @@ legacy_flags=( enterprise-package enterprise-routing mcp-integration + misc proxy-db-auth-checks proxy-db-budgets proxy-db-custom-logging @@ -22,6 +23,7 @@ legacy_flags=( proxy-db-proxy-utils proxy-extras proxy-infra + responses-caching-types ) legacy_paths() { @@ -36,6 +38,7 @@ legacy_paths() { echo tests/unit/enterprise/proxy/test_audit_logging_endpoints.py echo tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py ;; enterprise-routing) + echo tests/unit/google_genai echo tests/unit/enterprise/enterprise_callbacks/send_emails echo tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py echo tests/unit/enterprise/proxy/test_batch_retrieve_input_file_id.py @@ -48,9 +51,28 @@ legacy_paths() { echo tests/unit/enterprise/proxy/test_managed_files_access_check.py echo tests/unit/enterprise/proxy/test_managed_files_hook.py ;; mcp-integration) + echo tests/unit/experimental_mcp_client echo tests/unit/proxy/_experimental/mcp_server echo tests/unit/responses/mcp echo tests/mcp_tests/test_proxy_mcp_e2e.py ;; + misc) + find tests/unit -maxdepth 1 -name 'test_*.py' + echo tests/unit/test_router + echo tests/unit/a2a_protocol + echo tests/unit/batches + echo tests/unit/chat_completions + echo tests/unit/completion_extras + echo tests/unit/containers + echo tests/unit/embeddings + echo tests/unit/endpoints + echo tests/unit/files + echo tests/unit/images + echo tests/unit/interactions + echo tests/unit/messages + echo tests/unit/rag + echo tests/unit/rerank_api + echo tests/unit/vector_stores + echo tests/unit/videos ;; proxy-db-auth-checks) echo tests/unit/proxy/auth/test_auth_checks.py echo tests/unit/proxy/auth/test_user_api_key_auth.py @@ -113,6 +135,7 @@ legacy_paths() { proxy-db-proxy-utils) echo tests/unit/proxy/test_proxy_utils.py ;; proxy-extras) echo tests/unit/litellm_proxy_extras ;; proxy-infra) echo tests/unit/gateway ;; + responses-caching-types) echo tests/unit/types ;; *) echo "unit_selection.sh: unknown flag $1" >&2; exit 1 ;; esac } diff --git a/.circleci/tests.yml b/.circleci/tests.yml index 264d7695a94..10ee19f146a 100644 --- a/.circleci/tests.yml +++ b/.circleci/tests.yml @@ -341,6 +341,7 @@ workflows: flag: - enterprise-package - proxy-infra + - responses-caching-types - proxy-db-auth-checks - proxy-db-jwt-and-keys - proxy-db-proxy-server-core @@ -353,6 +354,13 @@ workflows: - proxy-db-endpoints-and-responses base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-misc + flag: misc + shards: 2 + reruns: 2 + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - unit: name: unit-proxy-db-proxy-utils flag: proxy-db-proxy-utils diff --git a/.github/merge-smoke-tests.json b/.github/merge-smoke-tests.json index 6088953b7eb..8ed7b917460 100644 --- a/.github/merge-smoke-tests.json +++ b/.github/merge-smoke-tests.json @@ -5,8 +5,8 @@ "CHAT-TOOL-STREAM": "tests/test_litellm/llms/openai/test_openai.py::test_acompletion_streams_tool_call_arguments_over_injected_transport", "MODEL-ALLOW": "tests/test_litellm/proxy/auth/test_auth_checks.py::test_can_object_call_model_allows_listed_model_for_key", "MODEL-DENY": "tests/test_litellm/proxy/auth/test_auth_checks.py::test_can_object_call_model_denials_return_forbidden[key-key_model_access_denied]", - "COST-EXPLICIT": "tests/test_litellm/test_cost_calculator.py::test_completion_cost_charges_explicit_per_token_rates_over_registered_ones", - "COST-ZERO": "tests/test_litellm/test_cost_calculator.py::test_completion_cost_is_zero_when_explicit_rates_are_zero", + "COST-EXPLICIT": "tests/unit/test_cost_calculator.py::test_completion_cost_charges_explicit_per_token_rates_over_registered_ones", + "COST-ZERO": "tests/unit/test_cost_calculator.py::test_completion_cost_is_zero_when_explicit_rates_are_zero", "LOG-CONTENT-ON": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_keeps_message_content_when_message_logging_is_on", "LOG-CONTENT-OFF": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_redacts_message_content_when_message_logging_is_off", "CALLBACK-SUCCESS": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_async_success_handler_delivers_standard_logging_payload_to_custom_logger", diff --git a/.github/workflows/test-redis-compat.yml b/.github/workflows/test-redis-compat.yml index 25fb8f8bce3..2f5ce4d441a 100644 --- a/.github/workflows/test-redis-compat.yml +++ b/.github/workflows/test-redis-compat.yml @@ -10,7 +10,7 @@ on: - "litellm/_redis_credential_provider.py" - "litellm/caching/redis_cache.py" - "litellm/caching/evicted_client_closer.py" - - "tests/test_litellm/test_redis.py" + - "tests/unit/test_redis.py" - "tests/local_testing/test_caching.py" - "tests/test_litellm/caching/test_redis_connection_pool.py" - "tests/test_litellm/caching/test_redis_cluster_cache.py" @@ -84,7 +84,7 @@ jobs: run: | redis-server --version uv run --no-sync pytest \ - tests/test_litellm/test_redis.py \ + tests/unit/test_redis.py \ tests/test_litellm/caching/test_redis_connection_pool.py \ tests/test_litellm/caching/test_redis_cluster_cache.py \ tests/test_litellm/caching/test_evicted_client_closer.py \ diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index a60d230d05f..91b54f4ee70 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -52,7 +52,7 @@ jobs: include: - shard: mcp-integration artifact-name: mcp-integration - test-path: "tests/mcp_tests tests/test_litellm/experimental_mcp_client" + test-path: "tests/mcp_tests" unit-flag: mcp-integration workers: 2 reruns: 0 @@ -70,7 +70,6 @@ jobs: - shard: enterprise-routing artifact-name: enterprise-routing test-path: >- - tests/test_litellm/google_genai tests/test_litellm/router_utils tests/test_litellm/router_strategy unit-flag: enterprise-routing @@ -106,26 +105,13 @@ jobs: - shard: misc artifact-name: misc test-path: >- - tests/test_litellm/batches tests/test_litellm/secret_managers - tests/test_litellm/a2a_protocol - tests/test_litellm/chat_completions - tests/test_litellm/completion_extras - tests/test_litellm/containers - tests/test_litellm/endpoints - tests/test_litellm/files - tests/test_litellm/images tests/test_litellm/interactions - tests/test_litellm/messages - tests/test_litellm/embeddings tests/test_litellm/ocr tests/test_litellm/passthrough - tests/test_litellm/rag - tests/test_litellm/rerank_api tests/test_litellm/rust_bridge - tests/test_litellm/vector_stores - tests/test_litellm/videos tests/test_litellm/test_*.py + unit-flag: misc workers: 2 reruns: 2 timeout-minutes: 20 @@ -243,7 +229,7 @@ jobs: test-path: >- tests/test_litellm/responses tests/test_litellm/caching - tests/test_litellm/types + unit-flag: responses-caching-types workers: 2 reruns: 2 timeout-minutes: 20 diff --git a/Makefile b/Makefile index 28daf589a23..62e6ae53275 100644 --- a/Makefile +++ b/Makefile @@ -332,10 +332,10 @@ test-unit-core-utils: install-test-deps $(UV_RUN) pytest tests/test_litellm/litellm_core_utils --tb=short -vv -n 2 --durations=20 test-unit-other: install-test-deps - $(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/test_litellm/secret_managers tests/test_litellm/vector_stores tests/test_litellm/a2a_protocol tests/test_litellm/anthropic_interface tests/test_litellm/completion_extras tests/test_litellm/containers tests/unit/enterprise tests/test_litellm/experimental_mcp_client tests/test_litellm/google_genai tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/test_litellm/types --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/test_litellm/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/test_litellm/anthropic_interface tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20 test-unit-root: install-test-deps - $(UV_RUN) pytest tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/unit/test_*.py tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20 # Proxy unit tests (tests/unit/proxy split alphabetically) test-proxy-unit-a: install-test-deps diff --git a/tests/_vcr_conftest_common.py b/tests/_vcr_conftest_common.py index ab046674eb6..3adc671021b 100644 --- a/tests/_vcr_conftest_common.py +++ b/tests/_vcr_conftest_common.py @@ -52,7 +52,7 @@ from tests._vcr_redis_persister import ( # network call entirely, so skip tests record nothing (NOOP) and passing tests # stop carrying a volatile github episode. This matches the established idiom in # the unit-test suite, which sets the same flag (see e.g. -# tests/test_litellm/test_cost_calculator.py). ``setdefault`` so an explicit +# tests/unit/test_cost_calculator.py). ``setdefault`` so an explicit # override still wins. os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True") diff --git a/tests/code_coverage_tests/code_qa_check_tests.py b/tests/code_coverage_tests/code_qa_check_tests.py index 025f836511c..6c620a02522 100644 --- a/tests/code_coverage_tests/code_qa_check_tests.py +++ b/tests/code_coverage_tests/code_qa_check_tests.py @@ -13,15 +13,16 @@ def check_for_litellm_module_deletion(base_dir): del sys.modules[module] """ problematic_files = [] - test_dir = os.path.join(base_dir, "test_litellm") + candidate_dirs = [os.path.join(base_dir, name) for name in ("test_litellm", "unit")] + test_dirs = [test_dir for test_dir in candidate_dirs if os.path.exists(test_dir)] - if not os.path.exists(test_dir): - print(f"Warning: Directory {test_dir} does not exist.") + if not test_dirs: + print(f"Warning: None of {candidate_dirs} exist.") return [] - print(f"Checking directory: {test_dir}") + print(f"Checking directories: {test_dirs}") - for root, _, files in os.walk(test_dir): + for root, _, files in (entry for test_dir in test_dirs for entry in os.walk(test_dir)): for file in files: if file.endswith(".py"): file_path = os.path.join(root, file) @@ -173,7 +174,7 @@ def main(): f"This can cause import issues and test failures. Files: {problematic_files}" ) else: - print("✓ No litellm module deletion patterns found in test_litellm directory.") + print("✓ No litellm module deletion patterns found in tests/test_litellm or tests/unit.") if __name__ == "__main__": diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 7332a533872..06e5b020836 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -31,7 +31,7 @@ def get_all_functions_called_in_tests(base_dir): specifically in files containing the word 'router'. """ called_functions = set() - test_dirs = ["local_testing", "router_unit_tests", "test_litellm"] + test_dirs = ["local_testing", "router_unit_tests", "test_litellm", "unit"] for test_dir in test_dirs: dir_path = os.path.join(base_dir, test_dir) diff --git a/tests/llm_translation/test_skills_api.py b/tests/llm_translation/test_skills_api.py index aeab5f0da3e..d21e7376ea7 100644 --- a/tests/llm_translation/test_skills_api.py +++ b/tests/llm_translation/test_skills_api.py @@ -277,4 +277,4 @@ class BaseSkillsAPITest(ABC): # # Transformation logic (URL construction, headers, request/response parsing) is # covered by unit tests in: -# tests/test_litellm/test_anthropic_skills_transformation.py +# tests/unit/test_anthropic_skills_transformation.py diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py deleted file mode 100644 index 0b2bfe9d266..00000000000 --- a/tests/test_litellm/batches/test_batch_utils.py +++ /dev/null @@ -1,387 +0,0 @@ -import json - -import pytest - -import litellm -import litellm.batches.batch_utils as bu -from litellm.types.llms.openai import Batch - -GROUNDED_USAGE_METADATA = { - "promptTokenCount": 19, - "candidatesTokenCount": 59, - "thoughtsTokenCount": 406, - "toolUsePromptTokenCount": 73, - "totalTokenCount": 557, - "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 19}], - "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 59}], - "toolUsePromptTokensDetails": [{"modality": "TEXT", "tokenCount": 73}], - "trafficType": "ON_DEMAND", -} -PASSTHROUGH_OUTPUT_URI = ( - "gs://litellm-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/u/" - "predictions.jsonl" -) -UNGROUNDED_USAGE_METADATA = { - "promptTokenCount": 20, - "candidatesTokenCount": 48, - "thoughtsTokenCount": 195, - "toolUsePromptTokenCount": 73, - "totalTokenCount": 336, - "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 20}], - "trafficType": "ON_DEMAND", -} - - -def _batch(output_file_id: str) -> Batch: - return Batch( - id="b", - completion_window="24h", - created_at=1, - endpoint="/v1/chat/completions", - input_file_id="f", - object="batch", - status="completed", - output_file_id=output_file_id, - ) - - -def _vertex_jsonl(rows: list[dict]) -> bytes: - return "\n".join(json.dumps(row) for row in rows).encode() - - -def _vertex_openai_row(custom_id: str, model: str, prompt_tokens: int, completion_tokens: int) -> dict: - return { - "id": f"batch_req_{custom_id}", - "custom_id": custom_id, - "response": { - "status_code": 200, - "request_id": custom_id, - "body": { - "id": f"chatcmpl-{custom_id}", - "object": "chat.completion", - "model": model, - "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], - "usage": { - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "total_tokens": prompt_tokens + completion_tokens, - }, - }, - }, - "error": None, - } - - -def _native_vertex_row(usage_metadata: dict, *, grounded: bool, model_version: str | None = "gemini-2.5-flash"): - candidate = {"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"} - grounding = {"groundingMetadata": {"webSearchQueries": ["q"]}} if grounded else {} - response = {"candidates": [{**candidate, **grounding}], "usageMetadata": usage_metadata} - return { - "request": {"contents": [{"role": "user", "parts": [{"text": "q"}]}], "tools": [{"googleSearch": {}}]}, - "status": "", - "response": {**response, **({"modelVersion": model_version} if model_version else {})}, - "processed_time": "2026-09-23T19:02:00.000+00:00", - } - - -def _capture_cost_calls(monkeypatch, prompt_cost=0.5, completion_cost=0.25) -> list: - import litellm.cost_calculator as cc - - calls: list = [] - - def _calc(**kw): - calls.append(kw) - return (prompt_cost, completion_cost) - - monkeypatch.setattr(cc, "batch_cost_calculator", _calc) - return calls - - -def test_vertex_native_cost_bills_embedding_rows(monkeypatch): - monkeypatch.setitem(litellm.model_cost, "vertex_ai/gemini-embedding-2", {"input_cost_per_token_batches": 1e-7}) - rows = [ - { - "key": "id_1", - "status": "", - "request": {"content": {"parts": [{"text": "hello world"}]}}, - "response": {"embedding": {"values": [0.1, 0.2]}, "usageMetadata": {"promptTokenCount": 2}}, - }, - { - "key": "id_2", - "status": "", - "request": {"content": {"parts": [{"text": "hello"}]}}, - "response": {"embedding": {"values": [0.3]}, "tokenCount": "3"}, - }, - {"key": "id_3", "status": "INVALID_ARGUMENT", "request": {"content": {"parts": [{"text": ""}]}}}, - ] - - result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-embedding-2") - - assert (result.successful_requests, result.failed_requests) == (2, 1) - assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (5, 0, 5) - assert result.cost == pytest.approx(5 * 1e-7) - assert result.models == ["gemini-embedding-2"] - - -@pytest.mark.asyncio -async def test_native_vertex_rows_route_to_vertex_cost_path_without_flag(monkeypatch): - monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) - monkeypatch.setattr( - bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") - ) - calls = _capture_cost_calls(monkeypatch) - rows = [ - _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), - _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False), - ] - - result = await bu.calculate_batch_cost_and_usage( - file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" - ) - - assert result.cost == pytest.approx(1.5) - assert (result.successful_requests, result.failed_requests) == (2, 0) - assert result.models == ["gemini-2.5-flash"] - assert {(call["model"], call["custom_llm_provider"]) for call in calls} == {("gemini-2.5-flash", "vertex_ai")} - - -@pytest.mark.asyncio -async def test_openai_shaped_vertex_rows_keep_the_generic_path_without_flag(monkeypatch): - monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) - monkeypatch.setattr( - bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run") - ) - _capture_cost_calls(monkeypatch) - rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)] - - result = await bu.calculate_batch_cost_and_usage( - file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" - ) - - assert result.successful_requests == 1 - - -@pytest.mark.asyncio -async def test_native_vertex_rows_on_another_provider_keep_the_generic_path(monkeypatch): - monkeypatch.setattr( - bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run") - ) - _capture_cost_calls(monkeypatch) - - result = await bu.calculate_batch_cost_and_usage( - file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], - custom_llm_provider="openai", - ) - - assert result.successful_requests == 0 - - -@pytest.mark.asyncio -async def test_handle_completed_batch_routes_native_rows_without_flag(monkeypatch): - monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) - raw_rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)] - - async def fake_fetch(batch, custom_llm_provider, litellm_params=None): - return _vertex_jsonl(raw_rows) - - monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) - monkeypatch.setattr( - bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") - ) - calls = _capture_cost_calls(monkeypatch, prompt_cost=0.7, completion_cost=0.3) - deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} - - result = await bu._handle_completed_batch( - _batch(PASSTHROUGH_OUTPUT_URI), - custom_llm_provider="vertex_ai", - model_name="gemini-2.5-flash", - model_info=deployment_model_info, - ) - - assert result.cost == pytest.approx(1.0) - assert result.usage.total_tokens == 557 - assert [call["model_info"] for call in calls] == [deployment_model_info] - - -def test_native_vertex_usage_is_billed_like_the_online_path(monkeypatch): - calls = _capture_cost_calls(monkeypatch) - grounded = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True) - ungrounded = _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False) - - result = bu.calculate_vertex_ai_batch_cost_and_usage([grounded, ungrounded], "gemini-2.5-flash") - - grounded_usage, ungrounded_usage = (call["usage"] for call in calls) - assert grounded_usage.prompt_tokens == 19 - assert grounded_usage.completion_tokens == 59 + 406 - assert grounded_usage.completion_tokens_details.reasoning_tokens == 406 - assert ungrounded_usage.prompt_tokens == 20 + 73 - assert ungrounded_usage.completion_tokens == 48 + 195 - assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( - 19 + 93, - 465 + 243, - 557 + 336, - ) - - -def test_native_vertex_rows_are_priced_by_model_version_without_a_model_name(monkeypatch): - calls = _capture_cost_calls(monkeypatch) - rows = [ - _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"), - _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-pro"), - _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None), - ] - - result = bu.calculate_vertex_ai_batch_cost_and_usage(rows) - - assert [call["model"] for call in calls] == ["gemini-2.5-flash", "gemini-2.5-pro"] - assert result.models == ["gemini-2.5-flash", "gemini-2.5-pro"] - assert result.cost == pytest.approx(1.5) - assert result.successful_requests == 3 - assert result.usage.total_tokens == 557 + 336 + 336 - - -def test_native_vertex_rows_without_usage_metadata_count_as_failed(monkeypatch): - _capture_cost_calls(monkeypatch) - rows = [ - {"request": {"contents": []}, "status": "Error: bad request", "processed_time": "t"}, - {"request": {"contents": []}, "response": {"candidates": []}}, - _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), - ] - - result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") - - assert (result.successful_requests, result.failed_requests) == (1, 2) - assert result.usage.total_tokens == 557 - - -def test_native_vertex_batch_whose_rows_all_failed_still_names_the_deployment_model(monkeypatch): - calls = _capture_cost_calls(monkeypatch) - rows = [{"request": {"contents": []}, "status": "Error: quota exceeded", "processed_time": "t"}] * 2 - - result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") - - assert result.models == ["gemini-2.5-flash"] - assert (result.successful_requests, result.failed_requests, result.cost) == (0, 2, 0.0) - assert calls == [] - - -def test_native_vertex_rows_are_priced_with_the_deployment_model_info(monkeypatch): - calls = _capture_cost_calls(monkeypatch) - deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} - - bu.calculate_vertex_ai_batch_cost_and_usage( - [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], - "gemini-2.5-flash", - model_info=deployment_model_info, - ) - - assert [call["model_info"] for call in calls] == [deployment_model_info] - - -@pytest.mark.asyncio -async def test_native_vertex_rows_keep_the_deployment_model_info_through_the_batch_entrypoint(monkeypatch): - calls = _capture_cost_calls(monkeypatch) - deployment_model_info = {"input_cost_per_token_batches": 1e-6} - - await bu.calculate_batch_cost_and_usage( - file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], - custom_llm_provider="vertex_ai", - model_name="gemini-2.5-flash", - model_info=deployment_model_info, - ) - - assert [call["model_info"] for call in calls] == [deployment_model_info] - - -def test_native_vertex_rows_are_priced_by_the_deployment_model_over_model_version(monkeypatch): - calls = _capture_cost_calls(monkeypatch) - rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-pro")] - - result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") - - assert [call["model"] for call in calls] == ["gemini-2.5-flash"] - assert result.models == ["gemini-2.5-flash"] - - -def test_native_vertex_rows_that_fail_response_validation_count_as_failed(monkeypatch): - calls = _capture_cost_calls(monkeypatch) - rows = [ - {"request": {"contents": []}, "response": {"candidates": "nope", "usageMetadata": GROUNDED_USAGE_METADATA}}, - _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), - ] - - result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") - - assert (result.successful_requests, result.failed_requests) == (1, 1) - assert result.usage.total_tokens == 557 - assert len(calls) == 1 - - -@pytest.mark.parametrize("wildcard_model", ["*", "vertex_ai/*"]) -def test_native_vertex_rows_under_a_wildcard_deployment_are_priced_by_model_version(monkeypatch, wildcard_model): - calls = _capture_cost_calls(monkeypatch) - rows = [ - _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"), - _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None), - ] - - result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, wildcard_model) - - assert [call["model"] for call in calls] == ["gemini-2.5-flash", wildcard_model] - assert result.cost == pytest.approx(1.5) - assert (result.successful_requests, result.failed_requests) == (2, 0) - assert result.usage.total_tokens == 557 + 336 - - -def test_native_vertex_row_without_model_version_under_a_wildcard_deployment_bills_its_explicit_prices(): - deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} - with_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash") - without_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version=None) - - twin = bu.calculate_vertex_ai_batch_cost_and_usage([with_version], "vertex_ai/*", model_info=deployment_model_info) - both = bu.calculate_vertex_ai_batch_cost_and_usage( - [with_version, without_version], "vertex_ai/*", model_info=deployment_model_info - ) - - assert twin.cost > 0 - assert both.cost == pytest.approx(2 * twin.cost) - assert (both.successful_requests, both.failed_requests) == (2, 0) - - -def test_native_vertex_row_the_cost_map_cannot_price_is_billed_at_zero_and_the_rest_still_bills(monkeypatch): - import litellm.cost_calculator as cc - - def _calc(**kw): - if kw["model"] == "gemini-unpriced": - raise ValueError("no pricing") - return (0.5, 0.25) - - monkeypatch.setattr(cc, "batch_cost_calculator", _calc) - rows = [ - _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-unpriced"), - _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-flash"), - ] - - result = bu.calculate_vertex_ai_batch_cost_and_usage(rows) - - assert result.cost == pytest.approx(0.75) - assert (result.successful_requests, result.failed_requests) == (2, 0) - assert result.usage.total_tokens == 557 + 336 - assert result.models == ["gemini-unpriced", "gemini-2.5-flash"] - - -@pytest.mark.asyncio -async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_known(monkeypatch): - monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False) - monkeypatch.setattr( - bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") - ) - calls = _capture_cost_calls(monkeypatch) - rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)] - - result = await bu.calculate_batch_cost_and_usage( - file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" - ) - - assert calls == [] - assert (result.successful_requests, result.failed_requests) == (0, 1) diff --git a/tests/test_litellm/chat_completions/test_dispatch.py b/tests/test_litellm/chat_completions/test_dispatch.py deleted file mode 100644 index ddb6e827309..00000000000 --- a/tests/test_litellm/chat_completions/test_dispatch.py +++ /dev/null @@ -1,117 +0,0 @@ -from __future__ import annotations - -from collections.abc import Mapping -from typing import Final - -import pytest - -import litellm -from litellm.chat_completions import dispatch -from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, RouteRule, Rules -from litellm.rust_bridge.chat_completions.entrypoints import ( - LiteLLMChatCompletionsRequest, - NativeAcompletion, - NativeCompletion, -) -from litellm.rust_bridge.configuration import Rollout -from litellm.types.utils import ModelResponse - -MESSAGES: Final = [{"role": "user", "content": "hi"}] - - -@pytest.mark.asyncio -async def test_public_completion_calls_keep_the_python_result() -> None: - sync_response: Final = litellm.completion(model="openai/test-model", messages=MESSAGES, mock_response="ok") - async_response: Final = await litellm.acompletion(model="openai/test-model", messages=MESSAGES, mock_response="ok") - - assert isinstance(sync_response, ModelResponse) - assert isinstance(async_response, ModelResponse) - assert sync_response.choices[0].message.content == "ok" - assert async_response.choices[0].message.content == "ok" - - -def test_sync_completion_request_projects_public_arguments() -> None: - rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),) - expected: Final = ModelResponse() - - def native( - request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object] - ) -> ModelResponse: - assert request.model == "test-model" - assert request.messages == MESSAGES - assert request.custom_llm_provider == "openai" - assert request.stream is True - return expected - - binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None) - binding.override(native) - response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - ("test-model", MESSAGES), - {"custom_llm_provider": "openai", "stream": True}, - python=lambda *args, **kwargs: pytest.fail("required native route must handle this call"), - binding=binding, - native=lambda hook, request, args, kwargs: hook(request, args, kwargs), - rules=rules, - ) - - assert response is expected - - -@pytest.mark.asyncio -async def test_async_completion_falls_back_after_native_declines() -> None: - from litellm.rust_bridge.bindings import native_exception_types - - native_types: Final = native_exception_types() - if native_types is None: - pytest.skip("native bridge is unavailable") - declined, _ = native_types - expected: Final = ModelResponse() - rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_OPT_OUT),) - - async def native( - request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object] - ) -> ModelResponse: - raise declined("unsupported") - - async def python(*args: object, **kwargs: object) -> ModelResponse: - return expected - - binding: Final[NativeBinding[NativeAcompletion]] = NativeBinding("acompletion", validate=lambda _: None) - binding.override(native) - response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - ("test-model", MESSAGES), - {}, - python=python, - binding=binding, - native=lambda hook, request, args, kwargs: hook(request, args, kwargs), - rules=rules, - ) - - assert response is expected - - -def test_internal_acompletion_marker_bypasses_native() -> None: - rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),) - expected: Final = ModelResponse() - - def python(*args: object, **kwargs: object) -> ModelResponse: - return expected - - def native( - request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object] - ) -> ModelResponse: - pytest.fail("acompletion's inner completion call must stay on Python") - - binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None) - binding.override(native) - response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - ("test-model", MESSAGES), - {"custom_llm_provider": "openai", "acompletion": True}, - python=python, - binding=binding, - native=lambda hook, request, args, kwargs: hook(request, args, kwargs), - rules=rules, - ) - - assert response is expected diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index beca10d5555..f8c7d5273d1 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -14,6 +14,7 @@ from pathlib import Path from types import SimpleNamespace import httpx import pytest +from pytest_socket import _remove_restrictions import asyncio @@ -509,6 +510,14 @@ def setup_and_teardown(): print(f"[conftest] Module teardown complete (worker: {worker_id or 'master'})") +def pytest_collectstart(): + _remove_restrictions() + + +def pytest_runtest_setup(): + _remove_restrictions() + + def pytest_collection_modifyitems(config, items): """ Customize test collection order. diff --git a/tests/test_litellm/interactions/test_litellm_responses_bridge.py b/tests/test_litellm/interactions/test_litellm_responses_bridge.py index 8400f2c4840..17e7f9fc4ff 100644 --- a/tests/test_litellm/interactions/test_litellm_responses_bridge.py +++ b/tests/test_litellm/interactions/test_litellm_responses_bridge.py @@ -7,10 +7,6 @@ the litellm_responses bridge provider, which calls litellm.responses() internall import os -from litellm.interactions.litellm_responses_transformation.transformation import ( - LiteLLMResponsesInteractionsConfig, -) -from litellm.types.interactions import Turn from tests.test_litellm.interactions.base_interactions_test import ( BaseInteractionsTest, ) @@ -30,71 +26,3 @@ class TestLiteLLMResponsesBridge(BaseInteractionsTest): def get_api_key(self) -> str: """Return the OpenAI API key from environment.""" return os.getenv("OPENAI_API_KEY", "") - - -class TestBridgeInputTransformation: - """Regression tests for translating Interactions input into Responses API input. - - The bridge used to pass Google content parts through raw ({"type": "text"}), - which the Responses API rejects with a 400, and it dropped the role encoded - in step types and in the legacy "model" turn role. - """ - - def test_step_input_maps_roles_and_content_types(self): - transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( - [ - {"type": "user_input", "content": [{"type": "text", "text": "I like apples."}]}, - {"type": "model_output", "content": [{"type": "text", "text": "I like oranges."}]}, - {"type": "user_input", "content": [{"type": "text", "text": "What did you say?"}]}, - ] - ) - assert transformed == [ - {"role": "user", "content": [{"type": "input_text", "text": "I like apples."}]}, - {"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]}, - {"role": "user", "content": [{"type": "input_text", "text": "What did you say?"}]}, - ] - - def test_legacy_turn_input_maps_model_role_to_assistant(self): - transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( - [ - {"role": "user", "content": [{"type": "text", "text": "I like apples."}]}, - {"role": "model", "content": [{"type": "text", "text": "I like oranges."}]}, - ] - ) - assert transformed == [ - {"role": "user", "content": [{"type": "input_text", "text": "I like apples."}]}, - {"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]}, - ] - - def test_turn_pydantic_model_with_string_content(self): - transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( - [Turn(role="model", content="I like oranges.")] - ) - assert transformed == [ - {"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]} - ] - - def test_string_input_passes_through(self): - transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input("Hello") - assert transformed == "Hello" - - def test_content_list_input_becomes_single_user_message(self): - transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( - [{"type": "text", "text": "Hello"}, "world"] - ) - assert transformed == [ - { - "role": "user", - "content": [ - {"type": "input_text", "text": "Hello"}, - {"type": "input_text", "text": "world"}, - ], - } - ] - - def test_non_text_content_passes_through_unchanged(self): - image_part = {"type": "image", "data": "base64data", "mime_type": "image/png"} - transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( - [{"type": "user_input", "content": [image_part]}] - ) - assert transformed == [{"role": "user", "content": [image_part]}] diff --git a/tests/test_litellm/messages/__init__.py b/tests/test_litellm/messages/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/messages/test_dispatch.py b/tests/test_litellm/messages/test_dispatch.py deleted file mode 100644 index 4da060f809a..00000000000 --- a/tests/test_litellm/messages/test_dispatch.py +++ /dev/null @@ -1,155 +0,0 @@ -from __future__ import annotations - -from collections.abc import Mapping -from typing import Final - -import pytest -from pydantic import TypeAdapter - -import litellm -from litellm.messages import dispatch -from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, RouteRule, Rules -from litellm.rust_bridge.configuration import Rollout -from litellm.rust_bridge.messages.entrypoints import ( - LiteLLMMessagesRequest, - NativeAmessages, - NativeMessages, -) -from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse - -MESSAGES: Final = [{"role": "user", "content": "hi"}] - - -@pytest.mark.asyncio -async def test_public_anthropic_messages_keeps_the_python_result() -> None: - response: Final = await litellm.anthropic_messages( - model="anthropic/claude-sonnet-4-5", messages=MESSAGES, max_tokens=10, mock_response="ok" - ) - - assert isinstance(response, dict) - content: Final = TypeAdapter(list[dict[str, object]]).validate_python(response.get("content", [])) - assert content[0]["text"] == "ok" - - -def test_sync_messages_request_projects_public_arguments() -> None: - rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),) - expected: Final = AnthropicMessagesResponse(model="claude-test") - - def native( - request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object] - ) -> AnthropicMessagesResponse: - assert request.model == "claude-test" - assert request.messages == MESSAGES - assert request.max_tokens == 10 - assert request.custom_llm_provider == "anthropic" - return expected - - binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None) - binding.override(native) - response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - (), - { - "model": "claude-test", - "messages": MESSAGES, - "max_tokens": 10, - "custom_llm_provider": "anthropic", - }, - python=lambda *args, **kwargs: pytest.fail("required native route must handle this call"), - binding=binding, - native=lambda hook, request, args, kwargs: hook(request, args, kwargs), - rules=rules, - ) - - assert response is expected - - -def test_messages_binding_error_delegates_unchanged_to_python() -> None: - rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),) - expected: Final = AnthropicMessagesResponse(model="claude-test") - - def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse: - return expected - - def native( - request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object] - ) -> AnthropicMessagesResponse: - pytest.fail("a call without max_tokens cannot project a request and must stay on Python") - - binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None) - binding.override(native) - response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - (), - {"model": "claude-test", "messages": MESSAGES, "custom_llm_provider": "anthropic"}, - python=python, - binding=binding, - native=lambda hook, request, args, kwargs: hook(request, args, kwargs), - rules=rules, - ) - - assert response is expected - - -@pytest.mark.asyncio -async def test_async_messages_falls_back_after_native_declines() -> None: - from litellm.rust_bridge.bindings import native_exception_types - - native_types: Final = native_exception_types() - if native_types is None: - pytest.skip("native bridge is unavailable") - declined, _ = native_types - expected: Final = AnthropicMessagesResponse(model="claude-test") - rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_OUT),) - - async def native( - request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object] - ) -> AnthropicMessagesResponse: - raise declined("unsupported") - - async def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse: - return expected - - binding: Final[NativeBinding[NativeAmessages]] = NativeBinding("amessages", validate=lambda _: None) - binding.override(native) - response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - (), - {"model": "claude-test", "messages": MESSAGES, "max_tokens": 10}, - python=python, - binding=binding, - native=lambda hook, request, args, kwargs: hook(request, args, kwargs), - rules=rules, - ) - - assert response is expected - - -def test_internal_is_async_marker_bypasses_native() -> None: - rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),) - expected: Final = AnthropicMessagesResponse(model="claude-test") - - def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse: - return expected - - def native( - request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object] - ) -> AnthropicMessagesResponse: - pytest.fail("anthropic_messages' inner handler call must stay on Python") - - binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None) - binding.override(native) - response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - (), - { - "model": "claude-test", - "messages": MESSAGES, - "max_tokens": 10, - "custom_llm_provider": "anthropic", - "is_async": True, - }, - python=python, - binding=binding, - native=lambda hook, request, args, kwargs: hook(request, args, kwargs), - rules=rules, - ) - - assert response is expected diff --git a/tests/test_litellm/rag/__init__.py b/tests/test_litellm/rag/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/rag/ingestion/__init__.py b/tests/test_litellm/rag/ingestion/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/rerank_api/__init__.py b/tests/test_litellm/rerank_api/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/router_strategy/test_router_tag_routing.py b/tests/test_litellm/router_strategy/test_router_tag_routing.py index 16c641b8d29..e4b8860a7a6 100644 --- a/tests/test_litellm/router_strategy/test_router_tag_routing.py +++ b/tests/test_litellm/router_strategy/test_router_tag_routing.py @@ -2823,7 +2823,7 @@ def test_update_router_config_schema_includes_tag_routing_prefix(): # UpdateRouterConfig before calling update_settings; a field missing here # causes model_dump(exclude_none=True) to silently drop it before # update_settings is ever called -- the same bug shape LIT-3152 fixed for - # retry_policy (see tests/test_litellm/test_router_retry_policy_update.py). + # retry_policy (see tests/unit/test_router_retry_policy_update.py). from litellm.types.router import UpdateRouterConfig config = UpdateRouterConfig(tag_routing_prefix="route:") diff --git a/tests/test_litellm/test_compression.py b/tests/test_litellm/test_compression.py index 4fbcd4ed30d..997778d0a1b 100644 --- a/tests/test_litellm/test_compression.py +++ b/tests/test_litellm/test_compression.py @@ -3,20 +3,13 @@ Unit tests for litellm.compress(). """ import os -import importlib import pytest import litellm -from litellm.compression.scoring.bm25 import bm25_score_messages -from litellm.compression.scoring.embedding_scorer import embedding_score_messages -from litellm.compression.content_detection import detect_content_type -from litellm.compression.message_stubbing import extract_key, stub_message -from litellm.compression.retrieval_tool import build_retrieval_tool from litellm.types.utils import CallTypes CALL_TYPE = CallTypes.completion -ANTHROPIC_CALL_TYPE = CallTypes.anthropic_messages # --------------------------------------------------------------------------- @@ -24,420 +17,26 @@ ANTHROPIC_CALL_TYPE = CallTypes.anthropic_messages # --------------------------------------------------------------------------- -def test_bm25_relevance_ranking(): - query = "Fix the authentication bug in the login handler" - messages = [ - { - "role": "user", - "content": "def login_handler(): authentication check bug fix", - }, - {"role": "user", "content": "def render_template(name): css styling layout"}, - {"role": "user", "content": "def verify(): authentication token bug handler"}, - ] - scores = bm25_score_messages(query, messages) - # Messages sharing query terms should score higher than unrelated ones - assert scores[0] > scores[1] - assert scores[2] > scores[1] - - -def test_bm25_empty_query(): - scores = bm25_score_messages("", [{"role": "user", "content": "hello"}]) - assert scores == [0.0] - - -def test_bm25_empty_messages(): - scores = bm25_score_messages("query", []) - assert scores == [] - - -def test_bm25_empty_content(): - scores = bm25_score_messages("query", [{"role": "user", "content": ""}]) - assert scores == [0.0] - - # --------------------------------------------------------------------------- # Content detection # --------------------------------------------------------------------------- -def test_detect_code(): - code = """ -import os -from pathlib import Path - -def main(): - class Foo: - pass - return Foo() -""" - assert detect_content_type(code) == "code" - - -def test_detect_json(): - assert detect_content_type('{"key": "value", "num": 42}') == "json" - assert detect_content_type("[1, 2, 3]") == "json" - - -def test_detect_text(): - assert detect_content_type("This is a plain text paragraph about dogs.") == "text" - - -def test_detect_empty(): - assert detect_content_type("") == "text" - - # --------------------------------------------------------------------------- # Message stubbing # --------------------------------------------------------------------------- -def test_extract_key_with_filename(): - msg = {"role": "user", "content": "# auth.py\ndef authenticate():\n pass"} - used: set = set() - key = extract_key(msg, fallback_index=0, used_keys=used) - assert key == "auth.py" - - -def test_extract_key_fallback(): - msg = {"role": "user", "content": "Some random content without a filename"} - used: set = set() - key = extract_key(msg, fallback_index=5, used_keys=used) - assert key == "message_5" - - -def test_extract_key_duplicates(): - used: set = set() - msg = {"role": "user", "content": "# auth.py\ncode here"} - k1 = extract_key(msg, fallback_index=0, used_keys=used) - k2 = extract_key(msg, fallback_index=1, used_keys=used) - assert k1 == "auth.py" - assert k2 == "auth.py_2" - - -def test_stub_message(): - msg = {"role": "user", "content": "line1\nline2\nline3"} - stubbed = stub_message(msg, "test_key") - assert stubbed["role"] == "user" - assert "test_key" in stubbed["content"] - assert "litellm_content_retrieve" in stubbed["content"] - assert "3 lines" in stubbed["content"] - - # --------------------------------------------------------------------------- # Retrieval tool # --------------------------------------------------------------------------- -def test_retrieval_tool_schema(): - tool = build_retrieval_tool(["auth.py", "utils.py"]) - assert tool["type"] == "function" - assert tool["function"]["name"] == "litellm_content_retrieve" - assert "key" in tool["function"]["parameters"]["properties"] - assert tool["function"]["parameters"]["properties"]["key"]["enum"] == [ - "auth.py", - "utils.py", - ] - assert tool["function"]["parameters"]["required"] == ["key"] - - -def test_retrieval_tool_description_lists_keys(): - tool = build_retrieval_tool(["foo.py", "bar.js"]) - desc = tool["function"]["description"] - assert "foo.py" in desc - assert "bar.js" in desc - - # --------------------------------------------------------------------------- # compress() — end-to-end # --------------------------------------------------------------------------- -def test_compress_below_trigger_passthrough(): - messages = [{"role": "user", "content": "hello"}] - result = litellm.compress(messages, model="gpt-4o", call_type=CALL_TYPE) - assert result["messages"] == messages - assert result["cache"] == {} - assert result["tools"] == [] - assert result["compression_ratio"] == 0.0 - assert result["original_tokens"] == result["compressed_tokens"] - - -def test_compress_above_trigger(): - big_messages = [ - {"role": "system", "content": "You are a coding assistant."}, - { - "role": "user", - "content": "# auth.py\n" + "def authenticate():\n pass\n" * 2000, - }, - { - "role": "user", - "content": "# utils.py\n" + "def helper():\n pass\n" * 2000, - }, - { - "role": "user", - "content": "# readme.md\n" + "This is documentation. " * 2000, - }, - {"role": "user", "content": "Fix the bug in auth.py"}, - ] - - result = litellm.compress( - big_messages, - model="gpt-4o", - call_type=CALL_TYPE, - compression_trigger=1000, - compression_target=500, - ) - - assert result["compressed_tokens"] < result["original_tokens"] - assert result["compression_ratio"] > 0 - assert len(result["cache"]) > 0 - assert len(result["tools"]) == 1 - assert result["tools"][0]["function"]["name"] == "litellm_content_retrieve" - - -def test_compress_anthropic_list_content_is_boundary_stable(): - messages = [ - {"role": "system", "content": [{"type": "text", "text": "System prompt"}]}, - { - "role": "user", - "content": [ - {"type": "text", "text": "# a.py\n" + "alpha " * 2000}, - { - "type": "image_url", - "image_url": {"url": "https://example.com/a.png"}, - }, - ], - }, - { - "role": "user", - "content": [ - {"type": "text", "text": "# b.py\n" + "beta " * 2000}, - { - "type": "image_url", - "image_url": {"url": "https://example.com/b.png"}, - }, - ], - }, - { - "role": "user", - "content": [{"type": "text", "text": "Fix alpha bug in a.py"}], - }, - ] - - result = litellm.compress( - messages=messages, - model="claude-sonnet-4-20250514", - call_type=ANTHROPIC_CALL_TYPE, - compression_trigger=1000, - compression_target=500, - ) - - assert result["compressed_tokens"] < result["original_tokens"] - assert len(result["messages"]) == len(messages) - assert [m["role"] for m in result["messages"]] == [m["role"] for m in messages] - assert len(result["cache"]) > 0 - assert len(result["tools"]) == 1 - assert result["tools"][0]["type"] == "custom" - assert result["tools"][0]["name"] == "litellm_content_retrieve" - assert "input_schema" in result["tools"][0] - - -def test_compress_preserves_system_message(): - messages = [ - {"role": "system", "content": "System prompt. " * 500}, - {"role": "user", "content": "Large file content. " * 5000}, - {"role": "user", "content": "Fix the bug"}, - ] - result = litellm.compress( - messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 - ) - assert result["messages"][0]["role"] == "system" - assert "System prompt" in result["messages"][0]["content"] - - -def test_compress_preserves_last_user_message(): - messages = [ - {"role": "user", "content": "Big context " * 5000}, - {"role": "user", "content": "Fix the bug in auth.py"}, - ] - result = litellm.compress( - messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 - ) - last_user = [m for m in result["messages"] if m["role"] == "user"][-1] - assert "Fix the bug in auth.py" in last_user["content"] - - -def test_compress_preserves_last_assistant_message(): - messages = [ - {"role": "user", "content": "Big context " * 5000}, - {"role": "assistant", "content": "I'll help with that. " * 2000}, - {"role": "user", "content": "Now fix the bug"}, - ] - result = litellm.compress( - messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 - ) - assistant_msgs = [m for m in result["messages"] if m["role"] == "assistant"] - assert len(assistant_msgs) >= 1 - # The last assistant message should be preserved (not stubbed) - last_assistant = assistant_msgs[-1] - assert "I'll help with that" in last_assistant["content"] - - -def test_cache_keys_match_stubs(): - messages = [ - {"role": "user", "content": "# auth.py\n" + "code " * 5000}, - {"role": "user", "content": "Fix it"}, - ] - result = litellm.compress( - messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 - ) - if result["tools"]: - tool_desc = result["tools"][0]["function"]["description"] - for key in result["cache"]: - assert key in tool_desc - - -def test_compress_default_target(): - """compression_target defaults to compression_trigger // 2.""" - messages = [ - {"role": "user", "content": "content " * 5000}, - {"role": "user", "content": "query"}, - ] - result = litellm.compress( - messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=2000 - ) - # Should have compressed — target = 1000 - assert result["compressed_tokens"] <= result["original_tokens"] - - -def test_compress_nested_tool_result_extracts_text_only(): - messages = [ - {"role": "system", "content": [{"type": "text", "text": "System rules"}]}, - { - "role": "user", - "content": [ - {"type": "text", "text": "prefix"}, - { - "type": "tool_result", - "tool_use_id": "toolu_1", - "content": [ - {"type": "text", "text": "nested text fragment"}, - { - "type": "image_url", - "image_url": { - "url": "https://example.com/secret-tool.png", - }, - }, - ], - }, - { - "type": "image_url", - "image_url": {"url": "https://example.com/top.png"}, - }, - {"type": "text", "text": " " + ("irrelevant " * 3000)}, - ], - }, - { - "role": "user", - "content": [{"type": "text", "text": "final query that must remain"}], - }, - ] - - result = litellm.compress( - messages=messages, - model="claude-sonnet-4-20250514", - call_type=ANTHROPIC_CALL_TYPE, - compression_trigger=500, - compression_target=100, - ) - - cached_text = " ".join(result["cache"].values()) - assert "nested text fragment" in cached_text - assert "https://example.com/secret-tool.png" not in cached_text - assert "https://example.com/top.png" not in cached_text - - -def test_compress_default_call_type_is_completion(): - result = litellm.compress( - messages=[ - {"role": "user", "content": "Large context " * 4000}, - {"role": "user", "content": "query"}, - ], - model="gpt-4o", - compression_trigger=1000, - compression_target=500, - ) - - assert result["compressed_tokens"] <= result["original_tokens"] - assert isinstance(result["tools"], list) - - -def test_compress_forwards_embedding_model_params(monkeypatch): - captured = {} - - def fake_embedding_score_messages( - query, messages, model, cache=None, embedding_model_params=None - ): - captured["query"] = query - captured["model"] = model - captured["embedding_model_params"] = embedding_model_params - return [0.0] * len(messages) - - monkeypatch.setattr( - "litellm.compression.scoring.embedding_scorer.embedding_score_messages", - fake_embedding_score_messages, - ) - - result = litellm.compress( - messages=[ - {"role": "user", "content": "Authentication code " * 2000}, - {"role": "user", "content": "Fix auth"}, - ], - model="gpt-4o", - call_type=CALL_TYPE, - compression_trigger=1000, - embedding_model="text-embedding-3-small", - embedding_model_params={"api_base": "https://example-embeddings.test"}, - ) - - assert result["compressed_tokens"] <= result["original_tokens"] - assert captured["model"] == "text-embedding-3-small" - assert captured["embedding_model_params"] == { - "api_base": "https://example-embeddings.test" - } - - -def test_embedding_scorer_forwards_embedding_model_params(monkeypatch): - captured = {} - - class _MockResponse: - data = [ - {"embedding": [1.0, 0.0]}, - {"embedding": [1.0, 0.0]}, - {"embedding": [0.0, 1.0]}, - ] - - def fake_embedding(**kwargs): - captured.update(kwargs) - return _MockResponse() - - monkeypatch.setattr(litellm, "embedding", fake_embedding) - - scores = embedding_score_messages( - query="auth", - messages=[ - {"role": "user", "content": "auth code"}, - {"role": "user", "content": "cooking recipe"}, - ], - model="text-embedding-3-small", - embedding_model_params={"api_base": "https://example-embeddings.test"}, - ) - - assert len(scores) == 2 - assert captured["model"] == "text-embedding-3-small" - assert captured["api_base"] == "https://example-embeddings.test" - - # --------------------------------------------------------------------------- # Embedding scorer — integration test (skipped without API key) # --------------------------------------------------------------------------- @@ -458,210 +57,3 @@ def test_embedding_scorer(): ) assert result["compression_ratio"] > 0 assert len(result["cache"]) > 0 - - -@pytest.mark.parametrize( - "final_user_message, expected_content", - [ - ("How to cook?", "Unrelated cooking recipes "), - ("Fix auth", "Authentication code "), - ], -) -def test_simple_compression(final_user_message, expected_content): - messages = [ - {"role": "user", "content": "Authentication code " * 2000}, - {"role": "user", "content": "Unrelated cooking recipes " * 2000}, - {"role": "user", "content": final_user_message}, - ] - result = litellm.compress( - messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 - ) - if expected_content == "Unrelated cooking recipes ": - assert "Unrelated cooking recipes " in result["messages"][1]["content"] - assert "Authentication code " not in result["messages"][0]["content"] - elif expected_content == "Authentication code ": - assert "Authentication code " in result["messages"][0]["content"] - assert "Unrelated cooking recipes " not in result["messages"][1]["content"] - else: - raise ValueError(f"Unexpected expected_content: {expected_content}") - - -def test_compress_anthropic_drops_irrelevant_tool_exchange_span(monkeypatch): - compress_module = importlib.import_module("litellm.compression.compress") - - def fake_bm25_score_messages(query, messages): - assert "final query" in query - assert len(messages) == 5 - # Prefer idx=0 and de-prioritize the tool exchange span (idx=1,2) - return [0.95, 0.01, 0.02, 0.8, 1.0] - - def fake_token_counter(model, messages=None, text=None): - if messages is not None: - return 1000 - if text is None: - return 0 - if "final query" in text: - return 50 - if "assistant_tail" in text: - return 20 - if "other_blob" in text: - return 220 - if "tool_payload_relevant" in text: - return 200 - if text == "": - return 1 - return 10 - - monkeypatch.setattr( - compress_module, "bm25_score_messages", fake_bm25_score_messages - ) - monkeypatch.setattr(compress_module, "token_counter", fake_token_counter) - - messages = [ - {"role": "user", "content": "other_blob " * 300}, - { - "role": "assistant", - "content": [ - { - "type": "tool_use", - "id": "toolu_drop", - "name": "litellm_content_retrieve", - "input": {"key": "message_1"}, - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "tool_result", - "tool_use_id": "toolu_drop", - "content": [{"type": "text", "text": "tool_payload_relevant"}], - } - ], - }, - {"role": "assistant", "content": "assistant_tail"}, - {"role": "user", "content": "final query"}, - ] - - result = litellm.compress( - messages=messages, - model="claude-sonnet-4-20250514", - call_type=ANTHROPIC_CALL_TYPE, - compression_trigger=100, - compression_target=280, - ) - - # idx=1,2 should be dropped atomically (no orphan tool blocks left behind) - assert len(result["messages"]) == 3 - assert result["messages"][0]["role"] == "user" - assert "other_blob" in result["messages"][0]["content"] - assert result["messages"][1]["content"] == "assistant_tail" - assert result["messages"][2]["content"] == "final query" - assert result["cache"] == {} - - -def test_compress_anthropic_keeps_relevant_tool_exchange_span(monkeypatch): - compress_module = importlib.import_module("litellm.compression.compress") - - def fake_bm25_score_messages(query, messages): - assert "final query" in query - assert len(messages) == 5 - # Prefer the tool exchange span over idx=0 - return [0.05, 0.01, 0.92, 0.8, 1.0] - - def fake_token_counter(model, messages=None, text=None): - if messages is not None: - return 1000 - if text is None: - return 0 - if "final query" in text: - return 50 - if "assistant_tail" in text: - return 20 - if "other_blob" in text: - return 220 - if "tool_payload_relevant" in text: - return 200 - if text == "": - return 1 - return 10 - - monkeypatch.setattr( - compress_module, "bm25_score_messages", fake_bm25_score_messages - ) - monkeypatch.setattr(compress_module, "token_counter", fake_token_counter) - - messages = [ - {"role": "user", "content": "other_blob " * 300}, - { - "role": "assistant", - "content": [ - { - "type": "tool_use", - "id": "toolu_keep", - "name": "litellm_content_retrieve", - "input": {"key": "message_1"}, - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "tool_result", - "tool_use_id": "toolu_keep", - "content": [{"type": "text", "text": "tool_payload_relevant"}], - } - ], - }, - {"role": "assistant", "content": "assistant_tail"}, - {"role": "user", "content": "final query"}, - ] - - result = litellm.compress( - messages=messages, - model="claude-sonnet-4-20250514", - call_type=ANTHROPIC_CALL_TYPE, - compression_trigger=100, - compression_target=280, - ) - - assert len(result["messages"]) == 5 - assert result["messages"][1]["role"] == "assistant" - assert result["messages"][2]["role"] == "user" - # idx=0 should be compressed instead - assert "litellm_content_retrieve" in result["messages"][0]["content"] - assert len(result["cache"]) == 1 - - -def test_compress_anthropic_malformed_tool_sequence_passes_through(): - messages = [ - {"role": "user", "content": "other_blob " * 300}, - { - "role": "assistant", - "content": [ - { - "type": "tool_use", - "id": "toolu_broken", - "name": "litellm_content_retrieve", - "input": {"key": "message_1"}, - } - ], - }, - {"role": "user", "content": [{"type": "text", "text": "missing tool_result"}]}, - {"role": "user", "content": "final query"}, - ] - - result = litellm.compress( - messages=messages, - model="claude-sonnet-4-20250514", - call_type=ANTHROPIC_CALL_TYPE, - compression_trigger=100, - compression_target=280, - ) - - assert result["messages"] == messages - assert result["cache"] == {} - assert result["tools"] == [] - assert result["compression_skipped_reason"] == "invalid_anthropic_tool_sequence" diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 227fb48bb08..78728d6fd58 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -1,31 +1,12 @@ -import asyncio -import base64 -from datetime import datetime -import contextlib -import copy import json -import logging import os -from collections.abc import Mapping -from dataclasses import dataclass -from typing import Final -import httpx import pytest -import respx -from fastapi.testclient import TestClient -import urllib.parse -from importlib import import_module from unittest.mock import MagicMock, patch import litellm -from litellm import main as litellm_main -from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage async def _async_fake_bedrock_image_details(image_url): @@ -61,111 +42,6 @@ def add_api_keys_to_env(monkeypatch): monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) -@pytest.fixture -def openai_api_response(): - mock_response_data = { - "id": "chatcmpl-B0W3vmiM78Xkgx7kI7dr7PC949DMS", - "choices": [ - { - "finish_reason": "stop", - "index": 0, - "logprobs": None, - "message": { - "content": "", - "refusal": None, - "role": "assistant", - "audio": None, - "function_call": None, - "tool_calls": None, - }, - } - ], - "created": 1739462947, - "model": "gpt-4o-mini-2024-07-18", - "object": "chat.completion", - "service_tier": "default", - "system_fingerprint": "fp_bd83329f63", - "usage": { - "completion_tokens": 1, - "prompt_tokens": 121, - "total_tokens": 122, - "completion_tokens_details": { - "accepted_prediction_tokens": 0, - "audio_tokens": 0, - "reasoning_tokens": 0, - "rejected_prediction_tokens": 0, - }, - "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, - }, - } - - return mock_response_data - - -def test_completion_missing_role(openai_api_response): - from openai import OpenAI - - from litellm.types.utils import ModelResponse - - client = OpenAI(api_key="test_api_key") - - mock_raw_response = MagicMock() - mock_raw_response.headers = { - "x-request-id": "123", - "openai-organization": "org-123", - "x-ratelimit-limit-requests": "100", - "x-ratelimit-remaining-requests": "99", - } - mock_raw_response.parse.return_value = ModelResponse(**openai_api_response) - - print(f"openai_api_response: {openai_api_response}") - - with patch.object( - client.chat.completions.with_raw_response, "create", MagicMock(return_value=mock_raw_response) - ) as mock_create: - litellm.completion( - model="gpt-4o-mini", - messages=[ - {"role": "user", "content": "Hey"}, - { - "content": "", - "tool_calls": [ - { - "id": "call_m0vFJjQmTH1McvaHBPR2YFwY", - "function": { - "arguments": '{"input": "dksjsdkjdhskdjshdskhjkhlk"}', - "name": "tool_name", - }, - "type": "function", - "index": 0, - }, - { - "id": "call_Vw6RaqV2n5aaANXEdp5pYxo2", - "function": { - "arguments": '{"input": "jkljlkjlkjlkjlk"}', - "name": "tool_name", - }, - "type": "function", - "index": 1, - }, - { - "id": "call_hBIKwldUEGlNh6NlSXil62K4", - "function": { - "arguments": '{"input": "jkjlkjlkjlkj;lj"}', - "name": "tool_name", - }, - "type": "function", - "index": 2, - }, - ], - }, - ], - client=client, - ) - - mock_create.assert_called_once() - - @pytest.mark.parametrize( "model", [ @@ -277,210 +153,6 @@ async def test_url_with_format_param(model, sync_mode, monkeypatch): assert "jpeg" not in json_str -@pytest.mark.parametrize("model", ["gpt-4o-mini"]) -@pytest.mark.parametrize("sync_mode", [True, False]) -@pytest.mark.asyncio -async def test_url_with_format_param_openai(model, sync_mode): - from openai import AsyncOpenAI, OpenAI - - from litellm import acompletion, completion - - if sync_mode: - client = OpenAI() - else: - client = AsyncOpenAI() - - args = { - "model": model, - "messages": [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "https://awsmp-logos.s3.amazonaws.com/seller-xw5kijmvmzasy/c233c9ade2ccb5491072ae232c814942.png", - "format": "image/png", - }, - }, - {"type": "text", "text": "Describe this image"}, - ], - } - ], - } - with patch.object( - client.chat.completions.with_raw_response, "create" - ) as mock_client: - try: - if sync_mode: - response = completion(**args, client=client) - else: - response = await acompletion(**args, client=client) - print(response) - except Exception as e: - print(e) - - mock_client.assert_called() - - print(mock_client.call_args.kwargs) - - json_str = json.dumps(mock_client.call_args.kwargs) - - assert "format" not in json_str - - -def test_bedrock_latency_optimized_inference(): - from litellm.llms.custom_httpx.http_handler import HTTPHandler - - client = HTTPHandler() - with patch.object(client, "post") as mock_post: - try: - response = litellm.completion( - model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - messages=[{"role": "user", "content": "Hello, how are you?"}], - performanceConfig={"latency": "optimized"}, - client=client, - ) - except Exception as e: - print(e) - - mock_post.assert_called_once() - json_data = json.loads(mock_post.call_args.kwargs["data"]) - assert json_data["performanceConfig"]["latency"] == "optimized" - - -@pytest.mark.parametrize( - ("custom_llm_provider", "model", "expected"), - [ - ("anthropic", "claude-sonnet-5", True), - ("bedrock", "us.anthropic.claude-sonnet-5-20260501-v1:0", True), - ("bedrock", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", True), - ("bedrock", "us.amazon.nova-2-lite-v1:0", False), - ("vertex_ai", "claude-sonnet-5", True), - ("vertex_ai", "gemini-3.8-flash", False), - ("azure_ai", "claude-sonnet-4-6", True), - ("azure_ai", "gpt-5.6", False), - ("openai", "gpt-5.6", False), - ("gemini", "gemini-3.8-flash", False), - ], -) -def test_is_claude_tool_target(custom_llm_provider: str, model: str, expected: bool): - assert litellm_main._is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model) is expected - - -@pytest.mark.parametrize("key", ["input_examples", "eager_input_streaming"]) -def test_drop_anthropic_only_tool_keys_strips_tool_and_function_levels(key: str): - tools = [ - {"type": "function", "name": "example_tool", key: True, "function": {"name": "example_tool", key: True}}, - "opaque_tool", - ] - - cleaned = litellm_main._drop_anthropic_only_tool_keys(tools=tools) - - assert cleaned == [ - {"type": "function", "name": "example_tool", "function": {"name": "example_tool"}}, - "opaque_tool", - ] - assert tools[0][key] is True - assert tools[0]["function"][key] is True - - -def test_completion_strips_eager_input_streaming_before_openai(respx_mock: respx.MockRouter, openai_api_response): - api_base: Final = "http://localhost:12346/v1" - mock_route: Final = respx_mock.post(url__regex=rf"{api_base}/chat/completions.*").mock( - return_value=httpx.Response(status_code=200, json=openai_api_response) - ) - - litellm.completion( - model="openai/gpt-5.6", - messages=[{"role": "user", "content": "Write the file"}], - tools=[ - { - "type": "function", - "function": {"name": "write_file", "parameters": {"type": "object", "properties": {}}}, - "eager_input_streaming": True, - } - ], - api_base=api_base, - api_key="fake_openai_api_key", - ) - - assert mock_route.called - sent_tool: Final = json.loads(respx_mock.calls[0].request.content)["tools"][0] - assert "eager_input_streaming" not in sent_tool - assert sent_tool["function"]["name"] == "write_file" - - -def test_custom_provider_with_extra_headers(): - from litellm.llms.custom_httpx.http_handler import HTTPHandler - - with patch.object( - litellm.llms.custom_httpx.http_handler.HTTPHandler, "post" - ) as mock_post: - response = litellm.completion( - model="custom/custom", - messages=[{"role": "user", "content": "Hello, how are you?"}], - headers={"X-Custom-Header": "custom-value"}, - api_base="https://example.com/api/v1", - ) - - mock_post.assert_called_once() - assert mock_post.call_args[1]["headers"]["X-Custom-Header"] == "custom-value" - - -def test_custom_provider_with_extra_body(): - from litellm.llms.custom_httpx.http_handler import HTTPHandler - - with patch.object( - litellm.llms.custom_httpx.http_handler.HTTPHandler, "post" - ) as mock_post: - response = litellm.completion( - model="custom/custom", - messages=[{"role": "user", "content": "Hello, how are you?"}], - extra_body={ - "X-Custom-BodyValue": "custom-value", - "X-Custom-BodyValue2": "custom-value2", - }, - api_base="https://example.com/api/v1", - ) - mock_post.assert_called_once() - - assert mock_post.call_args[1]["json"]["X-Custom-BodyValue"] == "custom-value" - assert mock_post.call_args[1]["json"] == { - "model": "custom", - "params": { - "prompt": ["Hello, how are you?"], - "max_tokens": None, - "temperature": None, - "top_p": None, - "top_k": None, - }, - "X-Custom-BodyValue": "custom-value", - "X-Custom-BodyValue2": "custom-value2", - } - - # test that extra_body is not passed if not provided - with patch.object( - litellm.llms.custom_httpx.http_handler.HTTPHandler, "post" - ) as mock_post: - response = litellm.completion( - model="custom/custom", - messages=[{"role": "user", "content": "Hello, how are you?"}], - api_base="https://example.com/api/v1", - ) - mock_post.assert_called_once() - assert mock_post.call_args[1]["json"] == { - "model": "custom", - "params": { - "prompt": ["Hello, how are you?"], - "max_tokens": None, - "temperature": None, - "top_p": None, - "top_k": None, - }, - } - - @pytest.fixture(autouse=True) def set_openrouter_api_key(): original_api_key = os.environ.get("OPENROUTER_API_KEY") @@ -490,3753 +162,3 @@ def set_openrouter_api_key(): os.environ["OPENROUTER_API_KEY"] = original_api_key else: del os.environ["OPENROUTER_API_KEY"] - - -@pytest.mark.asyncio -async def test_extra_body_with_fallback( - respx_mock: respx.MockRouter, set_openrouter_api_key, monkeypatch -): - """ - test regression for https://github.com/BerriAI/litellm/issues/8425. - - This was perhaps a wider issue with the acompletion function not passing kwargs such as extra_body correctly when fallbacks are specified. - """ - - # Save original state to restore after test - original_disable_aiohttp = litellm.disable_aiohttp_transport - - try: - # since this uses respx, we need to set use_aiohttp_transport to False - # Set both the global variable and environment variable to ensure it takes effect - litellm.disable_aiohttp_transport = True - monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") - # Flush cache to ensure no stale aiohttp clients are used - litellm.in_memory_llm_clients_cache.flush_cache() - - # Set up test parameters - model = "openrouter/deepseek/deepseek-chat" - messages = [{"role": "user", "content": "Hello, world!"}] - extra_body = { - "provider": { - "order": ["DeepSeek"], - "allow_fallbacks": False, - "require_parameters": True, - } - } - fallbacks = [{"model": "openrouter/google/gemini-flash-1.5-8b"}] - - # Set up mock to respond to any POST request to the OpenRouter endpoint - # This ensures it works for both primary and fallback models - mock_route = respx_mock.post("https://openrouter.ai/api/v1/chat/completions") - mock_route.return_value = httpx.Response( - 200, - json={ - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677652288, - "model": model, - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": "Hello from mocked response!", - }, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 9, - "completion_tokens": 12, - "total_tokens": 21, - }, - }, - ) - - response = await litellm.acompletion( - model=model, - messages=messages, - extra_body=extra_body, - fallbacks=fallbacks, - api_key="fake-openrouter-api-key", - ) - - # Verify the response - assert response is not None - assert ( - len(respx_mock.calls) > 0 - ), "Mock was not called - check if aiohttp transport is properly disabled" - - # Get the request from the mock - request: httpx.Request = respx_mock.calls[0].request - request_body = request.read() - request_body = json.loads(request_body) - - # Verify basic parameters - assert request_body["model"] == "deepseek/deepseek-chat" - assert request_body["messages"] == messages - - # Verify the extra_body parameters remain under the provider key - assert request_body["provider"]["order"] == ["DeepSeek"] - assert request_body["provider"]["allow_fallbacks"] is False - assert request_body["provider"]["require_parameters"] is True - finally: - # Restore original state to prevent test pollution - litellm.disable_aiohttp_transport = original_disable_aiohttp - litellm.in_memory_llm_clients_cache.flush_cache() - - -@pytest.mark.parametrize("env_base", ["OPENAI_BASE_URL", "OPENAI_API_BASE"]) -@pytest.mark.asyncio -@pytest.mark.flaky(retries=3, delay=1) -async def test_openai_env_base( - respx_mock: respx.MockRouter, env_base, openai_api_response, monkeypatch -): - "This tests OpenAI env variables are honored, including legacy OPENAI_API_BASE" - # Ensure aiohttp transport is disabled to use httpx which respx can mock - litellm.disable_aiohttp_transport = True - - expected_base_url = "http://localhost:12345/v1" - - # Assign the environment variable based on env_base, and use a fake API key. - monkeypatch.setenv(env_base, expected_base_url) - monkeypatch.setenv("OPENAI_API_KEY", "fake_openai_api_key") - - model = "gpt-4o" - messages = [{"role": "user", "content": "Hello, how are you?"}] - - # Configure respx mock to intercept the request - mock_route = respx_mock.post( - url__regex=r"http://localhost:12345/v1/chat/completions.*" - ).mock( - return_value=httpx.Response( - status_code=200, - json={ - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677652288, - "model": model, - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": "Hello from mocked response!", - }, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 9, - "completion_tokens": 12, - "total_tokens": 21, - }, - }, - ) - ) - - try: - response = await litellm.acompletion(model=model, messages=messages) - - # verify we had a response - assert response.choices[0].message.content == "Hello from mocked response!" - - # Verify the mock was called - assert ( - mock_route.called - ), "Mock route was not called - request may have bypassed respx" - finally: - # Clean up to avoid affecting other tests - litellm.disable_aiohttp_transport = False - - -def build_database_url(username, password, host, dbname): - username_enc = urllib.parse.quote_plus(username) - password_enc = urllib.parse.quote_plus(password) - dbname_enc = urllib.parse.quote_plus(dbname) - return f"postgresql://{username_enc}:{password_enc}@{host}/{dbname_enc}" - - -def test_build_database_url(): - url = build_database_url("user@name", "p@ss:word", "localhost", "db/name") - assert url == "postgresql://user%40name:p%40ss%3Aword@localhost/db%2Fname" - - -def test_bedrock_llama(): - litellm._turn_on_debug() - from litellm.types.utils import CallTypes - from litellm.utils import return_raw_request - - model = "bedrock/invoke/us.meta.llama4-scout-17b-instruct-v1:0" - - request = return_raw_request( - endpoint=CallTypes.completion, - kwargs={ - "model": model, - "messages": [ - {"role": "user", "content": "hi"}, - ], - }, - ) - print(request) - - assert ( - request["raw_request_body"]["prompt"] - == "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\nhi<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" - ) - - -def _mocked_openai_chat_response(model: str) -> httpx.Response: - return httpx.Response( - status_code=200, - json={ - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677652288, - "model": model, - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": "Hello from mocked response!", - }, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 9, - "completion_tokens": 12, - "total_tokens": 21, - }, - }, - ) - - -def test_return_raw_request_does_not_call_provider(respx_mock: respx.MockRouter): - """Regression for #33952: return_raw_request must transform without contacting the provider. - - Previously return_raw_request invoked the real endpoint with a fake key and relied on the - provider rejecting it, which sent an unintended inference request and (in the async proxy - route) blocked the event loop on provider I/O. - """ - from litellm.types.utils import CallTypes - from litellm.utils import return_raw_request - - model = "gpt-4o" - route = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - return_value=_mocked_openai_chat_response(model) - ) - - request = return_raw_request( - endpoint=CallTypes.completion, - kwargs={ - "model": model, - "messages": [{"role": "user", "content": "hi"}], - }, - ) - - assert route.call_count == 0 - assert request.get("error") is None - assert request["raw_request_body"]["model"] == model - assert request["raw_request_body"]["messages"] == [ - {"role": "user", "content": "hi"} - ] - - -def test_completion_forwards_verbosity_in_raw_request(respx_mock: respx.MockRouter): - """Regression test: completion() must forward the verbosity param to the provider request body.""" - from litellm.types.utils import CallTypes - from litellm.utils import return_raw_request - - model = "gpt-5.2" - messages = [{"role": "user", "content": "hi"}] - respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - return_value=_mocked_openai_chat_response(model) - ) - - request = return_raw_request( - endpoint=CallTypes.completion, - kwargs={ - "model": model, - "messages": messages, - "verbosity": "high", - }, - ) - - assert request["raw_request_body"]["verbosity"] == "high" - assert request["raw_request_body"]["model"] == model - assert request["raw_request_body"]["messages"] == messages - - -@pytest.mark.asyncio -async def test_acompletion_forwards_verbosity_to_provider_request( - respx_mock: respx.MockRouter, monkeypatch -): - """Regression test: acompletion() must forward the verbosity param to the provider request body.""" - original_disable_aiohttp = litellm.disable_aiohttp_transport - try: - litellm.disable_aiohttp_transport = True - monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") - litellm.in_memory_llm_clients_cache.flush_cache() - - model = "gpt-5.2" - messages = [{"role": "user", "content": "hi"}] - mock_route = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - return_value=_mocked_openai_chat_response(model) - ) - - response = await litellm.acompletion( - model=model, - messages=messages, - verbosity="low", - api_key="fake-openai-api-key", - ) - - assert response.choices[0].message.content == "Hello from mocked response!" - assert mock_route.called - request_body = json.loads(respx_mock.calls[0].request.read()) - assert request_body["verbosity"] == "low" - assert request_body["model"] == model - assert request_body["messages"] == messages - finally: - litellm.disable_aiohttp_transport = original_disable_aiohttp - litellm.in_memory_llm_clients_cache.flush_cache() - - -def test_responses_api_bridge_check_strips_responses_prefix(): - """Test that responses_api_bridge_check strips 'responses/' prefix and sets mode.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 4096} - - model_info, model = responses_api_bridge_check( - model="responses/gpt-4-responses", - custom_llm_provider="openai", - ) - - assert model == "gpt-4-responses" - assert model_info["mode"] == "responses" - - -def test_responses_api_bridge_check_gpt_5_4_pro(): - """Test that gpt-5.4-pro routes through responses API bridge, not chat completions. - - Regression test for https://github.com/BerriAI/litellm/issues/23014 - gpt-5.4-pro is a responses-only model and must not be sent to /v1/chat/completions. - """ - from litellm.main import responses_api_bridge_check - - for model_name in ["gpt-5.4-pro", "gpt-5.4-pro-2026-03-05"]: - model_info, model = responses_api_bridge_check( - model=model_name, - custom_llm_provider="openai", - ) - assert ( - model_info.get("mode") == "responses" - ), f"{model_name} should have mode='responses', got '{model_info.get('mode')}'" - - -def test_responses_api_bridge_check_gpt_5_4_tools_plus_reasoning_routes_to_responses(): - """gpt-5.4 with both tools and reasoning_effort should route to Responses API.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.4", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort="xhigh", - ) - - assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_gpt_6_astra_tools_with_default_reasoning_routes_to_responses(): - from litellm.main import responses_api_bridge_check - - model_info, model = responses_api_bridge_check( - model="gpt-6-astra", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - ) - - assert model == "gpt-6-astra" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_gpt_5_5_tools_plus_reasoning_routes_to_responses(): - """gpt-5.5+ with both tools and reasoning_effort should route to Responses API.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.5-pro", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort="xhigh", - ) - - assert model == "gpt-5.5-pro" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_azure_gpt_5_4_tools_plus_reasoning_routes_to_responses(): - """Azure gpt-5.4 with both tools and reasoning_effort should route to Responses API.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.4", - custom_llm_provider="azure", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort="high", - ) - - assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_azure_gpt_5_4_tools_with_default_reasoning_routes_to_responses(): - """ - Azure gpt-5.4 with tools and UNSET reasoning_effort must bridge: OpenAI enables - reasoning by default for gpt-5.4+, and Chat Completions rejects function tools - whenever reasoning is on. - """ - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.4", - custom_llm_provider="azure", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - ) - - assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_gpt_5_4_tools_with_default_reasoning_routes_to_responses(): - """ - gpt-5.4 with tools and UNSET reasoning_effort must bridge: OpenAI enables reasoning - by default for gpt-5.4+, and Chat Completions rejects function tools whenever - reasoning is on ("use /v1/responses or set reasoning_effort to 'none'"). - """ - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.4", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - ) - - assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" - - -@pytest.mark.parametrize( - "model_name, expected_mode", - [ - pytest.param("gpt-5.6-sol", "responses", id="above-boundary-bridges"), - pytest.param("gpt-5.1", None, id="below-boundary-stays-chat"), - ], -) -def test_responses_api_bridge_check_gpt_5_6_tools_with_default_reasoning_routes_to_responses( - monkeypatch, model_name, expected_mode -): - """ - gpt-5.6 must bridge on function tools alone. The bridge used to require an explicit - reasoning_effort, so a gpt-5.6 call carrying tools and no effort was rejected with - "Function tools with reasoning_effort are not supported for gpt-5.6-sol in - /v1/chat/completions". - - Paired with a model below the gpt-5.4 boundary, which must still stay on chat. The - gate parses the version and drops any suffix, so the family members bridge - identically and only the boundary distinguishes behaviour. - """ - import litellm - from litellm.main import responses_api_bridge_check - - monkeypatch.delenv("OPENAI_BASE_URL", raising=False) - monkeypatch.delenv("OPENAI_API_BASE", raising=False) - monkeypatch.setattr(litellm, "api_base", None) - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model=model_name, - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - ) - - assert model == model_name - assert model_info.get("mode") == expected_mode - - -def test_responses_api_bridge_check_gpt_5_4_tools_with_reasoning_none_stays_chat(): - """ - Explicit reasoning_effort "none" is OpenAI's documented escape hatch that keeps - function tools servable on Chat Completions; the bridge must not fire. - """ - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.4", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort="none", - ) - - assert model == "gpt-5.4" - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_reasoning_none_with_summary_still_routes_to_responses(): - """A reasoning summary is Responses-only regardless of effort value.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.4", - custom_llm_provider="openai", - reasoning_effort="none", - reasoning_summary="detailed", - ) - - assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_gpt_5_4_custom_tools_only_stays_chat(): - """ - Chat Completions serves custom (grammar) tools natively with reasoning on; only - FUNCTION tools trigger the OpenAI rejection. Custom-only requests must stay on chat - so responses keep the native custom tool_call shape instead of the bridge's - function-shaped mapping. - """ - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "custom", "custom": {"name": "ApplyPatch", "description": "V4A patch"}}], - reasoning_effort=None, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_gpt_5_4_mixed_function_and_custom_tools_routes_to_responses(): - """One function tool in the mix is enough to make chat unservable with reasoning on.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[ - {"type": "custom", "custom": {"name": "ApplyPatch"}}, - {"type": "function", "function": {"name": "shell"}}, - ], - reasoning_effort=None, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_gpt_5_4_flat_function_tool_routes_to_responses(): - """Responses-style flat function tool defs still count as function tools.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "name": "shell", "parameters": {"type": "object"}}], - reasoning_effort=None, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") == "responses" - - -@pytest.mark.parametrize( - "custom_llm_provider, model_name, api_base", - [ - pytest.param("openai", "gpt-5.6", None, id="openai"), - pytest.param("azure_ai", "gpt-6-astra", "https://myproject.services.ai.azure.com", id="azure-ai-foundry"), - ], -) -def test_responses_api_bridge_check_function_tool_without_body_stays_chat( - monkeypatch, custom_llm_provider, model_name, api_base -): - import litellm - from litellm.main import responses_api_bridge_check - - monkeypatch.delenv("OPENAI_BASE_URL", raising=False) - monkeypatch.delenv("OPENAI_API_BASE", raising=False) - monkeypatch.setattr(litellm, "api_base", None) - - model_info, model = responses_api_bridge_check( - model=model_name, - custom_llm_provider=custom_llm_provider, - tools=[{"type": "function"}], - reasoning_effort=None, - api_base=api_base, - ) - - assert model == model_name - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_dict_effort_none_stays_chat(): - """The escape hatch must honor litellm's dict form: {"effort": "none"} means reasoning off.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort={"effort": "none"}, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_dict_effort_active_routes_to_responses(): - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort={"effort": "low"}, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_dict_effort_none_with_summary_routes_to_responses(): - """A summary inside the dict form is Responses-only even when effort is none.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort={"effort": "none", "summary": "concise"}, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") == "responses" - - -@pytest.mark.parametrize("blank_api_base", [None, "", " ", "\t"]) -def test_responses_api_bridge_check_blank_api_base_is_default_openai(blank_api_base): - """ - A blank api_base (None, empty, or whitespace) resolves to the default OpenAI - endpoint downstream, which enforces the reasoning+tools constraint, so gpt-5.4+ - function-tool requests with unset reasoning_effort must still auto-bridge. - """ - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - api_base=blank_api_base, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_custom_api_base_with_unset_effort_stays_chat(): - """ - Chat-only OpenAI-compatible backends registered under the openai provider with a - custom api_base and gpt-5.4+ model names serve tools-without-reasoning fine and - have no /responses route; the unset-effort arm must not reroute them. - """ - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - api_base="http://vllm.internal:8000/v1", - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_custom_api_base_via_global_with_unset_effort_stays_chat(monkeypatch): - """ - A custom base set through the litellm.api_base global (not the call arg) is resolved the - same way the chat handler resolves it, so the unset-effort arm must not reroute a chat-only - backend to a /responses route it lacks. Regression guard: the gate previously inspected only - the call-level api_base and bridged these requests. - """ - import litellm - from litellm.main import responses_api_bridge_check - - monkeypatch.setattr(litellm, "api_base", "http://vllm.internal:8000/v1") - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - api_base=None, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") != "responses" - - -@pytest.mark.parametrize("env_var", ["OPENAI_BASE_URL", "OPENAI_API_BASE"]) -def test_responses_api_bridge_check_custom_api_base_via_env_with_unset_effort_stays_chat(monkeypatch, env_var): - """ - A custom base set via OPENAI_BASE_URL/OPENAI_API_BASE env is resolved identically to the chat - handler, so the unset-effort arm leaves the request on chat instead of bridging it. - """ - import litellm - from litellm.main import responses_api_bridge_check - - monkeypatch.setattr(litellm, "api_base", None) - monkeypatch.delenv("OPENAI_BASE_URL", raising=False) - monkeypatch.delenv("OPENAI_API_BASE", raising=False) - monkeypatch.setenv(env_var, "http://vllm.internal:8000/v1") - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - api_base=None, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") != "responses" - - -@pytest.mark.parametrize( - "api_base", - [ - "https://southcentralus.privatelink.api.openai.com/v1", - "https://privatelink.corp.api.openai.com/v1", - "https://api.openai.com:443/v1", - "https://api.openai.com/v1/", - "HTTPS://API.OPENAI.COM/v1", - ], -) -def test_responses_api_bridge_check_openai_backed_custom_api_base_with_unset_effort_routes_to_responses(api_base): - """ - A custom api_base whose host is api.openai.com or a subdomain of it (a PrivateLink hostname, a - port-qualified or trailing-slash default) still reaches the real OpenAI backend, which rejects - function tools with reasoning on Chat Completions, so the unset-effort arm must bridge exactly as - it does for the literal default URL. Regression guard for GH #39353. - """ - from litellm.main import responses_api_bridge_check - - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - api_base=api_base, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") == "responses" - - -@pytest.mark.parametrize( - "api_base", - [ - "https://api.openai.com.evil.example/v1", - "https://notapi.openai.com/v1", - "https://gateway.example/v1?upstream=api.openai.com", - "https://openai.internal.example/api.openai.com/v1", - ], -) -def test_responses_api_bridge_check_lookalike_custom_api_base_with_unset_effort_stays_chat(api_base): - """Only the host decides: api.openai.com appearing elsewhere in the URL is still a foreign backend.""" - from litellm.main import responses_api_bridge_check - - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - api_base=api_base, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_privatelink_api_base_via_env_with_unset_effort_routes_to_responses(monkeypatch): - """A PrivateLink base set through OPENAI_BASE_URL resolves the way the chat handler's does and still bridges.""" - import litellm - from litellm.main import responses_api_bridge_check - - monkeypatch.setattr(litellm, "api_base", None) - monkeypatch.delenv("OPENAI_API_BASE", raising=False) - monkeypatch.setenv("OPENAI_BASE_URL", "https://southcentralus.privatelink.api.openai.com/v1") - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - api_base=None, - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_custom_api_base_with_explicit_effort_still_routes(): - """Explicit reasoning_effort keeps its pre-existing bridging behavior on any api_base.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.6", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort="high", - api_base="http://vllm.internal:8000/v1", - ) - - assert model == "gpt-5.6" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_azure_with_api_base_and_unset_effort_routes(): - """Azure OpenAI always sets api_base and does enforce the constraint; keep bridging.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.4", - custom_llm_provider="azure", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - api_base="https://myresource.openai.azure.com", - ) - - assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" - - -_FOUNDRY_API_BASE: Final = "https://myproject.services.ai.azure.com" -_FOUNDRY_FUNCTION_TOOL: Final = ({"type": "function", "function": {"name": "get_weather"}},) - - -@pytest.mark.parametrize( - "model_name, api_base, reasoning_effort", - [ - pytest.param("gpt-6-astra", _FOUNDRY_API_BASE, None, id="gpt-6-unset-effort"), - pytest.param("gpt-6-astra", _FOUNDRY_API_BASE, "low", id="gpt-6-explicit-effort"), - pytest.param("gpt-6-astra", "https://myresource.openai.azure.com", None, id="gpt-6-azure-openai-host"), - pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, "low", id="gpt-5.6-explicit-effort"), - pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, {"effort": "high"}, id="gpt-5.6-explicit-effort-dict"), - ], -) -def test_responses_api_bridge_check_azure_ai_foundry_rejected_tools_route_to_responses( - model_name, api_base, reasoning_effort -): - from litellm.main import responses_api_bridge_check - - model_info, model = responses_api_bridge_check( - model=model_name, - custom_llm_provider="azure_ai", - tools=_FOUNDRY_FUNCTION_TOOL, - reasoning_effort=reasoning_effort, - api_base=api_base, - ) - - assert model == model_name - assert model_info.get("mode") == "responses" - - -@pytest.mark.parametrize( - "model_name, api_base, reasoning_effort", - [ - pytest.param("gpt-6-astra", _FOUNDRY_API_BASE, "none", id="explicit-none-stays-chat"), - pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, None, id="gpt-5.6-unset-effort-stays-chat"), - pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, "none", id="gpt-5.6-explicit-none-stays-chat"), - pytest.param("gpt-5.5", _FOUNDRY_API_BASE, "high", id="gpt-5.5-explicit-effort-stays-chat"), - pytest.param("gpt-5.4-mini", _FOUNDRY_API_BASE, None, id="gpt-5.4-mini-unset-effort-stays-chat"), - pytest.param("gpt-5.4-mini", _FOUNDRY_API_BASE, "low", id="gpt-5.4-mini-explicit-effort-stays-chat"), - pytest.param("gpt-6-astra", "https://myproject.models.ai.azure.com", None, id="serverless-host-stays-chat"), - pytest.param("Mistral-large-2411", _FOUNDRY_API_BASE, None, id="non-gpt-5-model-stays-chat"), - pytest.param("claude-opus-4-1", _FOUNDRY_API_BASE, None, id="claude-on-foundry-stays-chat"), - ], -) -def test_responses_api_bridge_check_azure_ai_without_foundry_responses_route_stays_chat( - model_name, api_base, reasoning_effort -): - from litellm.main import responses_api_bridge_check - - model_info, model = responses_api_bridge_check( - model=model_name, - custom_llm_provider="azure_ai", - tools=_FOUNDRY_FUNCTION_TOOL, - reasoning_effort=reasoning_effort, - api_base=api_base, - ) - - assert model == model_name - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_older_gpt_5_tools_without_reasoning_stays_chat(): - """Pre-5.4 GPT-5 names keep the old boundary: tools alone never bridge.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.1", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort=None, - ) - - assert model == "gpt-5.1" - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_gpt_5_4_reasoning_summary_without_tools_routes_to_responses(): - """gpt-5.4+ with reasoning_effort + reasoningSummary but no tools should bridge (AI SDK).""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5.4", - custom_llm_provider="openai", - tools=None, - reasoning_effort="medium", - reasoning_summary="auto", - ) - - assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_gpt_5_reasoning_summary_routes_to_responses(): - """Bare ``gpt-5`` with reasoning_effort + reasoningSummary should bridge (not 5.4+).""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5", - custom_llm_provider="openai", - tools=None, - reasoning_effort="medium", - reasoning_summary="auto", - ) - - assert model == "gpt-5" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_gpt_5_tools_without_summary_stays_chat(): - """gpt-5 with tools + reasoning_effort but no summary should stay on chat.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-5", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - reasoning_effort="medium", - reasoning_summary=None, - ) - - assert model == "gpt-5" - assert model_info.get("mode") != "responses" - - -@patch("litellm.completion_extras.responses_api_bridge.completion") -def test_gpt_5_4_responses_bridge_preserves_reasoning_summary_dict( - mock_responses_completion, -): - """When routed to Responses, preserve reasoning_effort summary dict.""" - mock_responses_completion.return_value = MagicMock() - - import litellm - - litellm.completion( - model="gpt-5.4", - messages=[{"role": "user", "content": "What is the capital of France?"}], - tools=[ - { - "type": "function", - "function": { - "name": "get_capital", - "description": "Get the capital of a country", - "parameters": { - "type": "object", - "properties": {"country": {"type": "string"}}, - }, - }, - } - ], - reasoning_effort={"effort": "xhigh", "summary": "detailed"}, - api_key="fake-key", - ) - - assert mock_responses_completion.called is True - optional_params = mock_responses_completion.call_args.kwargs["optional_params"] - assert optional_params["reasoning_effort"] == { - "effort": "xhigh", - "summary": "detailed", - } - - -@pytest.mark.parametrize("reasoning_effort", ["high", {"effort": "high"}]) -def test_responses_bridge_preserves_reasoning_effort_with_drop_params( - reasoning_effort, - restore_model_registry, - respx_mock: respx.MockRouter, - monkeypatch: pytest.MonkeyPatch, -): - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - response_body: Final = { - "id": "resp_test", - "object": "response", - "created_at": 1734366691, - "status": "completed", - "model": "test-responses-bridge", - "output": [ - { - "type": "message", - "id": "msg_1", - "status": "completed", - "role": "assistant", - "content": [{"type": "output_text", "text": "Done.", "annotations": []}], - } - ], - "parallel_tool_calls": True, - "usage": { - "input_tokens": 1, - "output_tokens": 1, - "total_tokens": 2, - "output_tokens_details": {"reasoning_tokens": 0}, - }, - "error": None, - "incomplete_details": None, - "instructions": None, - "metadata": None, - "temperature": None, - "tool_choice": "auto", - "tools": [], - "top_p": None, - "max_output_tokens": None, - "previous_response_id": None, - "reasoning": None, - "truncation": None, - "user": None, - } - response_route: Final = respx_mock.post("https://api.perplexity.ai/v1/responses").respond(json=response_body) - model: Final = "perplexity/test-responses-bridge" - litellm.register_model( - { - model: { - "litellm_provider": "perplexity", - "mode": "responses", - "supports_reasoning": False, - "input_cost_per_token": 0.0, - "output_cost_per_token": 0.0, - } - }, - persist_across_reloads=False, - ) - - litellm.completion( - model=model, - messages=[{"role": "user", "content": "hello"}], - reasoning_effort=reasoning_effort, - drop_params=True, - api_key="fake-key", - api_base="https://api.perplexity.ai", - ) - - request_body: Final = json.loads(response_route.calls[0].request.content) - assert request_body["reasoning"] == {"effort": "high"} - - -_FOUNDRY_RESPONSES_FUNCTION_CALL_BODY: Final = { - "id": "resp_foundry", - "object": "response", - "created_at": 1789852145, - "status": "completed", - "model": "gpt-6-astra", - "output": [ - { - "id": "fc_1", - "type": "function_call", - "status": "completed", - "arguments": '{"city":"Paris"}', - "call_id": "call_1", - "name": "get_weather", - } - ], - "parallel_tool_calls": True, - "usage": { - "input_tokens": 53, - "output_tokens": 18, - "total_tokens": 71, - "output_tokens_details": {"reasoning_tokens": 0}, - }, - "error": None, - "incomplete_details": None, - "instructions": None, - "metadata": {}, - "temperature": 1.0, - "tool_choice": "auto", - "tools": [], - "top_p": 1.0, - "max_output_tokens": 200, - "previous_response_id": None, - "reasoning": {"effort": "medium", "summary": None}, - "truncation": "disabled", - "user": None, -} - - -def test_completion_bridges_azure_ai_foundry_gpt_5_4_plus_function_tools_to_responses( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch -): - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - responses_route: Final = respx_mock.post(f"{_FOUNDRY_API_BASE}/openai/v1/responses").respond( - json=_FOUNDRY_RESPONSES_FUNCTION_CALL_BODY - ) - - response: Final = litellm.completion( - model="azure_ai/gpt-6-astra", - messages=[{"role": "user", "content": "What is the weather in Paris? Use the tool."}], - tools=[ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get weather for a city", - "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, - }, - } - ], - max_tokens=200, - api_base=_FOUNDRY_API_BASE, - api_key="fake-foundry-key", - ) - - assert [str(call.request.url) for call in respx_mock.calls] == [f"{_FOUNDRY_API_BASE}/openai/v1/responses"] - request: Final = responses_route.calls[0].request - request_body: Final = json.loads(request.content) - assert request_body["tools"][0]["type"] == "function" - assert request_body["tools"][0]["name"] == "get_weather" - assert request.headers["api-key"] == "fake-foundry-key" - assert response.choices[0].finish_reason == "tool_calls" - assert response.choices[0].message.tool_calls[0].function.name == "get_weather" - - -@pytest.mark.parametrize( - "model, model_info, expected_model_param, expected_base_model_param", - [ - ("gemini/gemini-3.1-pro", None, "gemini-3.1-pro", None), - ( - "gemini/gemini-3.1-pro", - {"base_model": "gemini-3.1-pro-preview"}, - "gemini-3.1-pro", - "gemini-3.1-pro-preview", - ), - ], -) -def test_completion_optional_params_base_model( - model: str, - model_info: dict | None, - expected_model_param: str, - expected_base_model_param: str | None, -): - """``model_info.base_model`` must reach ``get_optional_params`` as ``base_model`` - (an additive capability hint), without overwriting ``model`` with the label. - - Regression for #29618: overwriting ``model`` with a friendly ``base_model`` - label made Bedrock drop ``tools``/``tool_choice`` under ``drop_params``.""" - with patch("litellm.main.get_optional_params") as mock_get_optional_params: - mock_get_optional_params.return_value = MagicMock() - - import litellm - - kwargs = { - "model": model, - "messages": [{"role": "user", "content": "What is the capital of France?"}], - "api_key": "fake-key", - "mock_response": "Hey, how's it going?", - } - if model_info is not None: - kwargs["model_info"] = model_info - - litellm.completion(**kwargs) - - assert mock_get_optional_params.called is True - call_kwargs = mock_get_optional_params.call_args.kwargs - assert call_kwargs["model"] == expected_model_param - assert call_kwargs["base_model"] == expected_base_model_param - - -@patch("litellm.completion_extras.responses_api_bridge.completion") -def test_gpt_5_4_responses_bridge_merges_reasoning_summary_kwarg_without_tools( - mock_responses_completion, -): - """reasoningSummary without tools should route and merge into reasoning_effort dict.""" - mock_responses_completion.return_value = MagicMock() - - import litellm - - litellm.completion( - model="gpt-5.4", - messages=[{"role": "user", "content": "ok"}], - reasoning_effort="medium", - reasoningSummary="auto", - api_key="fake-key", - ) - - assert mock_responses_completion.called is True - optional_params = mock_responses_completion.call_args.kwargs["optional_params"] - assert optional_params["reasoning_effort"] == { - "effort": "medium", - "summary": "auto", - } - assert "reasoningSummary" not in optional_params - assert "reasoning_summary" not in optional_params - - -@patch("litellm.completion_extras.responses_api_bridge.completion") -def test_responses_bridge_preserves_reasoning_summary_without_effort( - mock_responses_completion, -): - """Reasoning summary should survive responses routing even without effort.""" - mock_responses_completion.return_value = MagicMock() - - import litellm - - with patch.object(litellm, "route_all_chat_openai_to_responses", True): - litellm.completion( - model="gpt-4o", - messages=[{"role": "user", "content": "ok"}], - reasoningSummary="auto", - api_key="fake-key", - ) - - assert mock_responses_completion.called is True - optional_params = mock_responses_completion.call_args.kwargs["optional_params"] - assert optional_params["reasoning_effort"] == {"summary": "auto"} - assert "reasoningSummary" not in optional_params - assert "reasoning_summary" not in optional_params - - -@patch("litellm.completion_extras.responses_api_bridge.completion") -def test_gpt_5_responses_bridge_tools_and_reasoning_summary( - mock_responses_completion, -): - """Bare gpt-5 with tools + reasoningSummary should bridge (OpenCode-style).""" - mock_responses_completion.return_value = MagicMock() - - import litellm - - litellm.completion( - model="gpt-5", - messages=[{"role": "user", "content": "ok"}], - tools=[ - { - "type": "function", - "function": { - "name": "apply_patch", - "parameters": {"type": "object", "properties": {}}, - }, - } - ], - tool_choice="auto", - reasoning_effort="medium", - reasoningSummary="auto", - stream=True, - api_key="fake-key", - ) - - assert mock_responses_completion.called is True - optional_params = mock_responses_completion.call_args.kwargs["optional_params"] - assert optional_params.get("reasoning_effort") == { - "effort": "medium", - "summary": "auto", - } - - -def test_responses_api_bridge_check_handles_exception(): - """Test that responses_api_bridge_check handles exceptions and still processes responses/ models.""" - from litellm.main import responses_api_bridge_check - - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.side_effect = Exception("Model not found") - - model_info, model = responses_api_bridge_check( - model="responses/custom-model", custom_llm_provider="custom" - ) - - assert model == "custom-model" - assert model_info["mode"] == "responses" - - -def test_responses_api_bridge_check_global_flag_routes_openai(): - """When route_all_chat_openai_to_responses is True, any OpenAI model routes to responses.""" - from litellm.main import responses_api_bridge_check - - with patch.object(litellm, "route_all_chat_openai_to_responses", True): - model_info, model = responses_api_bridge_check( - model="gpt-4o", - custom_llm_provider="openai", - ) - - assert model == "gpt-4o" - assert model_info.get("mode") == "responses" - - -def test_responses_api_bridge_check_global_flag_does_not_affect_azure(): - """route_all_chat_openai_to_responses should not affect Azure models.""" - from litellm.main import responses_api_bridge_check - - with patch.object(litellm, "route_all_chat_openai_to_responses", True): - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 4096} - model_info, model = responses_api_bridge_check( - model="gpt-4o", - custom_llm_provider="azure", - ) - - assert model_info.get("mode") != "responses" - - -def test_responses_api_bridge_check_global_flag_default_false(): - """By default, route_all_chat_openai_to_responses is False and doesn't affect routing.""" - from litellm.main import responses_api_bridge_check - - with patch.object(litellm, "route_all_chat_openai_to_responses", False): - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 4096} - model_info, model = responses_api_bridge_check( - model="gpt-4o", - custom_llm_provider="openai", - ) - - assert model_info.get("mode") != "responses" - - -@pytest.mark.asyncio -async def test_async_mock_delay(): - """Use asyncio await for mock delay on acompletion""" - import time - - from litellm import acompletion - - start_time = time.time() - result = await acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hey, how's it going?"}], - mock_delay=0.01, - mock_response="Hello world", - ) - end_time = time.time() - delay = end_time - start_time - assert delay >= 0.01 - - -def test_stream_chunk_builder_keeps_tool_calls_carried_only_by_a_later_choice_of_a_multi_choice_chunk(): - from litellm import stream_chunk_builder - from litellm.types.utils import ( - ChatCompletionDeltaToolCall, - Delta, - Function, - ModelResponseStream, - StreamingChoices, - ) - - def chunk(choices: list[StreamingChoices]) -> ModelResponseStream: - return ModelResponseStream( - id="chatcmpl-multi-choice", - created=1751934860, - model="gpt-4.1-mini", - object="chat.completion.chunk", - choices=choices, - ) - - chunks = [ - chunk( - [ - StreamingChoices(index=0, delta=Delta(role="assistant", content="hello")), - StreamingChoices( - index=1, - delta=Delta( - role="assistant", - tool_calls=[ - ChatCompletionDeltaToolCall( - id="call_1", - index=0, - type="function", - function=Function(name="lookup_fruit", arguments='{"fruit":'), - ) - ], - ), - ), - ] - ), - chunk( - [ - StreamingChoices(index=0, delta=Delta(content=" world"), finish_reason="stop"), - StreamingChoices( - index=1, - delta=Delta( - tool_calls=[ChatCompletionDeltaToolCall(index=0, function=Function(arguments='"kiwi"}'))] - ), - finish_reason="tool_calls", - ), - ] - ), - ] - - response = stream_chunk_builder(chunks=chunks) - - tool_calls = response.choices[0].message.tool_calls - assert tool_calls is not None - assert [(call.id, call.function.name, call.function.arguments) for call in tool_calls] == [ - ("call_1", "lookup_fruit", '{"fruit":"kiwi"}') - ] - - -def test_stream_chunk_builder_thinking_blocks(): - from litellm import stream_chunk_builder - from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices - - chunks = [ - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_content="I need to summar", - thinking_blocks=[ - { - "type": "thinking", - "thinking": "I need to summar", - "signature": None, - } - ], - provider_specific_fields={ - "thinking_blocks": [ - { - "type": "thinking", - "thinking": "I need to summar", - "signature": None, - } - ] - }, - content="", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_content="ize the previous agent's thinking process into a", - thinking_blocks=[ - { - "type": "thinking", - "thinking": "ize the previous agent's thinking process into a", - "signature": None, - } - ], - provider_specific_fields={ - "thinking_blocks": [ - { - "type": "thinking", - "thinking": "ize the previous agent's thinking process into a", - "signature": None, - } - ] - }, - content="", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_content=" short description. Based on the input data provide", - thinking_blocks=[ - { - "type": "thinking", - "thinking": " short description. Based on the input data provide", - "signature": None, - } - ], - provider_specific_fields={ - "thinking_blocks": [ - { - "type": "thinking", - "thinking": " short description. Based on the input data provide", - "signature": None, - } - ] - }, - content="", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_content="d, it seems the agent was planning to refine their search", - thinking_blocks=[ - { - "type": "thinking", - "thinking": "d, it seems the agent was planning to refine their search", - "signature": None, - } - ], - provider_specific_fields={ - "thinking_blocks": [ - { - "type": "thinking", - "thinking": "d, it seems the agent was planning to refine their search", - "signature": None, - } - ] - }, - content="", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_content=" to focus more on technical aspects of home automation and home", - thinking_blocks=[ - { - "type": "thinking", - "thinking": " to focus more on technical aspects of home automation and home", - "signature": None, - } - ], - provider_specific_fields={ - "thinking_blocks": [ - { - "type": "thinking", - "thinking": " to focus more on technical aspects of home automation and home", - "signature": None, - } - ] - }, - content="", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_content=" energy system management.\n\nI'll create a brief", - thinking_blocks=[ - { - "type": "thinking", - "thinking": " energy system management.\n\nI'll create a brief", - "signature": None, - } - ], - provider_specific_fields={ - "thinking_blocks": [ - { - "type": "thinking", - "thinking": " energy system management.\n\nI'll create a brief", - "signature": None, - } - ] - }, - content="", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_content=" summary of what the agent was doing.", - thinking_blocks=[ - { - "type": "thinking", - "thinking": " summary of what the agent was doing.", - "signature": None, - } - ], - provider_specific_fields={ - "thinking_blocks": [ - { - "type": "thinking", - "thinking": " summary of what the agent was doing.", - "signature": None, - } - ] - }, - content="", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta( - reasoning_content="", - thinking_blocks=[ - { - "type": "thinking", - "thinking": "", - "signature": "ErUBCkYIBRgCIkAKBSMkB2+MBF643wiWxlERsGXVdlhbPx9lnTIbygzjFIeZ5uhTV+HNWDon9vQV4hmXvAKwQfwS8vkNFB366l05Egzt2U18IpRrZRyQn1UaDDdYvKHYP8Ps1IbWjSIw8eSYOU9gtqNcwR6D0wY7iOPx2GliDEatLI5rSs96CByoTIoADL2M5bX8KP0jEpbHKh0ccYryigdH/3J8EiFt/BmGUceVASP5l9r22dFWiBgC", - } - ], - provider_specific_fields={ - "thinking_blocks": [ - { - "type": "thinking", - "thinking": "", - "signature": "ErUBCkYIBRgCIkAKBSMkB2+MBF643wiWxlERsGXVdlhbPx9lnTIbygzjFIeZ5uhTV+HNWDon9vQV4hmXvAKwQfwS8vkNFB366l05Egzt2U18IpRrZRyQn1UaDDdYvKHYP8Ps1IbWjSIw8eSYOU9gtqNcwR6D0wY7iOPx2GliDEatLI5rSs96CByoTIoADL2M5bX8KP0jEpbHKh0ccYryigdH/3J8EiFt/BmGUceVASP5l9r22dFWiBgC", - } - ] - }, - content="", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=1, - delta=Delta( - provider_specific_fields=None, - content='{"a', - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=1, - delta=Delta( - provider_specific_fields=None, - content='gent_doing"', - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=1, - delta=Delta( - provider_specific_fields=None, - content=': "Re', - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=1, - delta=Delta( - provider_specific_fields=None, - content="searching", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=1, - delta=Delta( - provider_specific_fields=None, - content=" technic", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=1, - delta=Delta( - provider_specific_fields=None, - content="al aspect", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=1, - delta=Delta( - provider_specific_fields=None, - content="s of home au", - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason=None, - index=1, - delta=Delta( - provider_specific_fields=None, - content='tomation"}', - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - citations=None, - ), - ModelResponseStream( - id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", - created=1751934860, - model="claude-3-7-sonnet-latest", - object="chat.completion.chunk", - system_fingerprint=None, - choices=[ - StreamingChoices( - finish_reason="tool_calls", - index=0, - delta=Delta( - provider_specific_fields=None, - content=None, - role=None, - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields=None, - ), - ] - - response = stream_chunk_builder(chunks=chunks) - print(response) - - assert response is not None - assert response.choices[0].message.content is not None - assert response.choices[0].message.thinking_blocks is not None - - -from litellm.llms.openai.openai import OpenAIChatCompletion - - -def throw_retryable_error(*_, **__): - raise RuntimeError("BOOM") - - -@pytest.mark.asyncio -async def test_retrying() -> None: - litellm.num_retries = 10 - with ( - patch.object( - OpenAIChatCompletion, - "make_openai_chat_completion_request", - side_effect=throw_retryable_error, - ) as mock_request, - pytest.raises(litellm.InternalServerError, match="LiteLLM Retried: 10 times"), - ): - await litellm.acompletion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "Hello"}], - ) - - -def test_anthropic_disable_url_suffix_env_var(): - """Test that LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX prevents /v1/messages suffix.""" - import os - from unittest.mock import MagicMock, patch - - from litellm import completion - - # Test with environment variable disabled (default behavior) - with patch.dict(os.environ, {"ANTHROPIC_API_BASE": "https://api.example.com"}): - actual_api_base = None - - with patch("litellm.main.anthropic_chat_completions") as mock_anthropic: - - def capture_completion(**kwargs): - nonlocal actual_api_base - actual_api_base = kwargs.get("api_base") - mock_response = MagicMock() - mock_response.choices = [MagicMock()] - return mock_response - - mock_anthropic.completion = capture_completion - - # This should append /v1/messages - completion( - model="anthropic/claude-3-sonnet", - messages=[{"role": "user", "content": "test"}], - api_key="test-key", - ) - - # Verify the api_base has /v1/messages appended - assert actual_api_base.endswith("/v1/messages") - assert actual_api_base == "https://api.example.com/v1/messages" - - # Test with environment variable enabled - with patch.dict( - os.environ, - { - "ANTHROPIC_API_BASE": "https://api.example.com/custom/path", - "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX": "true", - }, - ): - actual_api_base = None - - with patch("litellm.main.anthropic_chat_completions") as mock_anthropic: - - def capture_completion(**kwargs): - nonlocal actual_api_base - actual_api_base = kwargs.get("api_base") - mock_response = MagicMock() - mock_response.choices = [MagicMock()] - return mock_response - - mock_anthropic.completion = capture_completion - - # This should NOT append /v1/messages - completion( - model="anthropic/claude-3-sonnet", - messages=[{"role": "user", "content": "test"}], - api_key="test-key", - ) - - # Verify the api_base does not have /v1/messages appended - assert actual_api_base == "https://api.example.com/custom/path" - assert not actual_api_base.endswith("/v1/messages") - - -def test_anthropic_text_disable_url_suffix_env_var(): - """Test that LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX prevents /v1/complete suffix for anthropic_text.""" - import os - from unittest.mock import MagicMock, patch - - from litellm import completion - - # Test with environment variable disabled (default behavior) - with patch.dict(os.environ, {"ANTHROPIC_API_BASE": "https://api.example.com"}): - actual_api_base = None - - with patch("litellm.main.base_llm_http_handler") as mock_handler: - - def capture_completion(**kwargs): - nonlocal actual_api_base - actual_api_base = kwargs.get("api_base") - return MagicMock() - - mock_handler.completion = capture_completion - - # This should append /v1/complete - completion( - model="anthropic_text/claude-instant-1", - messages=[{"role": "user", "content": "test"}], - api_key="test-key", - ) - - # Verify the api_base has /v1/complete appended - assert actual_api_base.endswith("/v1/complete") - assert actual_api_base == "https://api.example.com/v1/complete" - - # Test with environment variable enabled - with patch.dict( - os.environ, - { - "ANTHROPIC_API_BASE": "https://api.example.com/custom/complete", - "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX": "true", - }, - ): - actual_api_base = None - - with patch("litellm.main.base_llm_http_handler") as mock_handler: - - def capture_completion(**kwargs): - nonlocal actual_api_base - actual_api_base = kwargs.get("api_base") - return MagicMock() - - mock_handler.completion = capture_completion - - # This should NOT append /v1/complete - completion( - model="anthropic_text/claude-instant-1", - messages=[{"role": "user", "content": "test"}], - api_key="test-key", - ) - - # Verify the api_base does not have /v1/complete appended - assert actual_api_base == "https://api.example.com/custom/complete" - assert not actual_api_base.endswith("/v1/complete") - - -def test_image_edit_merges_headers_and_extra_headers(): - from litellm.images.main import base_llm_http_handler - - combined_headers = { - "x-test-header-one": "value-1", - "x-test-header-two": "value-2", - } - - mock_image_edit_config = MagicMock() - mock_image_edit_config.get_supported_openai_params.return_value = set() - mock_image_edit_config.map_openai_params.side_effect = lambda **kwargs: dict( - kwargs["image_edit_optional_params"] - ) - - with ( - patch( - "litellm.images.main.ProviderConfigManager.get_provider_image_edit_config", - return_value=mock_image_edit_config, - ) as mock_config, - patch.object( - base_llm_http_handler, - "image_edit_handler", - return_value="ok", - ) as mock_handler, - ): - response = litellm.image_edit( - image=MagicMock(name="image"), - prompt="test", - model="azure/gpt-image-1", - headers={"x-test-header-one": "value-1"}, - extra_headers={ - "x-test-header-two": "value-2", - }, - ) - - assert response == "ok" - mock_config.assert_called_once() - - handler_kwargs = mock_handler.call_args.kwargs - assert handler_kwargs["extra_headers"] == combined_headers - assert "extra_headers" not in handler_kwargs["image_edit_optional_request_params"] - - -@pytest.mark.parametrize("metadata_key", ("metadata", "litellm_metadata")) -@pytest.mark.parametrize("input_tokens", (51234, 0)) -def test_mock_completion_usage_reports_admission_input_tokens(metadata_key: str, input_tokens: int): - response = litellm.completion( - model="anthropic/claude-sonnet-5", - messages=[{"role": "user", "content": "hello"}], - mock_response="ok", - api_key="mock", - **{metadata_key: {"user_api_key_budget_reservation": {"reserved_cost": 1.0, "input_tokens": input_tokens}}}, - ) - - assert response.usage.prompt_tokens == input_tokens - assert response.usage.total_tokens == input_tokens + response.usage.completion_tokens - - -def test_mock_completion_usage_falls_back_to_default_without_admission_count(): - response = litellm.completion( - model="anthropic/claude-sonnet-5", - messages=[{"role": "user", "content": "hello"}], - mock_response="ok", - api_key="mock", - metadata={"user_api_key_budget_reservation": {"reserved_cost": 1.0}}, - ) - - assert response.usage.prompt_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT - - -_AZURE_AI_CUSTOM_PRICED_DEPLOYMENT: Final = { - "model_name": "azure-ai-custom-priced", - "litellm_params": { - "model": "azure_ai/gpt-5.6", - "api_key": "mock", - "api_base": "https://example.services.ai.azure.com", - "mock_response": "ok", - "input_cost_per_token": 3e-6, - "output_cost_per_token": 7e-6, - "cache_read_input_token_cost": 1e-7, - "cache_creation_input_token_cost": 5e-7, - }, - "model_info": {"id": "azure-ai-custom-priced-deployment-id"}, -} - - -def _expected_custom_price(response: litellm.ModelResponse) -> float: - params: Final = _AZURE_AI_CUSTOM_PRICED_DEPLOYMENT["litellm_params"] - return ( - response.usage.prompt_tokens * params["input_cost_per_token"] - + response.usage.completion_tokens * params["output_cost_per_token"] - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("use_async", (False, True)) -async def test_mock_completion_prices_azure_ai_router_deployment_with_custom_pricing(use_async: bool): - router: Final = litellm.Router(model_list=[_AZURE_AI_CUSTOM_PRICED_DEPLOYMENT]) - messages: Final = [{"role": "user", "content": "hello"}] - - response: Final = ( - await router.acompletion(model="azure-ai-custom-priced", messages=messages) - if use_async - else router.completion(model="azure-ai-custom-priced", messages=messages) - ) - - assert response._hidden_params["response_cost"] == pytest.approx(_expected_custom_price(response)) - assert response._hidden_params["custom_llm_provider"] == "azure_ai" - - -@pytest.mark.parametrize( - ("model", "expected_provider"), - (("anthropic/claude-sonnet-5", "anthropic"), ("no-such-provider-model", None)), -) -def test_mock_completion_infers_provider_when_called_directly_without_one(model: str, expected_provider: str | None): - response: Final = litellm.mock_completion( - model=model, - messages=[{"role": "user", "content": "hello"}], - mock_response="ok", - ) - - assert response.choices[0].message.content == "ok" - assert response._hidden_params.get("custom_llm_provider") == expected_provider - - -_ADMISSION_INPUT_TOKENS: Final = 51234 - - -def _admission_metadata(input_tokens: int) -> dict[str, object]: # mutable-ok: logging writes into metadata - return {"user_api_key_budget_reservation": {"reserved_cost": 1.0, "input_tokens": input_tokens}} - - -_ADMISSION_METADATA: Final = _admission_metadata(_ADMISSION_INPUT_TOKENS) -_MOCK_STREAM_MESSAGES: Final = [{"role": "user", "content": "hello " * 200}] -_STREAM_CHUNK_BUILDER_TOKEN_COUNTER: Final = "litellm.litellm_core_utils.streaming_chunk_builder_utils.token_counter" - - -def _prompt_token_counter_calls(token_counter: MagicMock) -> list[object]: - return [call for call in token_counter.call_args_list if call.kwargs.get("messages") is not None] - - -def _client_usage_chunks(chunks: list[ModelResponseStream]) -> list[Usage]: - return [chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None] - - -@pytest.mark.parametrize("n", (None, 2)) -def test_mock_completion_stream_usage_reports_admission_input_tokens_without_tokenizer_fallback(n: int | None): - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - chunks: Final = list( - litellm.completion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - stream=True, - n=n, - stream_options={"include_usage": True}, - metadata=_ADMISSION_METADATA, - ) - ) - - usage_chunks: Final = _client_usage_chunks(chunks) - assert len(usage_chunks) == 1 - assert usage_chunks[0].prompt_tokens == _ADMISSION_INPUT_TOKENS - assert usage_chunks[0].completion_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT - assert usage_chunks[0].total_tokens == _ADMISSION_INPUT_TOKENS + usage_chunks[0].completion_tokens - assert _prompt_token_counter_calls(token_counter) == [] - assert all(chunk.choices for chunk in chunks[:-1]) - assert {chunk.id for chunk in chunks} == {chunks[0].id} - - -@pytest.mark.asyncio -@pytest.mark.parametrize("n", (None, 2)) -async def test_mock_acompletion_stream_usage_reports_admission_input_tokens_without_tokenizer_fallback( - n: int | None, -): - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - response: Final = await litellm.acompletion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - stream=True, - n=n, - stream_options={"include_usage": True}, - litellm_metadata=_ADMISSION_METADATA, - ) - chunks: Final = [chunk async for chunk in response] - - usage_chunks: Final = _client_usage_chunks(chunks) - assert len(usage_chunks) == 1 - assert usage_chunks[0].prompt_tokens == _ADMISSION_INPUT_TOKENS - assert usage_chunks[0].total_tokens == _ADMISSION_INPUT_TOKENS + usage_chunks[0].completion_tokens - assert _prompt_token_counter_calls(token_counter) == [] - assert all(chunk.choices for chunk in chunks[:-1]) - assert {chunk.id for chunk in chunks} == {chunks[0].id} - - -def test_mock_completion_stream_without_include_usage_hides_usage_chunk_but_logs_admission_count(): - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - chunks: Final = list( - litellm.completion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - stream=True, - metadata=_ADMISSION_METADATA, - ) - ) - - assert _client_usage_chunks(chunks) == [] - assert all(len(chunk.choices) == 1 for chunk in chunks) - assert chunks[-1]._hidden_params["usage"].prompt_tokens == _ADMISSION_INPUT_TOKENS - assert _prompt_token_counter_calls(token_counter) == [] - - -def test_mock_completion_stream_with_empty_stream_options_completes_and_logs_admission_count(): - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - chunks: Final = list( - litellm.completion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - stream=True, - stream_options={}, - metadata=_ADMISSION_METADATA, - ) - ) - - assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok" - assert _client_usage_chunks(chunks) == [] - assert _prompt_token_counter_calls(token_counter) == [] - - -@pytest.mark.asyncio -async def test_mock_acompletion_stream_with_empty_stream_options_completes_and_logs_admission_count(): - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - response: Final = await litellm.acompletion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - stream=True, - stream_options={}, - litellm_metadata=_ADMISSION_METADATA, - ) - chunks: Final = [chunk async for chunk in response] - - assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok" - assert _client_usage_chunks(chunks) == [] - assert _prompt_token_counter_calls(token_counter) == [] - - -def test_mock_completion_stream_without_admission_count_falls_back_to_tokenizer(): - expected_prompt_tokens: Final = litellm.token_counter(model="openai/gpt-5.4-mini", messages=_MOCK_STREAM_MESSAGES) - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - chunks: Final = list( - litellm.completion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - stream=True, - stream_options={"include_usage": True}, - metadata={"user_api_key_budget_reservation": {"reserved_cost": 1.0}}, - ) - ) - - usage_chunks: Final = _client_usage_chunks(chunks) - assert len(usage_chunks) == 1 - assert usage_chunks[0].prompt_tokens == expected_prompt_tokens - assert usage_chunks[0].total_tokens == expected_prompt_tokens + usage_chunks[0].completion_tokens - assert len(_prompt_token_counter_calls(token_counter)) >= 1 - - -@pytest.mark.asyncio -async def test_mock_acompletion_stream_without_admission_count_falls_back_to_tokenizer(): - expected_prompt_tokens: Final = litellm.token_counter(model="openai/gpt-5.4-mini", messages=_MOCK_STREAM_MESSAGES) - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - response: Final = await litellm.acompletion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - stream=True, - stream_options={"include_usage": True}, - ) - chunks: Final = [chunk async for chunk in response] - - usage_chunks: Final = _client_usage_chunks(chunks) - assert len(usage_chunks) == 1 - assert usage_chunks[0].prompt_tokens == expected_prompt_tokens - assert len(_prompt_token_counter_calls(token_counter)) >= 1 - - -def _usage_triple(usage: Usage) -> tuple[int, int, int]: - return (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) - - -@pytest.mark.parametrize("input_tokens", (_ADMISSION_INPUT_TOKENS, 0)) -def test_mock_completion_stream_and_non_stream_report_the_same_admission_usage(input_tokens: int): - metadata: Final = _admission_metadata(input_tokens) - non_stream: Final = litellm.completion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - metadata=metadata, - ) - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - chunks: Final = list( - litellm.completion( - model="openai/gpt-5.4-mini", - messages=_MOCK_STREAM_MESSAGES, - mock_response="ok", - api_key="mock", - stream=True, - stream_options={"include_usage": True}, - metadata=metadata, - ) - ) - - assert _usage_triple(non_stream.usage) == _usage_triple(_client_usage_chunks(chunks)[0]) - assert non_stream.usage.prompt_tokens == input_tokens - assert _prompt_token_counter_calls(token_counter) == [] - - -@pytest.mark.asyncio -async def test_mock_acompletion_stream_reports_zero_admission_input_tokens_without_tokenizer_fallback(): - with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: - response: Final = await litellm.acompletion( - model="openai/gpt-5.4-mini", - messages=[{"role": "user", "content": ""}], - mock_response="ok", - api_key="mock", - stream=True, - stream_options={"include_usage": True}, - litellm_metadata=_admission_metadata(0), - ) - chunks: Final = [chunk async for chunk in response] - - usage_chunks: Final = _client_usage_chunks(chunks) - assert len(usage_chunks) == 1 - assert _usage_triple(usage_chunks[0]) == (0, usage_chunks[0].completion_tokens, usage_chunks[0].completion_tokens) - assert _prompt_token_counter_calls(token_counter) == [] - - -def test_mock_text_completion_stream_and_non_stream_report_the_same_zero_admission_usage(): - metadata: Final = _admission_metadata(0) - non_stream: Final = litellm.text_completion( - model="openai/gpt-5.4-mini", prompt="", mock_response="ok", api_key="mock", metadata=metadata - ) - chunks: Final = list( - litellm.text_completion( - model="openai/gpt-5.4-mini", - prompt="", - mock_response="ok", - api_key="mock", - stream=True, - stream_options={"include_usage": True}, - metadata=metadata, - ) - ) - - stream_usages: Final = tuple(chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None) - assert len(stream_usages) == 1 - assert _usage_triple(non_stream.usage) == _usage_triple(stream_usages[0]) - assert non_stream.usage.prompt_tokens == 0 - - -def test_mock_completion_stream_with_model_response(): - """Test that mock_completion correctly handles stream=True with a ModelResponse as mock_response.""" - from litellm import completion - from litellm.types.utils import Choices, Message, ModelResponse, Usage - - # Create a ModelResponse object - mock_model_response = ModelResponse( - id="chatcmpl-test-123", - created=1234567890, - model="gpt-4o-mini", - object="chat.completion", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="This is a test response", - role="assistant", - ), - ) - ], - usage=Usage( - prompt_tokens=10, - completion_tokens=20, - total_tokens=30, - ), - ) - - # Call completion with stream=True and mock_response as ModelResponse - response = completion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "Hello"}], - stream=True, - mock_response=mock_model_response, - ) - - # Verify that the response is a stream - assert response is not None - - # Collect all chunks from the stream - chunks = [] - for chunk in response: - chunks.append(chunk) - print(f"Chunk: {chunk}") - - # Verify we got chunks - assert len(chunks) > 0 - - # Verify the content is streamed correctly - accumulated_content = "" - for chunk in chunks: - if ( - hasattr(chunk.choices[0].delta, "content") - and chunk.choices[0].delta.content - ): - accumulated_content += chunk.choices[0].delta.content - - assert "This is a test response" in accumulated_content or len(chunks) > 0 - - -@pytest.mark.asyncio -async def test_async_mock_completion_stream_with_model_response(): - """Test that async mock_completion correctly handles stream=True with a ModelResponse as mock_response.""" - from litellm import acompletion - from litellm.types.utils import Choices, Message, ModelResponse, Usage - - # Create a ModelResponse object - mock_model_response = ModelResponse( - id="chatcmpl-test-456", - created=1234567890, - model="gpt-4o-mini", - object="chat.completion", - choices=[ - Choices( - finish_reason="stop", - index=0, - message=Message( - content="This is an async test response", - role="assistant", - ), - ) - ], - usage=Usage( - prompt_tokens=15, - completion_tokens=25, - total_tokens=40, - ), - ) - - # Call acompletion with stream=True and mock_response as ModelResponse - response = await acompletion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "Hello async"}], - stream=True, - mock_response=mock_model_response, - ) - - # Verify that the response is a stream - assert response is not None - - # Collect all chunks from the stream - chunks = [] - async for chunk in response: - chunks.append(chunk) - print(f"Async Chunk: {chunk}") - - # Verify we got chunks - assert len(chunks) > 0 - - # Verify the content is streamed correctly - accumulated_content = "" - for chunk in chunks: - if ( - hasattr(chunk.choices[0].delta, "content") - and chunk.choices[0].delta.content - ): - accumulated_content += chunk.choices[0].delta.content - - assert "This is an async test response" in accumulated_content or len(chunks) > 0 - - -class TestCallTypesOCR: - """Test that OCR call types are properly defined in CallTypes enum. - - Fixes https://github.com/BerriAI/litellm/issues/17381 - """ - - def test_ocr_call_type_exists(self): - """Test that CallTypes.ocr exists and has correct value.""" - from litellm.types.utils import CallTypes - - assert hasattr(CallTypes, "ocr") - assert CallTypes.ocr.value == "ocr" - - def test_aocr_call_type_exists(self): - """Test that CallTypes.aocr exists and has correct value.""" - from litellm.types.utils import CallTypes - - assert hasattr(CallTypes, "aocr") - assert CallTypes.aocr.value == "aocr" - - def test_ocr_call_type_from_string(self): - """Test that CallTypes can be constructed from 'ocr' string.""" - from litellm.types.utils import CallTypes - - call_type = CallTypes("ocr") - assert call_type == CallTypes.ocr - - def test_aocr_call_type_from_string(self): - """Test that CallTypes can be constructed from 'aocr' string. - - This is the actual use case that was failing - the OCR endpoint - uses route_type='aocr' and guardrails try to instantiate - CallTypes('aocr'). - """ - from litellm.types.utils import CallTypes - - call_type = CallTypes("aocr") - assert call_type == CallTypes.aocr - - -def test_stream_chunk_builder_text_completion_combines_text_and_usage(): - from litellm.main import stream_chunk_builder_text_completion - from litellm.types.utils import TextCompletionResponse - - chunks = [ - TextCompletionResponse( - id="cmpl-1", - object="text_completion", - created=1, - model="gpt-3.5-turbo-instruct", - choices=[{"text": "Hello", "index": 0, "logprobs": None, "finish_reason": None}], - ), - TextCompletionResponse( - id="cmpl-1", - object="text_completion", - created=1, - model="gpt-3.5-turbo-instruct", - choices=[{"text": " world", "index": 0, "logprobs": None, "finish_reason": "stop"}], - ), - ] - - response = stream_chunk_builder_text_completion( - chunks=chunks, messages=[{"role": "user", "content": "say hello"}] - ) - - assert response.choices[0].text == "Hello world" - assert response.choices[0].finish_reason == "stop" - assert response.usage.prompt_tokens > 0 - assert response.usage.completion_tokens > 0 - assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens - - -def test_completion_forwards_store_and_prompt_cache_key_to_openai(): - """ - Regression test for https://github.com/BerriAI/litellm/issues/33184 - - store and prompt_cache_key are documented OpenAI chat completion params that - were accepted as supported but silently dropped before the provider request - was built, because they were not named parameters of completion() and - get_optional_params() the way safety_identifier is. - """ - from openai import OpenAI - - client = OpenAI(api_key="fake-api-key") - - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: - try: - litellm.completion( - model="openai/gpt-4o", - messages=[{"role": "user", "content": "Hello"}], - store=False, - prompt_cache_key="test-cache-key", - client=client, - ) - except Exception as e: - print(e) - - mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs - assert request_body["store"] is False - assert request_body["prompt_cache_key"] == "test-cache-key" - - -@pytest.mark.asyncio -async def test_acompletion_forwards_store_and_prompt_cache_key_to_openai(): - """ - Async variant of the store/prompt_cache_key forwarding regression test for - https://github.com/BerriAI/litellm/issues/33184 - """ - from openai import AsyncOpenAI - - client = AsyncOpenAI(api_key="fake-api-key") - - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: - try: - await litellm.acompletion( - model="openai/gpt-4o", - messages=[{"role": "user", "content": "Hello"}], - store=False, - prompt_cache_key="test-cache-key", - client=client, - ) - except Exception as e: - print(e) - - mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs - assert request_body["store"] is False - assert request_body["prompt_cache_key"] == "test-cache-key" - - -def test_completion_omits_store_and_prompt_cache_key_when_not_passed(): - """ - When store and prompt_cache_key are not passed, they must not appear in the - outbound request body (guards against always forwarding None defaults). - """ - from openai import OpenAI - - client = OpenAI(api_key="fake-api-key") - - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: - try: - litellm.completion( - model="openai/gpt-4o", - messages=[{"role": "user", "content": "Hello"}], - client=client, - ) - except Exception as e: - print(e) - - mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs - assert "store" not in request_body - assert "prompt_cache_key" not in request_body - - -def test_completion_forwards_store_and_prompt_cache_key_to_mcp_gateway(): - """ - Regression test for the MCP gateway early-return in completion(): store and - prompt_cache_key are named params, so they no longer travel via **kwargs and - must be forwarded explicitly like safety_identifier and service_tier. - """ - with patch.object( - import_module("litellm.responses.mcp.chat_completions_handler"), "acompletion_with_mcp" - ) as mock_mcp: - result = litellm.completion( - model="openai/gpt-4o", - messages=[{"role": "user", "content": "Hello"}], - tools=[{"type": "mcp", "server_url": "litellm_proxy"}], - store=False, - prompt_cache_key="test-cache-key", - ) - - result.close() - mock_mcp.assert_called_once() - call_kwargs = mock_mcp.call_args.kwargs - assert call_kwargs["store"] is False - assert call_kwargs["prompt_cache_key"] == "test-cache-key" - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "aws_credential_kwargs", - [ - { - "aws_session_name": "litellm-gcp", - "aws_role_name": "arn:aws:iam::123456789012:role/litellm-bedrock-role", - "aws_web_identity_token": "oidc/google/108963886734710037768", - }, - { - "aws_access_key_id": "AKIASTATICKEYFORTEST", - "aws_secret_access_key": "static-secret-key", - "aws_session_token": "static-session-token", - }, - ], - ids=["web_identity", "static_keys"], -) -async def test_acompletion_forwards_aws_credentials_through_responses_bridge( - respx_mock: respx.MockRouter, monkeypatch, aws_credential_kwargs: dict -): - from botocore.credentials import Credentials - - from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM - - original_disable_aiohttp = litellm.disable_aiohttp_transport - try: - litellm.disable_aiohttp_transport = True - monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") - litellm.in_memory_llm_clients_cache.flush_cache() - monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) - monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) - - get_credentials_mock = MagicMock(return_value=Credentials("fake-key", "fake-secret")) - monkeypatch.setattr(BaseAWSLLM, "get_credentials", get_credentials_mock) - - respx_mock.post("https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses").respond( - json={ - "id": "resp_123", - "object": "response", - "created_at": 1760144904, - "status": "completed", - "model": "openai.gpt-5.4", - "output": [ - { - "type": "message", - "id": "msg_1", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": "ok", "annotations": []}], - } - ], - } - ) - - response = await litellm.acompletion( - model="bedrock_mantle/openai.gpt-5.4", - messages=[{"role": "user", "content": "hi"}], - api_base="https://bedrock-mantle.us-east-2.api.aws/v1", - aws_region_name="us-east-2", - num_retries=0, - **aws_credential_kwargs, - ) - - assert response.choices[0].message.content == "ok" - credential_kwargs = get_credentials_mock.call_args.kwargs - assert credential_kwargs["aws_region_name"] == "us-east-2" - for key, value in aws_credential_kwargs.items(): - assert credential_kwargs[key] == value - authorization = respx_mock.calls.last.request.headers["Authorization"] - assert authorization.startswith("AWS4-HMAC-SHA256") - assert "fake-key" in authorization - finally: - litellm.disable_aiohttp_transport = original_disable_aiohttp - litellm.in_memory_llm_clients_cache.flush_cache() - - -_GEMINI_RESPONSE_BODY = { - "candidates": [{"content": {"parts": [{"text": "hello"}], "role": "model"}, "finishReason": "STOP"}], - "usageMetadata": {"promptTokenCount": 2, "candidatesTokenCount": 1, "totalTokenCount": 3}, -} - - -def _gemini_client_returning_a_reply(): - """An injected HTTP client whose post() answers like generativelanguage does.""" - from litellm.llms.custom_httpx.http_handler import HTTPHandler - - client = HTTPHandler() - request = httpx.Request("POST", "https://generativelanguage.googleapis.com/") - post = MagicMock(return_value=httpx.Response(200, json=_GEMINI_RESPONSE_BODY, request=request)) - return client, post - - -@pytest.fixture -def restore_model_registry(): - """litellm.model_cost and the provider name sets are module-global. - - register_model merges into the existing entry in place, hence the deep copy. - """ - model_cost = copy.deepcopy(litellm.model_cost) - openai_models = set(litellm.open_ai_chat_completion_models) - yield - litellm.model_cost.clear() - litellm.model_cost.update(model_cost) - litellm.open_ai_chat_completion_models.clear() - litellm.open_ai_chat_completion_models.update(openai_models) - - -def test_openai_model_name_does_not_outrank_explicit_provider(): - """`gemini/gpt-4o` goes to Google, not to litellm's OpenAI handler. - - completion() checks `model in litellm.open_ai_chat_completion_models` ahead of - the gemini branch, so the call used to reach the OpenAI handler carrying - VertexGeminiConfig, whose transform_request raises NotImplementedError. - """ - assert "gpt-4o" in litellm.open_ai_chat_completion_models - client, post = _gemini_client_returning_a_reply() - - with patch.object(client, "post", new=post): - response = litellm.completion( - model="gemini/gpt-4o", - messages=[{"role": "user", "content": "hello"}], - api_key="test-api-key", - client=client, - ) - - assert "generativelanguage.googleapis.com" in post.call_args.kwargs["url"] - assert "models/gpt-4o" in post.call_args.kwargs["url"] - assert response.choices[0].message.content == "hello" - - -def test_mislabelled_pricing_entry_does_not_reroute_provider(restore_model_registry): - """register_model is the other way into the same failure. - - An entry claiming litellm_provider "openai" adds its name to - open_ai_chat_completion_models, so one mislabelled price reroutes every later - call to that model in the process. - """ - litellm.register_model( - { - "gemini-2.5-pro": { - "litellm_provider": "openai", - "mode": "chat", - "input_cost_per_token": 1e-06, - "output_cost_per_token": 4e-06, - } - } - ) - assert "gemini-2.5-pro" in litellm.open_ai_chat_completion_models - client, post = _gemini_client_returning_a_reply() - - with patch.object(client, "post", new=post): - response = litellm.completion( - model="gemini/gemini-2.5-pro", - messages=[{"role": "user", "content": "hello"}], - api_key="test-api-key", - client=client, - ) - - assert "generativelanguage.googleapis.com" in post.call_args.kwargs["url"] - assert response.choices[0].message.content == "hello" - - -def test_openai_model_without_a_provider_still_routes_to_openai(): - from openai import OpenAI - - client = OpenAI(api_key="fake-key") - raw_response = client.chat.completions.with_raw_response - with patch.object(raw_response, "create") as mock_create, contextlib.suppress(Exception): - litellm.completion( - model="gpt-4o", - messages=[{"role": "user", "content": "hello"}], - client=client, - ) - - mock_create.assert_called() - - -def _openai_chat_create_kwargs(client, **completion_kwargs): - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: - with contextlib.suppress(Exception): - litellm.completion( - messages=[{"role": "system", "content": "sys"}, {"role": "user", "content": "hi"}], - cache_control_injection_points=[{"location": "message", "role": "system"}], - client=client, - **completion_kwargs, - ) - - mock_client.assert_called_once() - return mock_client.call_args.kwargs - - -@pytest.fixture -def _no_openai_api_base_override(monkeypatch): - monkeypatch.delenv("OPENAI_BASE_URL", raising=False) - monkeypatch.delenv("OPENAI_API_BASE", raising=False) - monkeypatch.setattr(litellm, "api_base", None) - - -@pytest.mark.usefixtures("_no_openai_api_base_override") -def test_completion_custom_api_base_sends_no_prompt_cache_breakpoint_for_gpt_5_6(): - from openai import OpenAI - - client = OpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") - request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", api_base="http://127.0.0.1:9/v1") - - assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} - assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) - assert "prompt_cache_options" not in json.dumps(request_body) - - -@pytest.mark.usefixtures("_no_openai_api_base_override") -def test_completion_custom_base_url_sends_no_prompt_cache_breakpoint_for_gpt_5_6(): - from openai import OpenAI - - client = OpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") - request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", base_url="http://127.0.0.1:9/v1") - - assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} - assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) - assert "prompt_cache_options" not in json.dumps(request_body) - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("_no_openai_api_base_override") -async def test_acompletion_custom_base_url_sends_no_prompt_cache_breakpoint_for_gpt_5_6(): - from openai import AsyncOpenAI - - client = AsyncOpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_create: - with contextlib.suppress(Exception): - await litellm.acompletion( - model="gpt-5.6", - messages=[{"role": "system", "content": "sys"}, {"role": "user", "content": "hi"}], - cache_control_injection_points=[{"location": "message", "role": "system"}], - client=client, - base_url="http://127.0.0.1:9/v1", - ) - - mock_create.assert_called_once() - request_body = mock_create.call_args.kwargs - - assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} - assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) - assert "prompt_cache_options" not in json.dumps(request_body) - - -@pytest.mark.usefixtures("_no_openai_api_base_override") -def test_completion_default_api_base_sends_prompt_cache_breakpoint_for_gpt_5_6(): - from openai import OpenAI - - client = OpenAI(api_key="fake-api-key") - request_body = _openai_chat_create_kwargs(client, model="gpt-5.6") - - assert request_body["messages"][0]["content"] == [ - {"type": "text", "text": "sys", "prompt_cache_breakpoint": {"mode": "explicit"}} - ] - assert request_body["extra_body"]["prompt_cache_options"] == {"mode": "explicit"} - - -_SUBSCRIPTION_OAUTH_CREDENTIAL = "Bearer sk-ant-oat01-fake-subscription-token-for-testing-0123456789" - - -def _scoped_headers_for_oauth_request(): - from litellm.types.utils import ProviderSpecificHeader - - return [ - ProviderSpecificHeader( - custom_llm_provider="anthropic,bedrock,vertex_ai", - extra_headers={"anthropic-version": "2023-06-01"}, - ), - ProviderSpecificHeader( - custom_llm_provider="anthropic", - extra_headers={"authorization": _SUBSCRIPTION_OAUTH_CREDENTIAL}, - ), - ] - - -def _run_anthropic_hop_with_shared_headers(shared_headers): - litellm.completion( - model="anthropic/claude-3-5-sonnet-20240620", - messages=[{"role": "user", "content": "Say OK"}], - extra_headers=shared_headers, - provider_specific_header=_scoped_headers_for_oauth_request(), - api_key="sk-fake-anthropic-key", - mock_response="OK", - ) - - -def test_completion_does_not_mutate_caller_supplied_headers(): - shared_headers = {"x-tenant": "acme"} - - _run_anthropic_hop_with_shared_headers(shared_headers) - - assert shared_headers == {"x-tenant": "acme"} - - -def test_anthropic_oauth_credential_does_not_persist_into_next_provider_hop(): - shared_headers = {"x-tenant": "acme"} - - _run_anthropic_hop_with_shared_headers(shared_headers) - - leaked = [name for name, value in shared_headers.items() if value == _SUBSCRIPTION_OAUTH_CREDENTIAL] - assert leaked == [] - assert "anthropic-version" not in shared_headers - - -STREAM_COST_MODEL = "gpt-4o" -STREAMED_USAGE = {"prompt_tokens": 137, "completion_tokens": 42, "total_tokens": 179} - - -def _text_chunk(content, finish_reason=None, usage=None): - chunk = { - "id": "chatcmpl-stream-cost", - "object": "chat.completion.chunk", - "created": 1700000000, - "model": STREAM_COST_MODEL, - "choices": [ - { - "index": 0, - "delta": {"role": "assistant", "content": content}, - "finish_reason": finish_reason, - } - ], - } - if usage is not None: - chunk["usage"] = usage - return chunk - - -def _priced_at(prompt_tokens, completion_tokens): - prices = litellm.model_cost[STREAM_COST_MODEL] - return ( - prompt_tokens * prices["input_cost_per_token"] - + completion_tokens * prices["output_cost_per_token"] - ) - - -@pytest.fixture -def local_cost_map(monkeypatch): - """The prices these tests assert are the checked-in ones. Setting the environment - variable alone does not reload the map, so pin the map itself. - - Prices are read through two separate lru_caches, so pinning ``model_cost`` is not - enough on its own: an entry warmed against the network-fetched map keeps its old - prices and billing reads those while the assertions read the pinned map. - ``_invalidate_model_cost_lowercase_map`` clears both caches, where - ``get_model_info.cache_clear`` reaches only one. Invalidate on the way in and out - so entries never leak across tests in either direction.""" - from litellm.utils import _invalidate_model_cost_lowercase_map - - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - _invalidate_model_cost_lowercase_map() - yield - _invalidate_model_cost_lowercase_map() - - -def test_a_streamed_response_bills_the_usage_the_provider_reported(local_cost_map): - rebuilt = litellm.stream_chunk_builder( - chunks=[ - _text_chunk("Hello"), - _text_chunk(" there"), - _text_chunk(None, finish_reason="stop", usage=STREAMED_USAGE), - ], - messages=[{"role": "user", "content": "hi"}], - ) - - assert rebuilt.choices[0].message.content == "Hello there" - assert rebuilt.usage.prompt_tokens == STREAMED_USAGE["prompt_tokens"] - assert rebuilt.usage.completion_tokens == STREAMED_USAGE["completion_tokens"] - - cost = litellm.completion_cost(completion_response=rebuilt, model=STREAM_COST_MODEL) - - assert cost == pytest.approx(_priced_at(137, 42)) - - -def test_streaming_and_not_streaming_bill_the_same_usage_the_same(local_cost_map): - rebuilt = litellm.stream_chunk_builder( - chunks=[ - _text_chunk("Hello"), - _text_chunk(" there"), - _text_chunk(None, finish_reason="stop", usage=STREAMED_USAGE), - ], - messages=[{"role": "user", "content": "hi"}], - ) - whole = litellm.ModelResponse( - id="chatcmpl-stream-cost", - model=STREAM_COST_MODEL, - object="chat.completion", - created=1700000000, - choices=[ - { - "index": 0, - "message": {"role": "assistant", "content": "Hello there"}, - "finish_reason": "stop", - } - ], - usage=STREAMED_USAGE, - ) - - assert litellm.completion_cost( - completion_response=rebuilt, model=STREAM_COST_MODEL - ) == pytest.approx(litellm.completion_cost(completion_response=whole, model=STREAM_COST_MODEL)) - - -def test_a_stream_that_reported_no_usage_is_still_billed(local_cost_map): - rebuilt = litellm.stream_chunk_builder( - chunks=[ - _text_chunk("Hello"), - _text_chunk(" there"), - _text_chunk(None, finish_reason="stop"), - ], - messages=[{"role": "user", "content": "hi"}], - ) - - assert rebuilt.usage.prompt_tokens > 0 - assert rebuilt.usage.completion_tokens > 0 - - cost = litellm.completion_cost(completion_response=rebuilt, model=STREAM_COST_MODEL) - - assert cost > 0 - assert cost == pytest.approx( - _priced_at(rebuilt.usage.prompt_tokens, rebuilt.usage.completion_tokens) - ) - - -@pytest.mark.asyncio -async def test_acompletion_resolves_provider_from_api_base(): - response = await litellm.acompletion( - model="deepseek-chat", - api_base="https://api.deepseek.com/v1", - api_key="fake-key", - messages=[{"role": "user", "content": "hi"}], - mock_response="resolved", - ) - - assert response.choices[0].message.content == "resolved" - - -@dataclass(frozen=True, slots=True) -class _RecordedSpeechSuccess: - call_type: str | None - spend_metadata: Mapping[str, object] - response_cost: float | None - logged_response_cost: float | None - - -def _record_speech_success(payload: dict[str, object]) -> _RecordedSpeechSuccess: - call_type: Final = payload.get("call_type") - response_cost: Final = payload.get("response_cost") - logging_payload: Final = payload.get("standard_logging_object") - logged_cost: Final = logging_payload.get("response_cost") if isinstance(logging_payload, dict) else None - return _RecordedSpeechSuccess( - call_type=call_type if isinstance(call_type, str) else None, - spend_metadata=get_litellm_metadata_from_kwargs(payload), - response_cost=response_cost if isinstance(response_cost, float) else None, - logged_response_cost=logged_cost if isinstance(logged_cost, float) else None, - ) - - -class _SuccessEventRecorder(CustomLogger): - def __init__(self) -> None: - super().__init__() - self.events: list[_RecordedSpeechSuccess] = [] # mutable-ok: test recorder of success-callback events - - async def async_log_success_event( - self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object - ) -> None: - self.events.append(_record_speech_success(kwargs)) - - -async def _wait_for_success_event(recorder: _SuccessEventRecorder, call_type: str) -> _RecordedSpeechSuccess: - for _ in range(100): - if (event := next((e for e in recorder.events if e.call_type == call_type), None)) is not None: - return event - await asyncio.sleep(0.05) - pytest.fail(f"no {call_type} success event; got {[e.call_type for e in recorder.events]}") - - -def _gemini_tts_generate_content_response() -> dict[str, object]: - return { - "candidates": [ - { - "content": { - "parts": [ - { - "inlineData": { - "mimeType": "audio/L16;codec=pcm;rate=24000", - "data": base64.b64encode(b"pcm-audio-bytes").decode(), - } - } - ], - "role": "model", - }, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 5, - "candidatesTokenCount": 60, - "totalTokenCount": 65, - "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 5}], - "candidatesTokensDetails": [{"modality": "AUDIO", "tokenCount": 60}], - }, - "modelVersion": "gemini-2.5-flash-preview-tts", - } - - -@pytest.mark.asyncio -async def test_aspeech_gemini_bridge_keeps_proxy_metadata_for_spend_tracking( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - monkeypatch.delenv("GEMINI_API_KEY", raising=False) - monkeypatch.delenv("GOOGLE_API_KEY", raising=False) - recorder: Final = _SuccessEventRecorder() - monkeypatch.setattr(litellm, "callbacks", [recorder]) - mock_route: Final = respx_mock.post( - url__regex=r"https://generativelanguage\.googleapis\.com/v1beta/models/gemini-2\.5-flash-preview-tts:generateContent.*" - ).mock(return_value=httpx.Response(200, json=_gemini_tts_generate_content_response())) - - await litellm.aspeech( - model="gemini/gemini-2.5-flash-preview-tts", - input="spend tracking check", - voice="Kore", - api_key="fake-gemini-key", - metadata={"user_api_key": "hashed-virtual-key", "user_api_key_user_id": "user-1"}, - ) - - assert mock_route.called - assert mock_route.calls.last.request.headers["x-goog-api-key"] == "fake-gemini-key" - speech_event: Final = await _wait_for_success_event(recorder, call_type="aspeech") - assert speech_event.spend_metadata["user_api_key"] == "hashed-virtual-key" - assert speech_event.spend_metadata["user_api_key_user_id"] == "user-1" - expected_prompt_cost, expected_completion_cost = litellm.cost_per_token( - model="gemini/gemini-2.5-flash-preview-tts", - usage_object=Usage(prompt_tokens=5, completion_tokens=60, total_tokens=65), - ) - expected_cost: Final = expected_prompt_cost + expected_completion_cost - assert expected_cost > 0 - assert speech_event.response_cost == pytest.approx(expected_cost) - assert speech_event.logged_response_cost == pytest.approx(expected_cost) - - -def _stream_builder_text_chunk(model: str, content: str, finish_reason: str | None = None) -> ModelResponseStream: - return ModelResponseStream( - id="chatcmpl-cost", - created=1724900000, - model=model, - object="chat.completion.chunk", - choices=[StreamingChoices(finish_reason=finish_reason, index=0, delta=Delta(content=content, role="assistant"))], - ) - - -def test_stream_chunk_builder_sets_hidden_response_cost_for_known_model(): - chunks: Final = [ - _stream_builder_text_chunk("gpt-4o", "Hello "), - _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), - ] - - response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) - - assert response is not None - prompt_cost, completion_cost = litellm.cost_per_token(model="gpt-4o", usage_object=response.usage) - expected_cost: Final = prompt_cost + completion_cost - assert expected_cost > 0 - assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) - - -def test_stream_chunk_builder_unknown_model_leaves_response_cost_unset(): - chunks: Final = [ - _stream_builder_text_chunk("totally-unknown-model-xyz", "Hello "), - _stream_builder_text_chunk("totally-unknown-model-xyz", "world.", finish_reason="stop"), - ] - - response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) - - assert response is not None - assert response._hidden_params.get("response_cost") is None - assert response.choices[0].message.content == "Hello world." - - -def test_stream_chunk_builder_prices_proxy_alias_via_model_map(): - chunks: Final = [ - _stream_builder_text_chunk("claude-opus-5", "Hello "), - _stream_builder_text_chunk("claude-opus-5", "world.", finish_reason="stop"), - ] - for chunk in chunks: - chunk._hidden_params = {"custom_llm_provider": "openai"} - - response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) - - assert response is not None - assert response._hidden_params["custom_llm_provider"] == "openai" - prompt_cost, completion_cost = litellm.cost_per_token(model="claude-opus-5", usage_object=response.usage) - expected_cost: Final = prompt_cost + completion_cost - assert expected_cost > 0 - assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) - - -def _stream_builder_logging_obj(model: str = "gpt-4o", custom_llm_provider: str = "openai") -> LiteLLMLogging: - logging_obj: Final = LiteLLMLogging( - model=model, - messages=[{"role": "user", "content": "hi"}], - stream=True, - call_type="completion", - start_time=datetime.now(), - litellm_call_id="test-call-id", - function_id="test-function-id", - ) - logging_obj.update_environment_variables( - model=model, - user=None, - optional_params={}, - litellm_params={"custom_llm_provider": custom_llm_provider}, - custom_llm_provider=custom_llm_provider, - ) - return logging_obj - - -def test_stream_chunk_builder_stamps_streaming_usage_cost_by_default(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) - chunks: Final = [ - _stream_builder_text_chunk("gpt-4o", "Hello "), - _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), - ] - - response: Final = litellm.stream_chunk_builder( - chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=_stream_builder_logging_obj() - ) - - assert response is not None - usage_cost: Final = getattr(response.usage, "cost", None) - assert usage_cost is not None - assert usage_cost > 0 - assert response._hidden_params["response_cost"] == pytest.approx(usage_cost) - - -def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable(): - import time as time_module - - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging - - logging_obj: Final = LiteLLMLogging( - model="us.anthropic.claude-opus-5", - messages=[{"role": "user", "content": "hi"}], - stream=True, - call_type="completion", - start_time=time_module.time(), - litellm_call_id="stream-builder-alias-unpriceable", - function_id="1", - ) - logging_obj.model_call_details["custom_llm_provider"] = "bedrock" - logging_obj.optional_params = {} - usage_chunk: Final = _stream_builder_text_chunk("bedrock-claude-opus-5", "") - usage_chunk.usage = Usage(prompt_tokens=40, completion_tokens=5, total_tokens=45) - chunks: Final = [ - _stream_builder_text_chunk("bedrock-claude-opus-5", "Hello ", finish_reason="stop"), - usage_chunk, - ] - - response: Final = litellm.stream_chunk_builder( - chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj - ) - - assert response is not None - assert getattr(response.usage, "cost", None) is None - assert response._hidden_params.get("response_cost") is None - - -def test_stream_chunk_builder_keeps_provider_reported_usage_cost(): - usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") - usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5) - chunks: Final = [ - _stream_builder_text_chunk("gpt-4o", "Hello "), - _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), - usage_chunk, - ] - - response: Final = litellm.stream_chunk_builder( - chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=_stream_builder_logging_obj() - ) - - assert response is not None - assert getattr(response.usage, "cost", None) == pytest.approx(0.5) - assert response._hidden_params["response_cost"] == pytest.approx(0.5) - - -def test_stream_chunk_builder_prices_alias_from_openai_sdk_usage_chunk(): - from openai.types.completion_usage import CompletionUsage - - usage_chunk: Final = _stream_builder_text_chunk("mantle-claude", "") - usage_chunk.usage = CompletionUsage(prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704) - assert type(usage_chunk.usage) is CompletionUsage - chunks: Final = [ - _stream_builder_text_chunk("mantle-claude", "Hello "), - _stream_builder_text_chunk("mantle-claude", "world.", finish_reason="stop"), - usage_chunk, - ] - - response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) - - assert response is not None - assert response.usage.prompt_tokens == 20 - assert response.usage.completion_tokens == 60 - assert getattr(response.usage, "cost", None) == pytest.approx(0.000704) - assert response._hidden_params["response_cost"] == pytest.approx(0.000704) - - -def test_stream_chunk_builder_leaves_xai_reported_cost_to_the_calculator(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(litellm, "cost_margin_config", {"xai": 0.5}) - usage_chunk: Final = _stream_builder_text_chunk("grok-4", "") - usage_chunk.usage = Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7, cost=0.42) - chunks: Final = [ - _stream_builder_text_chunk("grok-4", "Hello "), - _stream_builder_text_chunk("grok-4", "world.", finish_reason="stop"), - usage_chunk, - ] - logging_obj: Final = _stream_builder_logging_obj(model="grok-4", custom_llm_provider="xai") - - response: Final = litellm.stream_chunk_builder( - chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj - ) - - assert response is not None - assert getattr(response.usage, "cost", None) == pytest.approx(0.42) - assert response._hidden_params.get("response_cost") is None - assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.63) - - -def test_speech_mistral_dispatches_and_decodes_audio(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") - audio_bytes: Final = b"ID3-fake-mp3-bytes" - mock_route: Final = respx_mock.post("https://api.mistral.ai/v1/audio/speech").mock( - return_value=httpx.Response(200, json={"audio_data": base64.b64encode(audio_bytes).decode()}) - ) - - response: Final = litellm.speech( - model="mistral/voxtral-mini-tts-2603", - input="hello from litellm", - voice="en_paul_neutral", - response_format="wav", - speed=2, - instructions="sound cheerful", - ) - - assert mock_route.called - request_body: Final = json.loads(mock_route.calls.last.request.content) - assert request_body == { - "model": "voxtral-mini-tts-2603", - "input": "hello from litellm", - "voice_id": "en_paul_neutral", - "response_format": "wav", - } - assert mock_route.calls.last.request.headers["authorization"] == "Bearer sk-mistral-test" - assert response.content == audio_bytes - - -def test_speech_mistral_routes_to_configured_api_base(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") - audio_bytes: Final = b"ID3-gateway-bytes" - gateway_route: Final = respx_mock.post("https://mistral.gateway.internal/v1/audio/speech").mock( - return_value=httpx.Response(200, json={"audio_data": base64.b64encode(audio_bytes).decode()}) - ) - - response: Final = litellm.speech( - model="mistral/voxtral-mini-tts-2603", - input="hello from litellm", - voice="en_paul_neutral", - api_base="https://mistral.gateway.internal", - ) - - assert gateway_route.called - assert response.content == audio_bytes - - -FOUNDRY_HOST: Final = "https://my-project.services.ai.azure.com" - - -def test_azure_ai_transcription_on_a_foundry_host_uses_the_azure_openai_deployment_route( - respx_mock: respx.MockRouter, -): - route: Final = respx_mock.post( - url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/whisper-1/audio/transcriptions\?api-version=.+" - ).mock(return_value=httpx.Response(200, json={"text": "hello"})) - - response: Final = litellm.transcription( - model="azure_ai/whisper-1", - file=("tone.wav", b"RIFF\x00\x00\x00\x00WAVE", "audio/wav"), - api_base=FOUNDRY_HOST, - api_key="fake-key", - ) - - assert route.called - assert response.text == "hello" - - -def test_azure_ai_speech_on_a_foundry_host_uses_the_azure_openai_deployment_route( - respx_mock: respx.MockRouter, -): - route: Final = respx_mock.post( - url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/tts-1/audio/speech\?api-version=.+" - ).mock(return_value=httpx.Response(200, content=b"mp3-bytes")) - - response: Final = litellm.speech( - model="azure_ai/tts-1", - input="hello", - voice="alloy", - api_base=FOUNDRY_HOST, - api_key="fake-key", - ) - - assert route.called - assert response.content == b"mp3-bytes" - - -FORWARDED_CLIENT_HEADERS: Final = {"x-forwarded-for": "10.0.0.1", "x-amzn-trace-id": "Root=1-lit7694"} - - -def _chat_completion_json() -> Mapping[str, object]: - return { - "id": "chatcmpl-lit7694", - "object": "chat.completion", - "created": 1, - "model": "gpt-5.4", - "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, - } - - -def _chat_completion_sse() -> bytes: - chunk: Final = { - "id": "chatcmpl-lit7694", - "object": "chat.completion.chunk", - "created": 1, - "model": "gpt-5.4", - "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], - } - return f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode() - - -@pytest.mark.parametrize("stream", [False, True]) -def test_bridged_responses_with_openai_http_handler_keeps_forwarded_headers_out_of_the_body( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, stream: bool -): - monkeypatch.setenv("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER", "true") - route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - return_value=httpx.Response(200, content=_chat_completion_sse(), headers={"content-type": "text/event-stream"}) - if stream - else httpx.Response(200, json=_chat_completion_json()) - ) - - response: Final = litellm.responses( - model="openai/gpt-5.4", - input="Reply with the single word ok", - stream=stream, - use_chat_completions_api=True, - headers=dict(FORWARDED_CLIENT_HEADERS), - api_key="sk-test", - ) - if stream: - list(response) - - assert route.called - request: Final = route.calls.last.request - body: Final = json.loads(request.content) - assert "extra_headers" not in body - assert body["model"] == "gpt-5.4" - assert {k: request.headers[k] for k in FORWARDED_CLIENT_HEADERS} == FORWARDED_CLIENT_HEADERS - - -@pytest.mark.parametrize("http2_on", [True, False]) -def test_aiohttp_openai_warns_only_when_http2_enabled( - monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, http2_on: bool -): - from litellm.main import base_llm_aiohttp_handler - - monkeypatch.setattr(litellm, "http2", http2_on) - monkeypatch.delenv("LITELLM_HTTP2", raising=False) - - handler_completion: Final = MagicMock(return_value=MagicMock()) - monkeypatch.setattr(base_llm_aiohttp_handler, "completion", handler_completion) - - with caplog.at_level(logging.WARNING, logger="LiteLLM"): - litellm.completion( - model="aiohttp_openai/gpt-4o", - messages=[{"role": "user", "content": "hi"}], - api_key="sk-test", - ) - - assert handler_completion.called - warned: Final = "aiohttp_openai/ always uses aiohttp" in caplog.text - assert warned is http2_on - - -@pytest.mark.parametrize("tool_choice", [{"type": "bogus"}, {"name": "lookup_fruit"}, {"type": "file_search"}]) -def test_completion_rejects_untranslatable_tool_choice_with_a_400(tool_choice): - with pytest.raises(litellm.BadRequestError) as exc_info: - litellm.completion( - model="anthropic/claude-haiku-4-5", - messages=[{"role": "user", "content": "Which fruit is red?"}], - tools=[{"type": "function", "function": {"name": "lookup_fruit", "parameters": {"type": "object"}}}], - tool_choice=tool_choice, - api_key="sk-unused", - ) - assert exc_info.value.status_code == 400 - assert f"tool_choice={tool_choice}" in str(exc_info.value) diff --git a/tests/test_litellm/types/__init__.py b/tests/test_litellm/types/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/types/proxy/__init__.py b/tests/test_litellm/types/proxy/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/types/proxy/policy_engine/__init__.py b/tests/test_litellm/types/proxy/policy_engine/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/vector_stores/__init__.py b/tests/test_litellm/vector_stores/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/videos/__init__.py b/tests/test_litellm/videos/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index d1572f4a7c9..dd95addac40 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -2072,3 +2072,348 @@ def test_chat_rows_from_mistral_still_use_token_pricing(monkeypatch): ) assert result.cost == pytest.approx((10 * 0.001 + 5 * 0.002) / 2) assert result.usage.total_tokens == 15 + + +GROUNDED_USAGE_METADATA = { + "promptTokenCount": 19, + "candidatesTokenCount": 59, + "thoughtsTokenCount": 406, + "toolUsePromptTokenCount": 73, + "totalTokenCount": 557, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 19}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 59}], + "toolUsePromptTokensDetails": [{"modality": "TEXT", "tokenCount": 73}], + "trafficType": "ON_DEMAND", +} + + +PASSTHROUGH_OUTPUT_URI = ( + "gs://litellm-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/u/" + "predictions.jsonl" +) + + +UNGROUNDED_USAGE_METADATA = { + "promptTokenCount": 20, + "candidatesTokenCount": 48, + "thoughtsTokenCount": 195, + "toolUsePromptTokenCount": 73, + "totalTokenCount": 336, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 20}], + "trafficType": "ON_DEMAND", +} + + +def _native_vertex_row(usage_metadata: dict, *, grounded: bool, model_version: str | None = "gemini-2.5-flash"): + candidate = {"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"} + grounding = {"groundingMetadata": {"webSearchQueries": ["q"]}} if grounded else {} + response = {"candidates": [{**candidate, **grounding}], "usageMetadata": usage_metadata} + return { + "request": {"contents": [{"role": "user", "parts": [{"text": "q"}]}], "tools": [{"googleSearch": {}}]}, + "status": "", + "response": {**response, **({"modelVersion": model_version} if model_version else {})}, + "processed_time": "2026-09-23T19:02:00.000+00:00", + } + + +def _capture_cost_calls(monkeypatch, prompt_cost=0.5, completion_cost=0.25) -> list: + import litellm.cost_calculator as cc + + calls: list = [] + + def _calc(**kw): + calls.append(kw) + return (prompt_cost, completion_cost) + + monkeypatch.setattr(cc, "batch_cost_calculator", _calc) + return calls + + +def test_vertex_native_cost_bills_embedding_rows(monkeypatch): + monkeypatch.setitem(litellm.model_cost, "vertex_ai/gemini-embedding-2", {"input_cost_per_token_batches": 1e-7}) + rows = [ + { + "key": "id_1", + "status": "", + "request": {"content": {"parts": [{"text": "hello world"}]}}, + "response": {"embedding": {"values": [0.1, 0.2]}, "usageMetadata": {"promptTokenCount": 2}}, + }, + { + "key": "id_2", + "status": "", + "request": {"content": {"parts": [{"text": "hello"}]}}, + "response": {"embedding": {"values": [0.3]}, "tokenCount": "3"}, + }, + {"key": "id_3", "status": "INVALID_ARGUMENT", "request": {"content": {"parts": [{"text": ""}]}}}, + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-embedding-2") + + assert (result.successful_requests, result.failed_requests) == (2, 1) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (5, 0, 5) + assert result.cost == pytest.approx(5 * 1e-7) + assert result.models == ["gemini-embedding-2"] + + +@pytest.mark.asyncio +async def test_native_vertex_rows_route_to_vertex_cost_path_without_flag(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) + monkeypatch.setattr( + bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") + ) + calls = _capture_cost_calls(monkeypatch) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False), + ] + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" + ) + + assert result.cost == pytest.approx(1.5) + assert (result.successful_requests, result.failed_requests) == (2, 0) + assert result.models == ["gemini-2.5-flash"] + assert {(call["model"], call["custom_llm_provider"]) for call in calls} == {("gemini-2.5-flash", "vertex_ai")} + + +@pytest.mark.asyncio +async def test_openai_shaped_vertex_rows_keep_the_generic_path_without_flag(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) + monkeypatch.setattr( + bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run") + ) + _capture_cost_calls(monkeypatch) + rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)] + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" + ) + + assert result.successful_requests == 1 + + +@pytest.mark.asyncio +async def test_native_vertex_rows_on_another_provider_keep_the_generic_path(monkeypatch): + monkeypatch.setattr( + bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run") + ) + _capture_cost_calls(monkeypatch) + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], + custom_llm_provider="openai", + ) + + assert result.successful_requests == 0 + + +@pytest.mark.asyncio +async def test_handle_completed_batch_routes_native_rows_without_flag(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) + raw_rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)] + + async def fake_fetch(batch, custom_llm_provider, litellm_params=None): + return _vertex_jsonl(raw_rows) + + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) + monkeypatch.setattr( + bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") + ) + calls = _capture_cost_calls(monkeypatch, prompt_cost=0.7, completion_cost=0.3) + deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} + + result = await bu._handle_completed_batch( + _batch(PASSTHROUGH_OUTPUT_URI), + custom_llm_provider="vertex_ai", + model_name="gemini-2.5-flash", + model_info=deployment_model_info, + ) + + assert result.cost == pytest.approx(1.0) + assert result.usage.total_tokens == 557 + assert [call["model_info"] for call in calls] == [deployment_model_info] + + +def test_native_vertex_usage_is_billed_like_the_online_path(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + grounded = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True) + ungrounded = _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False) + + result = bu.calculate_vertex_ai_batch_cost_and_usage([grounded, ungrounded], "gemini-2.5-flash") + + grounded_usage, ungrounded_usage = (call["usage"] for call in calls) + assert grounded_usage.prompt_tokens == 19 + assert grounded_usage.completion_tokens == 59 + 406 + assert grounded_usage.completion_tokens_details.reasoning_tokens == 406 + assert ungrounded_usage.prompt_tokens == 20 + 73 + assert ungrounded_usage.completion_tokens == 48 + 195 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( + 19 + 93, + 465 + 243, + 557 + 336, + ) + + +def test_native_vertex_rows_are_priced_by_model_version_without_a_model_name(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-pro"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows) + + assert [call["model"] for call in calls] == ["gemini-2.5-flash", "gemini-2.5-pro"] + assert result.models == ["gemini-2.5-flash", "gemini-2.5-pro"] + assert result.cost == pytest.approx(1.5) + assert result.successful_requests == 3 + assert result.usage.total_tokens == 557 + 336 + 336 + + +def test_native_vertex_rows_without_usage_metadata_count_as_failed(monkeypatch): + _capture_cost_calls(monkeypatch) + rows = [ + {"request": {"contents": []}, "status": "Error: bad request", "processed_time": "t"}, + {"request": {"contents": []}, "response": {"candidates": []}}, + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert (result.successful_requests, result.failed_requests) == (1, 2) + assert result.usage.total_tokens == 557 + + +def test_native_vertex_batch_whose_rows_all_failed_still_names_the_deployment_model(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [{"request": {"contents": []}, "status": "Error: quota exceeded", "processed_time": "t"}] * 2 + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert result.models == ["gemini-2.5-flash"] + assert (result.successful_requests, result.failed_requests, result.cost) == (0, 2, 0.0) + assert calls == [] + + +def test_native_vertex_rows_are_priced_with_the_deployment_model_info(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} + + bu.calculate_vertex_ai_batch_cost_and_usage( + [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], + "gemini-2.5-flash", + model_info=deployment_model_info, + ) + + assert [call["model_info"] for call in calls] == [deployment_model_info] + + +@pytest.mark.asyncio +async def test_native_vertex_rows_keep_the_deployment_model_info_through_the_batch_entrypoint(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + deployment_model_info = {"input_cost_per_token_batches": 1e-6} + + await bu.calculate_batch_cost_and_usage( + file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], + custom_llm_provider="vertex_ai", + model_name="gemini-2.5-flash", + model_info=deployment_model_info, + ) + + assert [call["model_info"] for call in calls] == [deployment_model_info] + + +def test_native_vertex_rows_are_priced_by_the_deployment_model_over_model_version(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-pro")] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert [call["model"] for call in calls] == ["gemini-2.5-flash"] + assert result.models == ["gemini-2.5-flash"] + + +def test_native_vertex_rows_that_fail_response_validation_count_as_failed(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [ + {"request": {"contents": []}, "response": {"candidates": "nope", "usageMetadata": GROUNDED_USAGE_METADATA}}, + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert (result.successful_requests, result.failed_requests) == (1, 1) + assert result.usage.total_tokens == 557 + assert len(calls) == 1 + + +@pytest.mark.parametrize("wildcard_model", ["*", "vertex_ai/*"]) +def test_native_vertex_rows_under_a_wildcard_deployment_are_priced_by_model_version(monkeypatch, wildcard_model): + calls = _capture_cost_calls(monkeypatch) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, wildcard_model) + + assert [call["model"] for call in calls] == ["gemini-2.5-flash", wildcard_model] + assert result.cost == pytest.approx(1.5) + assert (result.successful_requests, result.failed_requests) == (2, 0) + assert result.usage.total_tokens == 557 + 336 + + +def test_native_vertex_row_without_model_version_under_a_wildcard_deployment_bills_its_explicit_prices(): + deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} + with_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash") + without_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version=None) + + twin = bu.calculate_vertex_ai_batch_cost_and_usage([with_version], "vertex_ai/*", model_info=deployment_model_info) + both = bu.calculate_vertex_ai_batch_cost_and_usage( + [with_version, without_version], "vertex_ai/*", model_info=deployment_model_info + ) + + assert twin.cost > 0 + assert both.cost == pytest.approx(2 * twin.cost) + assert (both.successful_requests, both.failed_requests) == (2, 0) + + +def test_native_vertex_row_the_cost_map_cannot_price_is_billed_at_zero_and_the_rest_still_bills(monkeypatch): + import litellm.cost_calculator as cc + + def _calc(**kw): + if kw["model"] == "gemini-unpriced": + raise ValueError("no pricing") + return (0.5, 0.25) + + monkeypatch.setattr(cc, "batch_cost_calculator", _calc) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-unpriced"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-flash"), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows) + + assert result.cost == pytest.approx(0.75) + assert (result.successful_requests, result.failed_requests) == (2, 0) + assert result.usage.total_tokens == 557 + 336 + assert result.models == ["gemini-unpriced", "gemini-2.5-flash"] + + +@pytest.mark.asyncio +async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_known(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False) + monkeypatch.setattr( + bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") + ) + calls = _capture_cost_calls(monkeypatch) + rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)] + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" + ) + + assert calls == [] + assert (result.successful_requests, result.failed_requests) == (0, 1) diff --git a/tests/unit/chat_completions/test_dispatch.py b/tests/unit/chat_completions/test_dispatch.py index 2807ed7f8f7..40b1c0ef019 100644 --- a/tests/unit/chat_completions/test_dispatch.py +++ b/tests/unit/chat_completions/test_dispatch.py @@ -20,6 +20,8 @@ from litellm.rust_bridge.chat_completions.entrypoints import ( ) from litellm.rust_bridge.configuration import Rollout from litellm.types.utils import ModelResponse +from litellm.chat_completions import dispatch +from litellm.rust_bridge.catalog import Rules MESSAGES: Final = [{"role": "user", "content": "hi"}] PYTHON_RULES: Final = () @@ -256,3 +258,100 @@ async def test_public_acompletion_routes_through_dispatch(monkeypatch: pytest.Mo NATIVE_ACOMPLETION.reset() assert result is expected assert [request.model for request in captured] == ["gpt-4o"] + + +@pytest.mark.asyncio +async def test_public_completion_calls_keep_the_python_result() -> None: + sync_response: Final = litellm.completion(model="openai/test-model", messages=MESSAGES, mock_response="ok") + async_response: Final = await litellm.acompletion(model="openai/test-model", messages=MESSAGES, mock_response="ok") + + assert isinstance(sync_response, ModelResponse) + assert isinstance(async_response, ModelResponse) + assert sync_response.choices[0].message.content == "ok" + assert async_response.choices[0].message.content == "ok" + + +def test_sync_completion_request_projects_public_arguments() -> None: + rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),) + expected: Final = ModelResponse() + + def native( + request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object] + ) -> ModelResponse: + assert request.model == "test-model" + assert request.messages == MESSAGES + assert request.custom_llm_provider == "openai" + assert request.stream is True + return expected + + binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None) + binding.override(native) + response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision + ("test-model", MESSAGES), + {"custom_llm_provider": "openai", "stream": True}, + python=lambda *args, **kwargs: pytest.fail("required native route must handle this call"), + binding=binding, + native=lambda hook, request, args, kwargs: hook(request, args, kwargs), + rules=rules, + ) + + assert response is expected + + +@pytest.mark.asyncio +async def test_async_completion_falls_back_after_native_declines() -> None: + from litellm.rust_bridge.bindings import native_exception_types + + native_types: Final = native_exception_types() + if native_types is None: + pytest.skip("native bridge is unavailable") + declined, _ = native_types + expected: Final = ModelResponse() + rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_OPT_OUT),) + + async def native( + request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object] + ) -> ModelResponse: + raise declined("unsupported") + + async def python(*args: object, **kwargs: object) -> ModelResponse: + return expected + + binding: Final[NativeBinding[NativeAcompletion]] = NativeBinding("acompletion", validate=lambda _: None) + binding.override(native) + response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision + ("test-model", MESSAGES), + {}, + python=python, + binding=binding, + native=lambda hook, request, args, kwargs: hook(request, args, kwargs), + rules=rules, + ) + + assert response is expected + + +def test_internal_acompletion_marker_bypasses_native() -> None: + rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),) + expected: Final = ModelResponse() + + def python(*args: object, **kwargs: object) -> ModelResponse: + return expected + + def native( + request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object] + ) -> ModelResponse: + pytest.fail("acompletion's inner completion call must stay on Python") + + binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None) + binding.override(native) + response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision + ("test-model", MESSAGES), + {"custom_llm_provider": "openai", "acompletion": True}, + python=python, + binding=binding, + native=lambda hook, request, args, kwargs: hook(request, args, kwargs), + rules=rules, + ) + + assert response is expected diff --git a/tests/test_litellm/a2a_protocol/__init__.py b/tests/unit/completion_extras/litellm_responses_transformation/__init__.py similarity index 100% rename from tests/test_litellm/a2a_protocol/__init__.py rename to tests/unit/completion_extras/litellm_responses_transformation/__init__.py diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py similarity index 100% rename from tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py rename to tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py similarity index 100% rename from tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py rename to tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 202ecb80d7b..653b2c9914a 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,7 +1,11 @@ +import asyncio +import importlib import os -from collections.abc import Iterator +from collections.abc import Coroutine, Iterator +from pathlib import Path from typing import Final +import boto3 import pytest from pytest_socket import enable_socket, socket_allow_hosts @@ -10,6 +14,14 @@ os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" import litellm # noqa: E402 # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at import import litellm.router as litellm_router_module # noqa: E402 # same import-time dependency import litellm.utils as litellm_utils_module # noqa: E402 # same import-time dependency +from litellm._logging import ALL_LOGGERS # noqa: E402 # same import-time dependency +from litellm.litellm_core_utils.prompt_templates import ( # noqa: E402 # same import-time dependency + image_handling as image_handling_module, +) +from litellm.llms.custom_httpx.async_client_cleanup import ( # noqa: E402 # same import-time dependency + close_litellm_async_clients, +) +from litellm.proxy.db import tool_registry_writer as tool_registry_writer_module # noqa: E402 # same import-time dependency LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1", "localhost"] AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = ( @@ -20,6 +32,63 @@ AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = ( "AZURE_USERNAME", "AZURE_PASSWORD", ) +AMBIENT_AWS_ENV_VARS: Final = ( + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_CONTAINER_CREDENTIALS_FULL_URI", + "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", + "AWS_SESSION_TOKEN", + "AWS_ROLE_ARN", + "AWS_WEB_IDENTITY_TOKEN_FILE", + "AWS_BEARER_TOKEN_BEDROCK", + "AWS_REGION_NAME", + "AWS_DEFAULT_REGION", +) +MODULES_WITH_AWS_AUTH_HANDLERS: Final = ( + "litellm.main", + "litellm.files.main", + "litellm.rerank_api.main", + "litellm.realtime_api.main", +) +CALLBACK_LISTS: Final = ( + "callbacks", + "success_callback", + "failure_callback", + "input_callback", + "_async_success_callback", + "_async_failure_callback", + "_async_input_callback", +) +RESET_TO_NONE_GLOBALS: Final = ("model_fallbacks", "cache") +RESTORED_GLOBALS: Final = ( + "disable_aiohttp_transport", + "force_ipv4", + "drop_params", + "secret_manager_client", + "_key_management_system", + "_key_management_settings", + "api_base", + "num_retries", + "modify_params", + "ssl_verify", + "credential_list", + "model_group_settings", + "default_internal_user_params", + "default_team_params", + "prometheus_emit_stream_label", + "vector_store_registry", + "model_cost", + "cost_margin_config", + "cost_discount_config", + "disable_hf_tokenizer_download", + "disable_copilot_system_to_assistant", + "cohere_models", + "anthropic_models", + "token_counter", + "initialized_langfuse_clients", +) +MODULE_LEVEL_CLIENTS: Final = ("module_level_client", "module_level_aclient") +SESSION_CLIENTS: Final = ("base_llm_aiohttp_handler", "httpx_client", "aclient", "client") def _allow_loopback_only() -> None: @@ -29,11 +98,116 @@ def _allow_loopback_only() -> None: _allow_loopback_only() +def pytest_collectstart() -> None: + _allow_loopback_only() + + @pytest.hookimpl(trylast=True) def pytest_runtest_setup() -> None: _allow_loopback_only() +def _run_coroutine_if_needed(result: object) -> None: + if not asyncio.iscoroutine(result): + return + coroutine: Final[Coroutine[object, object, object]] = result + try: + asyncio.run(coroutine) + except RuntimeError: + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + coroutine.close() + return + loop.create_task(coroutine) + + +def _close_handler_if_needed(handler: object) -> None: + close: Final = getattr(handler, "close", None) + if not callable(close): + return + _run_coroutine_if_needed(close()) + + +def _reset_aws_auth_caches() -> None: + modules: Final = tuple(importlib.import_module(name) for name in MODULES_WITH_AWS_AUTH_HANDLERS) + flushes: Final = ( + getattr(getattr(getattr(module, attr_name), "iam_cache", None), "flush_cache", None) + for module in modules + for attr_name in dir(module) + ) + for flush in filter(callable, flushes): + flush() + boto3.DEFAULT_SESSION = None + + +def _flush_client_caches() -> None: + litellm.in_memory_llm_clients_cache.flush_cache() + image_handling_module.in_memory_cache.flush_cache() + _reset_aws_auth_caches() + + +@pytest.fixture(scope="session") +def isolated_aws_config_files(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: + aws_dir: Final = tmp_path_factory.mktemp("aws-config") + credentials: Final = aws_dir / "credentials" + config: Final = aws_dir / "config" + credentials.write_text("", encoding="utf-8") + config.write_text("", encoding="utf-8") + return credentials, config + + +@pytest.fixture(autouse=True) +def isolate_host_environment(isolated_aws_config_files: tuple[Path, Path]) -> Iterator[None]: + credentials, config = isolated_aws_config_files + with pytest.MonkeyPatch.context() as environment: + environment.setenv("AWS_SHARED_CREDENTIALS_FILE", str(credentials)) + environment.setenv("AWS_CONFIG_FILE", str(config)) + environment.setenv("AWS_EC2_METADATA_DISABLED", "true") + for name in AMBIENT_AWS_ENV_VARS: + environment.delenv(name, raising=False) + environment.delenv("PROXY_BASE_URL", raising=False) + environment.setenv("LITELLM_CLI_DISABLE_KEYRING", "1") + yield + + +@pytest.fixture(autouse=True) +def isolate_litellm_globals() -> Iterator[None]: + original_callbacks: Final = {name: list(getattr(litellm, name) or []) for name in CALLBACK_LISTS} + original_reset: Final = {name: getattr(litellm, name) for name in RESET_TO_NONE_GLOBALS} + original_restored: Final = {name: getattr(litellm, name) for name in RESTORED_GLOBALS if hasattr(litellm, name)} + original_clients: Final = {name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__} + original_loggers: Final = { + logger: (logger.level, logger.disabled, logger.propagate, list(logger.handlers), list(logger.filters)) + for logger in ALL_LOGGERS + } + original_tool_policy_registry: Final = tool_registry_writer_module._tool_policy_registry + _flush_client_caches() + for name in CALLBACK_LISTS: + setattr(litellm, name, []) + for name in RESET_TO_NONE_GLOBALS: + setattr(litellm, name, None) + for name in MODULE_LEVEL_CLIENTS: + litellm.__dict__.pop(name, None) + tool_registry_writer_module._tool_policy_registry = None + yield + _flush_client_caches() + leaked_clients: Final = tuple(litellm.__dict__.pop(name, None) for name in MODULE_LEVEL_CLIENTS) + for name, client in zip(MODULE_LEVEL_CLIENTS, leaked_clients): + if client is not original_clients.get(name): + _close_handler_if_needed(client) + litellm.__dict__.update(original_clients) + for name, value in (original_callbacks | original_reset | original_restored).items(): + setattr(litellm, name, value) + for logger, (level, disabled, propagate, handlers, filters) in original_loggers.items(): + logger.setLevel(level) + logger.disabled = disabled + logger.propagate = propagate + logger.handlers = handlers + logger.filters = filters + tool_registry_writer_module._tool_policy_registry = original_tool_policy_registry + + @pytest.fixture(autouse=True) def isolate_router_model_cost_state() -> Iterator[None]: original_live_routers: Final = frozenset(litellm_router_module._live_routers) @@ -41,6 +215,7 @@ def isolate_router_model_cost_state() -> Iterator[None]: model_key: dict(model_value) for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items() } + litellm_utils_module._invalidate_model_cost_lowercase_map() yield for router in tuple(litellm_router_module._live_routers): litellm_router_module._live_routers.discard(router) @@ -68,4 +243,9 @@ def no_ambient_azure_credentials(monkeypatch: pytest.MonkeyPatch) -> None: def pytest_sessionfinish() -> None: + for name in MODULE_LEVEL_CLIENTS: + _close_handler_if_needed(litellm.__dict__.pop(name, None)) + for name in SESSION_CLIENTS: + _close_handler_if_needed(getattr(litellm, name, None)) + _run_coroutine_if_needed(close_litellm_async_clients()) enable_socket() diff --git a/tests/test_litellm/a2a_protocol/providers/__init__.py b/tests/unit/containers/__init__.py similarity index 100% rename from tests/test_litellm/a2a_protocol/providers/__init__.py rename to tests/unit/containers/__init__.py diff --git a/tests/test_litellm/containers/test_azure_container_transformation.py b/tests/unit/containers/test_azure_container_transformation.py similarity index 100% rename from tests/test_litellm/containers/test_azure_container_transformation.py rename to tests/unit/containers/test_azure_container_transformation.py diff --git a/tests/test_litellm/containers/test_container_api.py b/tests/unit/containers/test_container_api.py similarity index 100% rename from tests/test_litellm/containers/test_container_api.py rename to tests/unit/containers/test_container_api.py diff --git a/tests/test_litellm/containers/test_container_handler_url.py b/tests/unit/containers/test_container_handler_url.py similarity index 100% rename from tests/test_litellm/containers/test_container_handler_url.py rename to tests/unit/containers/test_container_handler_url.py diff --git a/tests/test_litellm/containers/test_container_integration.py b/tests/unit/containers/test_container_integration.py similarity index 100% rename from tests/test_litellm/containers/test_container_integration.py rename to tests/unit/containers/test_container_integration.py diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/unit/containers/test_container_proxy_ownership.py similarity index 100% rename from tests/test_litellm/containers/test_container_proxy_ownership.py rename to tests/unit/containers/test_container_proxy_ownership.py diff --git a/tests/test_litellm/containers/test_container_regional_api_base.py b/tests/unit/containers/test_container_regional_api_base.py similarity index 100% rename from tests/test_litellm/containers/test_container_regional_api_base.py rename to tests/unit/containers/test_container_regional_api_base.py diff --git a/tests/test_litellm/containers/test_container_transformation.py b/tests/unit/containers/test_container_transformation.py similarity index 100% rename from tests/test_litellm/containers/test_container_transformation.py rename to tests/unit/containers/test_container_transformation.py diff --git a/tests/test_litellm/containers/test_container_utils.py b/tests/unit/containers/test_container_utils.py similarity index 100% rename from tests/test_litellm/containers/test_container_utils.py rename to tests/unit/containers/test_container_utils.py diff --git a/tests/test_litellm/containers/test_endpoint_factory.py b/tests/unit/containers/test_endpoint_factory.py similarity index 100% rename from tests/test_litellm/containers/test_endpoint_factory.py rename to tests/unit/containers/test_endpoint_factory.py diff --git a/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/__init__.py b/tests/unit/embeddings/__init__.py similarity index 100% rename from tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/__init__.py rename to tests/unit/embeddings/__init__.py diff --git a/tests/test_litellm/embeddings/test_dispatch.py b/tests/unit/embeddings/test_dispatch.py similarity index 100% rename from tests/test_litellm/embeddings/test_dispatch.py rename to tests/unit/embeddings/test_dispatch.py diff --git a/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py b/tests/unit/experimental_mcp_client/__init__.py similarity index 100% rename from tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py rename to tests/unit/experimental_mcp_client/__init__.py diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py similarity index 100% rename from tests/test_litellm/experimental_mcp_client/test_mcp_client.py rename to tests/unit/experimental_mcp_client/test_mcp_client.py diff --git a/tests/test_litellm/experimental_mcp_client/test_tools.py b/tests/unit/experimental_mcp_client/test_tools.py similarity index 100% rename from tests/test_litellm/experimental_mcp_client/test_tools.py rename to tests/unit/experimental_mcp_client/test_tools.py diff --git a/tests/test_litellm/batches/__init__.py b/tests/unit/files/__init__.py similarity index 100% rename from tests/test_litellm/batches/__init__.py rename to tests/unit/files/__init__.py diff --git a/tests/test_litellm/files/test_main.py b/tests/unit/files/test_main.py similarity index 100% rename from tests/test_litellm/files/test_main.py rename to tests/unit/files/test_main.py diff --git a/tests/test_litellm/chat_completions/__init__.py b/tests/unit/fixtures/__init__.py similarity index 100% rename from tests/test_litellm/chat_completions/__init__.py rename to tests/unit/fixtures/__init__.py diff --git a/tests/test_litellm/completion_extras/__init__.py b/tests/unit/fixtures/together_ai_sync/__init__.py similarity index 100% rename from tests/test_litellm/completion_extras/__init__.py rename to tests/unit/fixtures/together_ai_sync/__init__.py diff --git a/tests/test_litellm/fixtures/together_ai_sync/deprecations.md b/tests/unit/fixtures/together_ai_sync/deprecations.md similarity index 100% rename from tests/test_litellm/fixtures/together_ai_sync/deprecations.md rename to tests/unit/fixtures/together_ai_sync/deprecations.md diff --git a/tests/test_litellm/fixtures/together_ai_sync/models_serverless.json b/tests/unit/fixtures/together_ai_sync/models_serverless.json similarity index 100% rename from tests/test_litellm/fixtures/together_ai_sync/models_serverless.json rename to tests/unit/fixtures/together_ai_sync/models_serverless.json diff --git a/tests/test_litellm/containers/__init__.py b/tests/unit/google_genai/__init__.py similarity index 100% rename from tests/test_litellm/containers/__init__.py rename to tests/unit/google_genai/__init__.py diff --git a/tests/test_litellm/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py similarity index 100% rename from tests/test_litellm/google_genai/test_google_genai_adapter.py rename to tests/unit/google_genai/test_google_genai_adapter.py diff --git a/tests/test_litellm/google_genai/test_google_genai_adapter_fixes.py b/tests/unit/google_genai/test_google_genai_adapter_fixes.py similarity index 100% rename from tests/test_litellm/google_genai/test_google_genai_adapter_fixes.py rename to tests/unit/google_genai/test_google_genai_adapter_fixes.py diff --git a/tests/test_litellm/google_genai/test_google_genai_handler.py b/tests/unit/google_genai/test_google_genai_handler.py similarity index 76% rename from tests/test_litellm/google_genai/test_google_genai_handler.py rename to tests/unit/google_genai/test_google_genai_handler.py index bf037c59854..5361d91718d 100644 --- a/tests/test_litellm/google_genai/test_google_genai_handler.py +++ b/tests/unit/google_genai/test_google_genai_handler.py @@ -2,99 +2,13 @@ """ Test to verify the Google GenAI generate_content handler functionality """ -import json from unittest.mock import AsyncMock, MagicMock, patch import pytest -import litellm from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter -from litellm.types.utils import ModelResponse - - -def test_non_stream_response_when_stream_requested_sync(): - """ - Test that when a non-stream response is returned but streaming was requested, - the sync handler correctly transforms it to generate_content format. - """ - from litellm.types.utils import Choices - - # Mock a non-stream response (ModelResponse with valid choices) - mock_response = ModelResponse( - id="test-123", - choices=[ - Choices( - index=0, - message={"role": "assistant", "content": "Hello, world!"}, - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-3.5-turbo", - object="chat.completion", - ) - - # Create an instance of the adapter - adapter = GoogleGenAIAdapter() - - # Test the adapter's translate_completion_to_generate_content method directly - result = adapter.translate_completion_to_generate_content(mock_response) - - # Verify the result is a valid Google GenAI format response - assert "candidates" in result - assert isinstance(result["candidates"], list) - assert len(result["candidates"]) > 0 - candidate = result["candidates"][0] - assert "content" in candidate - assert "parts" in candidate["content"] - assert isinstance(candidate["content"]["parts"], list) - assert len(candidate["content"]["parts"]) > 0 - assert "text" in candidate["content"]["parts"][0] - assert candidate["content"]["parts"][0]["text"] == "Hello, world!" - - -@pytest.mark.asyncio -async def test_non_stream_response_when_stream_requested_async(): - """ - Test that when a non-stream response is returned but streaming was requested, - the async handler correctly transforms it to generate_content format. - """ - from litellm.types.utils import Choices - - # Mock a non-stream response (ModelResponse with valid choices) - mock_response = ModelResponse( - id="test-123", - choices=[ - Choices( - index=0, - message={"role": "assistant", "content": "Hello, world!"}, - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-3.5-turbo", - object="chat.completion", - ) - - # Create an instance of the adapter - adapter = GoogleGenAIAdapter() - - # Test the adapter's translate_completion_to_generate_content method directly - result = adapter.translate_completion_to_generate_content(mock_response) - - # Verify the result is a valid Google GenAI format response - assert "candidates" in result - assert isinstance(result["candidates"], list) - assert len(result["candidates"]) > 0 - candidate = result["candidates"][0] - assert "content" in candidate - assert "parts" in candidate["content"] - assert isinstance(candidate["content"]["parts"], list) - assert len(candidate["content"]["parts"]) > 0 - assert "text" in candidate["content"]["parts"][0] - assert candidate["content"]["parts"][0]["text"] == "Hello, world!" def test_stream_response_when_stream_requested_sync(): diff --git a/tests/test_litellm/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py similarity index 100% rename from tests/test_litellm/google_genai/test_google_genai_main.py rename to tests/unit/google_genai/test_google_genai_main.py diff --git a/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py b/tests/unit/google_genai/test_google_genai_streaming_iterator.py similarity index 100% rename from tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py rename to tests/unit/google_genai/test_google_genai_streaming_iterator.py diff --git a/tests/test_litellm/google_genai/test_google_genai_transformation.py b/tests/unit/google_genai/test_google_genai_transformation.py similarity index 100% rename from tests/test_litellm/google_genai/test_google_genai_transformation.py rename to tests/unit/google_genai/test_google_genai_transformation.py diff --git a/tests/test_litellm/endpoints/__init__.py b/tests/unit/images/__init__.py similarity index 100% rename from tests/test_litellm/endpoints/__init__.py rename to tests/unit/images/__init__.py diff --git a/tests/test_litellm/images/test_image_edit_extra_params.py b/tests/unit/images/test_image_edit_extra_params.py similarity index 100% rename from tests/test_litellm/images/test_image_edit_extra_params.py rename to tests/unit/images/test_image_edit_extra_params.py diff --git a/tests/test_litellm/images/test_image_edit_utils.py b/tests/unit/images/test_image_edit_utils.py similarity index 100% rename from tests/test_litellm/images/test_image_edit_utils.py rename to tests/unit/images/test_image_edit_utils.py diff --git a/tests/test_litellm/images/test_image_generation_extra_headers.py b/tests/unit/images/test_image_generation_extra_headers.py similarity index 100% rename from tests/test_litellm/images/test_image_generation_extra_headers.py rename to tests/unit/images/test_image_generation_extra_headers.py diff --git a/tests/test_litellm/endpoints/speech/__init__.py b/tests/unit/interactions/__init__.py similarity index 100% rename from tests/test_litellm/endpoints/speech/__init__.py rename to tests/unit/interactions/__init__.py diff --git a/tests/test_litellm/interactions/test_agents_http_handler.py b/tests/unit/interactions/test_agents_http_handler.py similarity index 100% rename from tests/test_litellm/interactions/test_agents_http_handler.py rename to tests/unit/interactions/test_agents_http_handler.py diff --git a/tests/test_litellm/interactions/test_agents_main_and_utils.py b/tests/unit/interactions/test_agents_main_and_utils.py similarity index 100% rename from tests/test_litellm/interactions/test_agents_main_and_utils.py rename to tests/unit/interactions/test_agents_main_and_utils.py diff --git a/tests/test_litellm/interactions/test_background_cost_polling.py b/tests/unit/interactions/test_background_cost_polling.py similarity index 100% rename from tests/test_litellm/interactions/test_background_cost_polling.py rename to tests/unit/interactions/test_background_cost_polling.py diff --git a/tests/test_litellm/interactions/test_gemini_interactions_transformation.py b/tests/unit/interactions/test_gemini_interactions_transformation.py similarity index 100% rename from tests/test_litellm/interactions/test_gemini_interactions_transformation.py rename to tests/unit/interactions/test_gemini_interactions_transformation.py diff --git a/tests/test_litellm/interactions/test_interactions_streaming_iterator.py b/tests/unit/interactions/test_interactions_streaming_iterator.py similarity index 100% rename from tests/test_litellm/interactions/test_interactions_streaming_iterator.py rename to tests/unit/interactions/test_interactions_streaming_iterator.py diff --git a/tests/unit/interactions/test_litellm_responses_bridge.py b/tests/unit/interactions/test_litellm_responses_bridge.py new file mode 100644 index 00000000000..3abd0a6ca98 --- /dev/null +++ b/tests/unit/interactions/test_litellm_responses_bridge.py @@ -0,0 +1,80 @@ +""" +Tests for LiteLLM Responses bridge provider. + +Inherits from BaseInteractionsTest to run the same test suite against +the litellm_responses bridge provider, which calls litellm.responses() internally. +""" + + +from litellm.interactions.litellm_responses_transformation.transformation import ( + LiteLLMResponsesInteractionsConfig, +) +from litellm.types.interactions import Turn + + +class TestBridgeInputTransformation: + """Regression tests for translating Interactions input into Responses API input. + + The bridge used to pass Google content parts through raw ({"type": "text"}), + which the Responses API rejects with a 400, and it dropped the role encoded + in step types and in the legacy "model" turn role. + """ + + def test_step_input_maps_roles_and_content_types(self): + transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( + [ + {"type": "user_input", "content": [{"type": "text", "text": "I like apples."}]}, + {"type": "model_output", "content": [{"type": "text", "text": "I like oranges."}]}, + {"type": "user_input", "content": [{"type": "text", "text": "What did you say?"}]}, + ] + ) + assert transformed == [ + {"role": "user", "content": [{"type": "input_text", "text": "I like apples."}]}, + {"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]}, + {"role": "user", "content": [{"type": "input_text", "text": "What did you say?"}]}, + ] + + def test_legacy_turn_input_maps_model_role_to_assistant(self): + transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( + [ + {"role": "user", "content": [{"type": "text", "text": "I like apples."}]}, + {"role": "model", "content": [{"type": "text", "text": "I like oranges."}]}, + ] + ) + assert transformed == [ + {"role": "user", "content": [{"type": "input_text", "text": "I like apples."}]}, + {"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]}, + ] + + def test_turn_pydantic_model_with_string_content(self): + transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( + [Turn(role="model", content="I like oranges.")] + ) + assert transformed == [ + {"role": "assistant", "content": [{"type": "output_text", "text": "I like oranges."}]} + ] + + def test_string_input_passes_through(self): + transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input("Hello") + assert transformed == "Hello" + + def test_content_list_input_becomes_single_user_message(self): + transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( + [{"type": "text", "text": "Hello"}, "world"] + ) + assert transformed == [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Hello"}, + {"type": "input_text", "text": "world"}, + ], + } + ] + + def test_non_text_content_passes_through_unchanged(self): + image_part = {"type": "image", "data": "base64data", "mime_type": "image/png"} + transformed = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( + [{"type": "user_input", "content": [image_part]}] + ) + assert transformed == [{"role": "user", "content": [image_part]}] diff --git a/tests/test_litellm/interactions/test_openapi_compliance.py b/tests/unit/interactions/test_openapi_compliance.py similarity index 99% rename from tests/test_litellm/interactions/test_openapi_compliance.py rename to tests/unit/interactions/test_openapi_compliance.py index 2665f8703a6..d3f1183cea6 100644 --- a/tests/test_litellm/interactions/test_openapi_compliance.py +++ b/tests/unit/interactions/test_openapi_compliance.py @@ -4,7 +4,7 @@ OpenAPI compliance tests for Google Interactions API. Validates that our SDK requests/responses match the OpenAPI spec at: https://ai.google.dev/static/api/interactions.openapi.json -Run with: pytest tests/test_litellm/interactions/test_openapi_compliance.py -v +Run with: pytest tests/unit/interactions/test_openapi_compliance.py -v """ import json diff --git a/tests/unit/messages/test_dispatch.py b/tests/unit/messages/test_dispatch.py index 88ef849f0e2..3d5059b200f 100644 --- a/tests/unit/messages/test_dispatch.py +++ b/tests/unit/messages/test_dispatch.py @@ -22,6 +22,8 @@ from litellm.rust_bridge.messages.entrypoints import ( NativeMessages, ) from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse +from pydantic import TypeAdapter +from litellm.messages import dispatch MESSAGES: Final = [{"role": "user", "content": "hi"}] PYTHON_RULES: Final[Rules] = () @@ -284,3 +286,137 @@ async def test_anthropic_acreate_routes_through_dispatch(monkeypatch: pytest.Mon NATIVE_AMESSAGES.reset() assert result is expected assert [request.model for request in captured] == ["claude-sonnet-4-5"] + + +@pytest.mark.asyncio +async def test_public_anthropic_messages_keeps_the_python_result() -> None: + response: Final = await litellm.anthropic_messages( + model="anthropic/claude-sonnet-4-5", messages=MESSAGES, max_tokens=10, mock_response="ok" + ) + + assert isinstance(response, dict) + content: Final = TypeAdapter(list[dict[str, object]]).validate_python(response.get("content", [])) + assert content[0]["text"] == "ok" + + +def test_sync_messages_request_projects_public_arguments() -> None: + rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),) + expected: Final = AnthropicMessagesResponse(model="claude-test") + + def native( + request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object] + ) -> AnthropicMessagesResponse: + assert request.model == "claude-test" + assert request.messages == MESSAGES + assert request.max_tokens == 10 + assert request.custom_llm_provider == "anthropic" + return expected + + binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None) + binding.override(native) + response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision + (), + { + "model": "claude-test", + "messages": MESSAGES, + "max_tokens": 10, + "custom_llm_provider": "anthropic", + }, + python=lambda *args, **kwargs: pytest.fail("required native route must handle this call"), + binding=binding, + native=lambda hook, request, args, kwargs: hook(request, args, kwargs), + rules=rules, + ) + + assert response is expected + + +def test_messages_binding_error_delegates_unchanged_to_python() -> None: + rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),) + expected: Final = AnthropicMessagesResponse(model="claude-test") + + def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse: + return expected + + def native( + request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object] + ) -> AnthropicMessagesResponse: + pytest.fail("a call without max_tokens cannot project a request and must stay on Python") + + binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None) + binding.override(native) + response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision + (), + {"model": "claude-test", "messages": MESSAGES, "custom_llm_provider": "anthropic"}, + python=python, + binding=binding, + native=lambda hook, request, args, kwargs: hook(request, args, kwargs), + rules=rules, + ) + + assert response is expected + + +@pytest.mark.asyncio +async def test_async_messages_falls_back_after_native_declines() -> None: + from litellm.rust_bridge.bindings import native_exception_types + + native_types: Final = native_exception_types() + if native_types is None: + pytest.skip("native bridge is unavailable") + declined, _ = native_types + expected: Final = AnthropicMessagesResponse(model="claude-test") + rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_OUT),) + + async def native( + request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object] + ) -> AnthropicMessagesResponse: + raise declined("unsupported") + + async def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse: + return expected + + binding: Final[NativeBinding[NativeAmessages]] = NativeBinding("amessages", validate=lambda _: None) + binding.override(native) + response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision + (), + {"model": "claude-test", "messages": MESSAGES, "max_tokens": 10}, + python=python, + binding=binding, + native=lambda hook, request, args, kwargs: hook(request, args, kwargs), + rules=rules, + ) + + assert response is expected + + +def test_internal_is_async_marker_bypasses_native() -> None: + rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),) + expected: Final = AnthropicMessagesResponse(model="claude-test") + + def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse: + return expected + + def native( + request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object] + ) -> AnthropicMessagesResponse: + pytest.fail("anthropic_messages' inner handler call must stay on Python") + + binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None) + binding.override(native) + response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision + (), + { + "model": "claude-test", + "messages": MESSAGES, + "max_tokens": 10, + "custom_llm_provider": "anthropic", + "is_async": True, + }, + python=python, + binding=binding, + native=lambda hook, request, args, kwargs: hook(request, args, kwargs), + rules=rules, + ) + + assert response is expected diff --git a/tests/test_litellm/rag/test_main.py b/tests/unit/rag/test_main.py similarity index 100% rename from tests/test_litellm/rag/test_main.py rename to tests/unit/rag/test_main.py diff --git a/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/__init__.py b/tests/unit/rerank_api/__init__.py similarity index 100% rename from tests/test_litellm/endpoints/speech/speech_to_completion_bridge/__init__.py rename to tests/unit/rerank_api/__init__.py diff --git a/tests/test_litellm/rerank_api/test_main.py b/tests/unit/rerank_api/test_main.py similarity index 100% rename from tests/test_litellm/rerank_api/test_main.py rename to tests/unit/rerank_api/test_main.py diff --git a/tests/test_litellm/test_a2a_registry_lookup.py b/tests/unit/test_a2a_registry_lookup.py similarity index 100% rename from tests/test_litellm/test_a2a_registry_lookup.py rename to tests/unit/test_a2a_registry_lookup.py diff --git a/tests/test_litellm/test_acompletion_session_reuse_e2e.py b/tests/unit/test_acompletion_session_reuse_e2e.py similarity index 100% rename from tests/test_litellm/test_acompletion_session_reuse_e2e.py rename to tests/unit/test_acompletion_session_reuse_e2e.py diff --git a/tests/test_litellm/test_add_deployment_no_master_key.py b/tests/unit/test_add_deployment_no_master_key.py similarity index 100% rename from tests/test_litellm/test_add_deployment_no_master_key.py rename to tests/unit/test_add_deployment_no_master_key.py diff --git a/tests/test_litellm/test_aembedding_session_reuse_e2e.py b/tests/unit/test_aembedding_session_reuse_e2e.py similarity index 100% rename from tests/test_litellm/test_aembedding_session_reuse_e2e.py rename to tests/unit/test_aembedding_session_reuse_e2e.py diff --git a/tests/test_litellm/test_anthropic_beta_headers_filtering.py b/tests/unit/test_anthropic_beta_headers_filtering.py similarity index 100% rename from tests/test_litellm/test_anthropic_beta_headers_filtering.py rename to tests/unit/test_anthropic_beta_headers_filtering.py diff --git a/tests/test_litellm/test_anthropic_skills_transformation.py b/tests/unit/test_anthropic_skills_transformation.py similarity index 100% rename from tests/test_litellm/test_anthropic_skills_transformation.py rename to tests/unit/test_anthropic_skills_transformation.py diff --git a/tests/test_litellm/test_assert_ci_coverage.py b/tests/unit/test_assert_ci_coverage.py similarity index 100% rename from tests/test_litellm/test_assert_ci_coverage.py rename to tests/unit/test_assert_ci_coverage.py diff --git a/tests/test_litellm/test_assert_workflow_dir_hygiene.py b/tests/unit/test_assert_workflow_dir_hygiene.py similarity index 100% rename from tests/test_litellm/test_assert_workflow_dir_hygiene.py rename to tests/unit/test_assert_workflow_dir_hygiene.py diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/unit/test_audio_transcription_rust_bridge.py similarity index 100% rename from tests/test_litellm/test_audio_transcription_rust_bridge.py rename to tests/unit/test_audio_transcription_rust_bridge.py diff --git a/tests/test_litellm/test_auto_update_price_and_context_window_file.py b/tests/unit/test_auto_update_price_and_context_window_file.py similarity index 100% rename from tests/test_litellm/test_auto_update_price_and_context_window_file.py rename to tests/unit/test_auto_update_price_and_context_window_file.py diff --git a/tests/test_litellm/test_azure_ad_token_credential_resolution.py b/tests/unit/test_azure_ad_token_credential_resolution.py similarity index 100% rename from tests/test_litellm/test_azure_ad_token_credential_resolution.py rename to tests/unit/test_azure_ad_token_credential_resolution.py diff --git a/tests/test_litellm/test_azure_ai_grok_4_3_model_metadata.py b/tests/unit/test_azure_ai_grok_4_3_model_metadata.py similarity index 100% rename from tests/test_litellm/test_azure_ai_grok_4_3_model_metadata.py rename to tests/unit/test_azure_ai_grok_4_3_model_metadata.py diff --git a/tests/test_litellm/test_azure_ai_grok_4_6_model_metadata.py b/tests/unit/test_azure_ai_grok_4_6_model_metadata.py similarity index 100% rename from tests/test_litellm/test_azure_ai_grok_4_6_model_metadata.py rename to tests/unit/test_azure_ai_grok_4_6_model_metadata.py diff --git a/tests/test_litellm/test_baseten_glm_5_3_model_metadata.py b/tests/unit/test_baseten_glm_5_3_model_metadata.py similarity index 100% rename from tests/test_litellm/test_baseten_glm_5_3_model_metadata.py rename to tests/unit/test_baseten_glm_5_3_model_metadata.py diff --git a/tests/test_litellm/test_batch_completion_models_all_responses.py b/tests/unit/test_batch_completion_models_all_responses.py similarity index 100% rename from tests/test_litellm/test_batch_completion_models_all_responses.py rename to tests/unit/test_batch_completion_models_all_responses.py diff --git a/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py b/tests/unit/test_bedrock_marengo_embed_3_model_metadata.py similarity index 100% rename from tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py rename to tests/unit/test_bedrock_marengo_embed_3_model_metadata.py diff --git a/tests/test_litellm/test_budget_ratchet_check.py b/tests/unit/test_budget_ratchet_check.py similarity index 100% rename from tests/test_litellm/test_budget_ratchet_check.py rename to tests/unit/test_budget_ratchet_check.py diff --git a/tests/test_litellm/test_chat_ui_responses_session.py b/tests/unit/test_chat_ui_responses_session.py similarity index 100% rename from tests/test_litellm/test_chat_ui_responses_session.py rename to tests/unit/test_chat_ui_responses_session.py diff --git a/tests/test_litellm/test_check_licenses.py b/tests/unit/test_check_licenses.py similarity index 100% rename from tests/test_litellm/test_check_licenses.py rename to tests/unit/test_check_licenses.py diff --git a/tests/test_litellm/test_check_mcp_operation_boundary.py b/tests/unit/test_check_mcp_operation_boundary.py similarity index 100% rename from tests/test_litellm/test_check_mcp_operation_boundary.py rename to tests/unit/test_check_mcp_operation_boundary.py diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/unit/test_check_migrations_no_data_rewrites.py similarity index 100% rename from tests/test_litellm/test_check_migrations_no_data_rewrites.py rename to tests/unit/test_check_migrations_no_data_rewrites.py diff --git a/tests/test_litellm/test_check_py310_typing_imports.py b/tests/unit/test_check_py310_typing_imports.py similarity index 100% rename from tests/test_litellm/test_check_py310_typing_imports.py rename to tests/unit/test_check_py310_typing_imports.py diff --git a/tests/test_litellm/test_check_test_quality.py b/tests/unit/test_check_test_quality.py similarity index 100% rename from tests/test_litellm/test_check_test_quality.py rename to tests/unit/test_check_test_quality.py diff --git a/tests/test_litellm/test_check_type_discipline.py b/tests/unit/test_check_type_discipline.py similarity index 100% rename from tests/test_litellm/test_check_type_discipline.py rename to tests/unit/test_check_type_discipline.py diff --git a/tests/test_litellm/test_circleci_path_filter.py b/tests/unit/test_circleci_path_filter.py similarity index 100% rename from tests/test_litellm/test_circleci_path_filter.py rename to tests/unit/test_circleci_path_filter.py diff --git a/tests/test_litellm/test_circleci_rust_toolchain.py b/tests/unit/test_circleci_rust_toolchain.py similarity index 100% rename from tests/test_litellm/test_circleci_rust_toolchain.py rename to tests/unit/test_circleci_rust_toolchain.py diff --git a/tests/test_litellm/test_claude_fable_5_config.py b/tests/unit/test_claude_fable_5_config.py similarity index 100% rename from tests/test_litellm/test_claude_fable_5_config.py rename to tests/unit/test_claude_fable_5_config.py diff --git a/tests/test_litellm/test_claude_opus_4_6_config.py b/tests/unit/test_claude_opus_4_6_config.py similarity index 100% rename from tests/test_litellm/test_claude_opus_4_6_config.py rename to tests/unit/test_claude_opus_4_6_config.py diff --git a/tests/test_litellm/test_claude_opus_4_8_config.py b/tests/unit/test_claude_opus_4_8_config.py similarity index 100% rename from tests/test_litellm/test_claude_opus_4_8_config.py rename to tests/unit/test_claude_opus_4_8_config.py diff --git a/tests/test_litellm/test_claude_opus_5_config.py b/tests/unit/test_claude_opus_5_config.py similarity index 100% rename from tests/test_litellm/test_claude_opus_5_config.py rename to tests/unit/test_claude_opus_5_config.py diff --git a/tests/test_litellm/test_claude_sonnet_5_config.py b/tests/unit/test_claude_sonnet_5_config.py similarity index 100% rename from tests/test_litellm/test_claude_sonnet_5_config.py rename to tests/unit/test_claude_sonnet_5_config.py diff --git a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py b/tests/unit/test_cloudflare_workers_ai_model_metadata.py similarity index 100% rename from tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py rename to tests/unit/test_cloudflare_workers_ai_model_metadata.py diff --git a/tests/test_litellm/test_completion_timeout_resolution.py b/tests/unit/test_completion_timeout_resolution.py similarity index 100% rename from tests/test_litellm/test_completion_timeout_resolution.py rename to tests/unit/test_completion_timeout_resolution.py diff --git a/tests/test_litellm/test_component_entrypoint.py b/tests/unit/test_component_entrypoint.py similarity index 100% rename from tests/test_litellm/test_component_entrypoint.py rename to tests/unit/test_component_entrypoint.py diff --git a/tests/unit/test_compression.py b/tests/unit/test_compression.py new file mode 100644 index 00000000000..be718f03963 --- /dev/null +++ b/tests/unit/test_compression.py @@ -0,0 +1,649 @@ +""" +Unit tests for litellm.compress(). +""" + +import importlib + +import pytest + +import litellm +from litellm.compression.scoring.bm25 import bm25_score_messages +from litellm.compression.scoring.embedding_scorer import embedding_score_messages +from litellm.compression.content_detection import detect_content_type +from litellm.compression.message_stubbing import extract_key, stub_message +from litellm.compression.retrieval_tool import build_retrieval_tool +from litellm.types.utils import CallTypes + +CALL_TYPE = CallTypes.completion +ANTHROPIC_CALL_TYPE = CallTypes.anthropic_messages + + +# --------------------------------------------------------------------------- +# BM25 scorer +# --------------------------------------------------------------------------- + + +def test_bm25_relevance_ranking(): + query = "Fix the authentication bug in the login handler" + messages = [ + { + "role": "user", + "content": "def login_handler(): authentication check bug fix", + }, + {"role": "user", "content": "def render_template(name): css styling layout"}, + {"role": "user", "content": "def verify(): authentication token bug handler"}, + ] + scores = bm25_score_messages(query, messages) + # Messages sharing query terms should score higher than unrelated ones + assert scores[0] > scores[1] + assert scores[2] > scores[1] + + +def test_bm25_empty_query(): + scores = bm25_score_messages("", [{"role": "user", "content": "hello"}]) + assert scores == [0.0] + + +def test_bm25_empty_messages(): + scores = bm25_score_messages("query", []) + assert scores == [] + + +def test_bm25_empty_content(): + scores = bm25_score_messages("query", [{"role": "user", "content": ""}]) + assert scores == [0.0] + + +# --------------------------------------------------------------------------- +# Content detection +# --------------------------------------------------------------------------- + + +def test_detect_code(): + code = """ +import os +from pathlib import Path + +def main(): + class Foo: + pass + return Foo() +""" + assert detect_content_type(code) == "code" + + +def test_detect_json(): + assert detect_content_type('{"key": "value", "num": 42}') == "json" + assert detect_content_type("[1, 2, 3]") == "json" + + +def test_detect_text(): + assert detect_content_type("This is a plain text paragraph about dogs.") == "text" + + +def test_detect_empty(): + assert detect_content_type("") == "text" + + +# --------------------------------------------------------------------------- +# Message stubbing +# --------------------------------------------------------------------------- + + +def test_extract_key_with_filename(): + msg = {"role": "user", "content": "# auth.py\ndef authenticate():\n pass"} + used: set = set() + key = extract_key(msg, fallback_index=0, used_keys=used) + assert key == "auth.py" + + +def test_extract_key_fallback(): + msg = {"role": "user", "content": "Some random content without a filename"} + used: set = set() + key = extract_key(msg, fallback_index=5, used_keys=used) + assert key == "message_5" + + +def test_extract_key_duplicates(): + used: set = set() + msg = {"role": "user", "content": "# auth.py\ncode here"} + k1 = extract_key(msg, fallback_index=0, used_keys=used) + k2 = extract_key(msg, fallback_index=1, used_keys=used) + assert k1 == "auth.py" + assert k2 == "auth.py_2" + + +def test_stub_message(): + msg = {"role": "user", "content": "line1\nline2\nline3"} + stubbed = stub_message(msg, "test_key") + assert stubbed["role"] == "user" + assert "test_key" in stubbed["content"] + assert "litellm_content_retrieve" in stubbed["content"] + assert "3 lines" in stubbed["content"] + + +# --------------------------------------------------------------------------- +# Retrieval tool +# --------------------------------------------------------------------------- + + +def test_retrieval_tool_schema(): + tool = build_retrieval_tool(["auth.py", "utils.py"]) + assert tool["type"] == "function" + assert tool["function"]["name"] == "litellm_content_retrieve" + assert "key" in tool["function"]["parameters"]["properties"] + assert tool["function"]["parameters"]["properties"]["key"]["enum"] == [ + "auth.py", + "utils.py", + ] + assert tool["function"]["parameters"]["required"] == ["key"] + + +def test_retrieval_tool_description_lists_keys(): + tool = build_retrieval_tool(["foo.py", "bar.js"]) + desc = tool["function"]["description"] + assert "foo.py" in desc + assert "bar.js" in desc + + +# --------------------------------------------------------------------------- +# compress() — end-to-end +# --------------------------------------------------------------------------- + + +def test_compress_below_trigger_passthrough(): + messages = [{"role": "user", "content": "hello"}] + result = litellm.compress(messages, model="gpt-4o", call_type=CALL_TYPE) + assert result["messages"] == messages + assert result["cache"] == {} + assert result["tools"] == [] + assert result["compression_ratio"] == 0.0 + assert result["original_tokens"] == result["compressed_tokens"] + + +def test_compress_above_trigger(): + big_messages = [ + {"role": "system", "content": "You are a coding assistant."}, + { + "role": "user", + "content": "# auth.py\n" + "def authenticate():\n pass\n" * 2000, + }, + { + "role": "user", + "content": "# utils.py\n" + "def helper():\n pass\n" * 2000, + }, + { + "role": "user", + "content": "# readme.md\n" + "This is documentation. " * 2000, + }, + {"role": "user", "content": "Fix the bug in auth.py"}, + ] + + result = litellm.compress( + big_messages, + model="gpt-4o", + call_type=CALL_TYPE, + compression_trigger=1000, + compression_target=500, + ) + + assert result["compressed_tokens"] < result["original_tokens"] + assert result["compression_ratio"] > 0 + assert len(result["cache"]) > 0 + assert len(result["tools"]) == 1 + assert result["tools"][0]["function"]["name"] == "litellm_content_retrieve" + + +def test_compress_anthropic_list_content_is_boundary_stable(): + messages = [ + {"role": "system", "content": [{"type": "text", "text": "System prompt"}]}, + { + "role": "user", + "content": [ + {"type": "text", "text": "# a.py\n" + "alpha " * 2000}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/a.png"}, + }, + ], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "# b.py\n" + "beta " * 2000}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/b.png"}, + }, + ], + }, + { + "role": "user", + "content": [{"type": "text", "text": "Fix alpha bug in a.py"}], + }, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=1000, + compression_target=500, + ) + + assert result["compressed_tokens"] < result["original_tokens"] + assert len(result["messages"]) == len(messages) + assert [m["role"] for m in result["messages"]] == [m["role"] for m in messages] + assert len(result["cache"]) > 0 + assert len(result["tools"]) == 1 + assert result["tools"][0]["type"] == "custom" + assert result["tools"][0]["name"] == "litellm_content_retrieve" + assert "input_schema" in result["tools"][0] + + +def test_compress_preserves_system_message(): + messages = [ + {"role": "system", "content": "System prompt. " * 500}, + {"role": "user", "content": "Large file content. " * 5000}, + {"role": "user", "content": "Fix the bug"}, + ] + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) + assert result["messages"][0]["role"] == "system" + assert "System prompt" in result["messages"][0]["content"] + + +def test_compress_preserves_last_user_message(): + messages = [ + {"role": "user", "content": "Big context " * 5000}, + {"role": "user", "content": "Fix the bug in auth.py"}, + ] + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) + last_user = [m for m in result["messages"] if m["role"] == "user"][-1] + assert "Fix the bug in auth.py" in last_user["content"] + + +def test_compress_preserves_last_assistant_message(): + messages = [ + {"role": "user", "content": "Big context " * 5000}, + {"role": "assistant", "content": "I'll help with that. " * 2000}, + {"role": "user", "content": "Now fix the bug"}, + ] + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) + assistant_msgs = [m for m in result["messages"] if m["role"] == "assistant"] + assert len(assistant_msgs) >= 1 + # The last assistant message should be preserved (not stubbed) + last_assistant = assistant_msgs[-1] + assert "I'll help with that" in last_assistant["content"] + + +def test_cache_keys_match_stubs(): + messages = [ + {"role": "user", "content": "# auth.py\n" + "code " * 5000}, + {"role": "user", "content": "Fix it"}, + ] + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) + if result["tools"]: + tool_desc = result["tools"][0]["function"]["description"] + for key in result["cache"]: + assert key in tool_desc + + +def test_compress_default_target(): + """compression_target defaults to compression_trigger // 2.""" + messages = [ + {"role": "user", "content": "content " * 5000}, + {"role": "user", "content": "query"}, + ] + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=2000 + ) + # Should have compressed — target = 1000 + assert result["compressed_tokens"] <= result["original_tokens"] + + +def test_compress_nested_tool_result_extracts_text_only(): + messages = [ + {"role": "system", "content": [{"type": "text", "text": "System rules"}]}, + { + "role": "user", + "content": [ + {"type": "text", "text": "prefix"}, + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": [ + {"type": "text", "text": "nested text fragment"}, + { + "type": "image_url", + "image_url": { + "url": "https://example.com/secret-tool.png", + }, + }, + ], + }, + { + "type": "image_url", + "image_url": {"url": "https://example.com/top.png"}, + }, + {"type": "text", "text": " " + ("irrelevant " * 3000)}, + ], + }, + { + "role": "user", + "content": [{"type": "text", "text": "final query that must remain"}], + }, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=500, + compression_target=100, + ) + + cached_text = " ".join(result["cache"].values()) + assert "nested text fragment" in cached_text + assert "https://example.com/secret-tool.png" not in cached_text + assert "https://example.com/top.png" not in cached_text + + +def test_compress_default_call_type_is_completion(): + result = litellm.compress( + messages=[ + {"role": "user", "content": "Large context " * 4000}, + {"role": "user", "content": "query"}, + ], + model="gpt-4o", + compression_trigger=1000, + compression_target=500, + ) + + assert result["compressed_tokens"] <= result["original_tokens"] + assert isinstance(result["tools"], list) + + +def test_compress_forwards_embedding_model_params(monkeypatch): + captured = {} + + def fake_embedding_score_messages( + query, messages, model, cache=None, embedding_model_params=None + ): + captured["query"] = query + captured["model"] = model + captured["embedding_model_params"] = embedding_model_params + return [0.0] * len(messages) + + monkeypatch.setattr( + "litellm.compression.scoring.embedding_scorer.embedding_score_messages", + fake_embedding_score_messages, + ) + + result = litellm.compress( + messages=[ + {"role": "user", "content": "Authentication code " * 2000}, + {"role": "user", "content": "Fix auth"}, + ], + model="gpt-4o", + call_type=CALL_TYPE, + compression_trigger=1000, + embedding_model="text-embedding-3-small", + embedding_model_params={"api_base": "https://example-embeddings.test"}, + ) + + assert result["compressed_tokens"] <= result["original_tokens"] + assert captured["model"] == "text-embedding-3-small" + assert captured["embedding_model_params"] == { + "api_base": "https://example-embeddings.test" + } + + +def test_embedding_scorer_forwards_embedding_model_params(monkeypatch): + captured = {} + + class _MockResponse: + data = [ + {"embedding": [1.0, 0.0]}, + {"embedding": [1.0, 0.0]}, + {"embedding": [0.0, 1.0]}, + ] + + def fake_embedding(**kwargs): + captured.update(kwargs) + return _MockResponse() + + monkeypatch.setattr(litellm, "embedding", fake_embedding) + + scores = embedding_score_messages( + query="auth", + messages=[ + {"role": "user", "content": "auth code"}, + {"role": "user", "content": "cooking recipe"}, + ], + model="text-embedding-3-small", + embedding_model_params={"api_base": "https://example-embeddings.test"}, + ) + + assert len(scores) == 2 + assert captured["model"] == "text-embedding-3-small" + assert captured["api_base"] == "https://example-embeddings.test" + + +# --------------------------------------------------------------------------- +# Embedding scorer — integration test (skipped without API key) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "final_user_message, expected_content", + [ + ("How to cook?", "Unrelated cooking recipes "), + ("Fix auth", "Authentication code "), + ], +) +def test_simple_compression(final_user_message, expected_content): + messages = [ + {"role": "user", "content": "Authentication code " * 2000}, + {"role": "user", "content": "Unrelated cooking recipes " * 2000}, + {"role": "user", "content": final_user_message}, + ] + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) + if expected_content == "Unrelated cooking recipes ": + assert "Unrelated cooking recipes " in result["messages"][1]["content"] + assert "Authentication code " not in result["messages"][0]["content"] + elif expected_content == "Authentication code ": + assert "Authentication code " in result["messages"][0]["content"] + assert "Unrelated cooking recipes " not in result["messages"][1]["content"] + else: + raise ValueError(f"Unexpected expected_content: {expected_content}") + + +def test_compress_anthropic_drops_irrelevant_tool_exchange_span(monkeypatch): + compress_module = importlib.import_module("litellm.compression.compress") + + def fake_bm25_score_messages(query, messages): + assert "final query" in query + assert len(messages) == 5 + # Prefer idx=0 and de-prioritize the tool exchange span (idx=1,2) + return [0.95, 0.01, 0.02, 0.8, 1.0] + + def fake_token_counter(model, messages=None, text=None): + if messages is not None: + return 1000 + if text is None: + return 0 + if "final query" in text: + return 50 + if "assistant_tail" in text: + return 20 + if "other_blob" in text: + return 220 + if "tool_payload_relevant" in text: + return 200 + if text == "": + return 1 + return 10 + + monkeypatch.setattr( + compress_module, "bm25_score_messages", fake_bm25_score_messages + ) + monkeypatch.setattr(compress_module, "token_counter", fake_token_counter) + + messages = [ + {"role": "user", "content": "other_blob " * 300}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_drop", + "name": "litellm_content_retrieve", + "input": {"key": "message_1"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_drop", + "content": [{"type": "text", "text": "tool_payload_relevant"}], + } + ], + }, + {"role": "assistant", "content": "assistant_tail"}, + {"role": "user", "content": "final query"}, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=100, + compression_target=280, + ) + + # idx=1,2 should be dropped atomically (no orphan tool blocks left behind) + assert len(result["messages"]) == 3 + assert result["messages"][0]["role"] == "user" + assert "other_blob" in result["messages"][0]["content"] + assert result["messages"][1]["content"] == "assistant_tail" + assert result["messages"][2]["content"] == "final query" + assert result["cache"] == {} + + +def test_compress_anthropic_keeps_relevant_tool_exchange_span(monkeypatch): + compress_module = importlib.import_module("litellm.compression.compress") + + def fake_bm25_score_messages(query, messages): + assert "final query" in query + assert len(messages) == 5 + # Prefer the tool exchange span over idx=0 + return [0.05, 0.01, 0.92, 0.8, 1.0] + + def fake_token_counter(model, messages=None, text=None): + if messages is not None: + return 1000 + if text is None: + return 0 + if "final query" in text: + return 50 + if "assistant_tail" in text: + return 20 + if "other_blob" in text: + return 220 + if "tool_payload_relevant" in text: + return 200 + if text == "": + return 1 + return 10 + + monkeypatch.setattr( + compress_module, "bm25_score_messages", fake_bm25_score_messages + ) + monkeypatch.setattr(compress_module, "token_counter", fake_token_counter) + + messages = [ + {"role": "user", "content": "other_blob " * 300}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_keep", + "name": "litellm_content_retrieve", + "input": {"key": "message_1"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_keep", + "content": [{"type": "text", "text": "tool_payload_relevant"}], + } + ], + }, + {"role": "assistant", "content": "assistant_tail"}, + {"role": "user", "content": "final query"}, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=100, + compression_target=280, + ) + + assert len(result["messages"]) == 5 + assert result["messages"][1]["role"] == "assistant" + assert result["messages"][2]["role"] == "user" + # idx=0 should be compressed instead + assert "litellm_content_retrieve" in result["messages"][0]["content"] + assert len(result["cache"]) == 1 + + +def test_compress_anthropic_malformed_tool_sequence_passes_through(): + messages = [ + {"role": "user", "content": "other_blob " * 300}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_broken", + "name": "litellm_content_retrieve", + "input": {"key": "message_1"}, + } + ], + }, + {"role": "user", "content": [{"type": "text", "text": "missing tool_result"}]}, + {"role": "user", "content": "final query"}, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=100, + compression_target=280, + ) + + assert result["messages"] == messages + assert result["cache"] == {} + assert result["tools"] == [] + assert result["compression_skipped_reason"] == "invalid_anthropic_tool_sequence" diff --git a/tests/test_litellm/test_conftest_isolation.py b/tests/unit/test_conftest_isolation.py similarity index 100% rename from tests/test_litellm/test_conftest_isolation.py rename to tests/unit/test_conftest_isolation.py diff --git a/tests/test_litellm/test_constants.py b/tests/unit/test_constants.py similarity index 100% rename from tests/test_litellm/test_constants.py rename to tests/unit/test_constants.py diff --git a/tests/test_litellm/test_container_router.py b/tests/unit/test_container_router.py similarity index 100% rename from tests/test_litellm/test_container_router.py rename to tests/unit/test_container_router.py diff --git a/tests/test_litellm/test_cost_calculation_log_level.py b/tests/unit/test_cost_calculation_log_level.py similarity index 100% rename from tests/test_litellm/test_cost_calculation_log_level.py rename to tests/unit/test_cost_calculation_log_level.py diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/unit/test_cost_calculator.py similarity index 100% rename from tests/test_litellm/test_cost_calculator.py rename to tests/unit/test_cost_calculator.py diff --git a/tests/test_litellm/test_cost_map_guard.py b/tests/unit/test_cost_map_guard.py similarity index 100% rename from tests/test_litellm/test_cost_map_guard.py rename to tests/unit/test_cost_map_guard.py diff --git a/tests/test_litellm/test_count_tokens_public_api.py b/tests/unit/test_count_tokens_public_api.py similarity index 100% rename from tests/test_litellm/test_count_tokens_public_api.py rename to tests/unit/test_count_tokens_public_api.py diff --git a/tests/test_litellm/test_dashscope_image_generation.py b/tests/unit/test_dashscope_image_generation.py similarity index 99% rename from tests/test_litellm/test_dashscope_image_generation.py rename to tests/unit/test_dashscope_image_generation.py index 1dd0b322623..6f91fe9a0e0 100644 --- a/tests/test_litellm/test_dashscope_image_generation.py +++ b/tests/unit/test_dashscope_image_generation.py @@ -2,7 +2,7 @@ Unit tests for DashScope image generation support (qwen-image-2.0, qwen-image-2.0-pro, qwen-image-3.0, qwen-image-3.0-pro). -Run in docker: pytest tests/test_litellm/test_dashscope_image_generation.py -v +Run in docker: pytest tests/unit/test_dashscope_image_generation.py -v """ from unittest.mock import MagicMock, patch diff --git a/tests/test_litellm/test_daybreak_model_metadata.py b/tests/unit/test_daybreak_model_metadata.py similarity index 100% rename from tests/test_litellm/test_daybreak_model_metadata.py rename to tests/unit/test_daybreak_model_metadata.py diff --git a/tests/test_litellm/test_deepseek_model_metadata.py b/tests/unit/test_deepseek_model_metadata.py similarity index 100% rename from tests/test_litellm/test_deepseek_model_metadata.py rename to tests/unit/test_deepseek_model_metadata.py diff --git a/tests/test_litellm/test_default_branch.py b/tests/unit/test_default_branch.py similarity index 100% rename from tests/test_litellm/test_default_branch.py rename to tests/unit/test_default_branch.py diff --git a/tests/test_litellm/test_detect_changes.py b/tests/unit/test_detect_changes.py similarity index 100% rename from tests/test_litellm/test_detect_changes.py rename to tests/unit/test_detect_changes.py diff --git a/tests/test_litellm/test_dockerfile_apk_repository.py b/tests/unit/test_dockerfile_apk_repository.py similarity index 100% rename from tests/test_litellm/test_dockerfile_apk_repository.py rename to tests/unit/test_dockerfile_apk_repository.py diff --git a/tests/test_litellm/test_dockerfile_bedrock_realtime_extra.py b/tests/unit/test_dockerfile_bedrock_realtime_extra.py similarity index 100% rename from tests/test_litellm/test_dockerfile_bedrock_realtime_extra.py rename to tests/unit/test_dockerfile_bedrock_realtime_extra.py diff --git a/tests/test_litellm/test_dockerfile_non_root.py b/tests/unit/test_dockerfile_non_root.py similarity index 100% rename from tests/test_litellm/test_dockerfile_non_root.py rename to tests/unit/test_dockerfile_non_root.py diff --git a/tests/test_litellm/test_drop_params_env_var.py b/tests/unit/test_drop_params_env_var.py similarity index 100% rename from tests/test_litellm/test_drop_params_env_var.py rename to tests/unit/test_drop_params_env_var.py diff --git a/tests/test_litellm/test_e2e_egress_sentinel.py b/tests/unit/test_e2e_egress_sentinel.py similarity index 100% rename from tests/test_litellm/test_e2e_egress_sentinel.py rename to tests/unit/test_e2e_egress_sentinel.py diff --git a/tests/test_litellm/test_eager_tiktoken_load.py b/tests/unit/test_eager_tiktoken_load.py similarity index 100% rename from tests/test_litellm/test_eager_tiktoken_load.py rename to tests/unit/test_eager_tiktoken_load.py diff --git a/tests/test_litellm/test_env_key_doc_gate.py b/tests/unit/test_env_key_doc_gate.py similarity index 100% rename from tests/test_litellm/test_env_key_doc_gate.py rename to tests/unit/test_env_key_doc_gate.py diff --git a/tests/test_litellm/test_exception_exports.py b/tests/unit/test_exception_exports.py similarity index 100% rename from tests/test_litellm/test_exception_exports.py rename to tests/unit/test_exception_exports.py diff --git a/tests/test_litellm/test_exception_header_preservation.py b/tests/unit/test_exception_header_preservation.py similarity index 100% rename from tests/test_litellm/test_exception_header_preservation.py rename to tests/unit/test_exception_header_preservation.py diff --git a/tests/test_litellm/test_exception_mapping_request_attribute.py b/tests/unit/test_exception_mapping_request_attribute.py similarity index 100% rename from tests/test_litellm/test_exception_mapping_request_attribute.py rename to tests/unit/test_exception_mapping_request_attribute.py diff --git a/tests/test_litellm/test_filter_out_litellm_params.py b/tests/unit/test_filter_out_litellm_params.py similarity index 100% rename from tests/test_litellm/test_filter_out_litellm_params.py rename to tests/unit/test_filter_out_litellm_params.py diff --git a/tests/test_litellm/test_fireworks_serverless_model_costs.py b/tests/unit/test_fireworks_serverless_model_costs.py similarity index 100% rename from tests/test_litellm/test_fireworks_serverless_model_costs.py rename to tests/unit/test_fireworks_serverless_model_costs.py diff --git a/tests/test_litellm/test_gate_slot_lock.py b/tests/unit/test_gate_slot_lock.py similarity index 100% rename from tests/test_litellm/test_gate_slot_lock.py rename to tests/unit/test_gate_slot_lock.py diff --git a/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py b/tests/unit/test_gemini_3_1_flash_lite_image_pricing.py similarity index 100% rename from tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py rename to tests/unit/test_gemini_3_1_flash_lite_image_pricing.py diff --git a/tests/test_litellm/test_gemini_tts_native_audio_pricing.py b/tests/unit/test_gemini_tts_native_audio_pricing.py similarity index 100% rename from tests/test_litellm/test_gemini_tts_native_audio_pricing.py rename to tests/unit/test_gemini_tts_native_audio_pricing.py diff --git a/tests/test_litellm/test_get_blog_posts.py b/tests/unit/test_get_blog_posts.py similarity index 100% rename from tests/test_litellm/test_get_blog_posts.py rename to tests/unit/test_get_blog_posts.py diff --git a/tests/test_litellm/test_git_hooks.py b/tests/unit/test_git_hooks.py similarity index 100% rename from tests/test_litellm/test_git_hooks.py rename to tests/unit/test_git_hooks.py diff --git a/tests/test_litellm/test_gpt_5_4_model_metadata.py b/tests/unit/test_gpt_5_4_model_metadata.py similarity index 100% rename from tests/test_litellm/test_gpt_5_4_model_metadata.py rename to tests/unit/test_gpt_5_4_model_metadata.py diff --git a/tests/test_litellm/test_gpt_5_5_model_metadata.py b/tests/unit/test_gpt_5_5_model_metadata.py similarity index 100% rename from tests/test_litellm/test_gpt_5_5_model_metadata.py rename to tests/unit/test_gpt_5_5_model_metadata.py diff --git a/tests/test_litellm/test_gpt_image_cost_calculator.py b/tests/unit/test_gpt_image_cost_calculator.py similarity index 100% rename from tests/test_litellm/test_gpt_image_cost_calculator.py rename to tests/unit/test_gpt_image_cost_calculator.py diff --git a/tests/test_litellm/test_gpt_realtime_mode.py b/tests/unit/test_gpt_realtime_mode.py similarity index 100% rename from tests/test_litellm/test_gpt_realtime_mode.py rename to tests/unit/test_gpt_realtime_mode.py diff --git a/tests/test_litellm/test_groq_streaming_encoding.py b/tests/unit/test_groq_streaming_encoding.py similarity index 100% rename from tests/test_litellm/test_groq_streaming_encoding.py rename to tests/unit/test_groq_streaming_encoding.py diff --git a/tests/test_litellm/test_guardrail_exception_status_codes.py b/tests/unit/test_guardrail_exception_status_codes.py similarity index 100% rename from tests/test_litellm/test_guardrail_exception_status_codes.py rename to tests/unit/test_guardrail_exception_status_codes.py diff --git a/tests/test_litellm/test_lazy_imports.py b/tests/unit/test_lazy_imports.py similarity index 100% rename from tests/test_litellm/test_lazy_imports.py rename to tests/unit/test_lazy_imports.py diff --git a/tests/test_litellm/test_lint_workflow_diff_gates.py b/tests/unit/test_lint_workflow_diff_gates.py similarity index 100% rename from tests/test_litellm/test_lint_workflow_diff_gates.py rename to tests/unit/test_lint_workflow_diff_gates.py diff --git a/tests/test_litellm/test_litellm_params_reserved_keys.py b/tests/unit/test_litellm_params_reserved_keys.py similarity index 100% rename from tests/test_litellm/test_litellm_params_reserved_keys.py rename to tests/unit/test_litellm_params_reserved_keys.py diff --git a/tests/test_litellm/test_logging.py b/tests/unit/test_logging.py similarity index 100% rename from tests/test_litellm/test_logging.py rename to tests/unit/test_logging.py diff --git a/tests/test_litellm/test_lowest_latency_zero_tokens.py b/tests/unit/test_lowest_latency_zero_tokens.py similarity index 100% rename from tests/test_litellm/test_lowest_latency_zero_tokens.py rename to tests/unit/test_lowest_latency_zero_tokens.py diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py new file mode 100644 index 00000000000..effc038f85b --- /dev/null +++ b/tests/unit/test_main.py @@ -0,0 +1,4124 @@ +import asyncio +import base64 +from datetime import datetime +import contextlib +import copy +import json +import logging +import os +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +import httpx +import pytest +import respx + + +import urllib.parse +from importlib import import_module +from unittest.mock import MagicMock, patch + +import litellm +from litellm import main as litellm_main +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage + + +@pytest.fixture(autouse=True) +def clear_client_cache(): + """ + Clear the HTTP client cache before each test to ensure mocks are used. + This prevents cached real clients from being reused across tests. + """ + cache = getattr(litellm, "in_memory_llm_clients_cache", None) + if cache is not None: + cache.flush_cache() + yield + if cache is not None: + cache.flush_cache() + + +@pytest.fixture(autouse=True) +def add_api_keys_to_env(monkeypatch): + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-api03-1234567890") + monkeypatch.setenv("OPENAI_API_KEY", "sk-openai-api03-1234567890") + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "my-fake-aws-access-key-id") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "my-fake-aws-secret-access-key") + monkeypatch.setenv("AWS_REGION", "us-east-1") + # Keep these transformation tests on the simple access-key path. A leaked + # session token or role/web-identity env var pushes Bedrock auth down a + # different branch and fails before the mocked HTTP client is exercised. + monkeypatch.delenv("AWS_SESSION_TOKEN", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + + +@pytest.fixture +def openai_api_response(): + mock_response_data = { + "id": "chatcmpl-B0W3vmiM78Xkgx7kI7dr7PC949DMS", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "content": "", + "refusal": None, + "role": "assistant", + "audio": None, + "function_call": None, + "tool_calls": None, + }, + } + ], + "created": 1739462947, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion", + "service_tier": "default", + "system_fingerprint": "fp_bd83329f63", + "usage": { + "completion_tokens": 1, + "prompt_tokens": 121, + "total_tokens": 122, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, + }, + } + + return mock_response_data + + +def test_completion_missing_role(openai_api_response): + from openai import OpenAI + + from litellm.types.utils import ModelResponse + + client = OpenAI(api_key="test_api_key") + + mock_raw_response = MagicMock() + mock_raw_response.headers = { + "x-request-id": "123", + "openai-organization": "org-123", + "x-ratelimit-limit-requests": "100", + "x-ratelimit-remaining-requests": "99", + } + mock_raw_response.parse.return_value = ModelResponse(**openai_api_response) + + print(f"openai_api_response: {openai_api_response}") + + with patch.object( + client.chat.completions.with_raw_response, "create", MagicMock(return_value=mock_raw_response) + ) as mock_create: + litellm.completion( + model="gpt-4o-mini", + messages=[ + {"role": "user", "content": "Hey"}, + { + "content": "", + "tool_calls": [ + { + "id": "call_m0vFJjQmTH1McvaHBPR2YFwY", + "function": { + "arguments": '{"input": "dksjsdkjdhskdjshdskhjkhlk"}', + "name": "tool_name", + }, + "type": "function", + "index": 0, + }, + { + "id": "call_Vw6RaqV2n5aaANXEdp5pYxo2", + "function": { + "arguments": '{"input": "jkljlkjlkjlkjlk"}', + "name": "tool_name", + }, + "type": "function", + "index": 1, + }, + { + "id": "call_hBIKwldUEGlNh6NlSXil62K4", + "function": { + "arguments": '{"input": "jkjlkjlkjlkj;lj"}', + "name": "tool_name", + }, + "type": "function", + "index": 2, + }, + ], + }, + ], + client=client, + ) + + mock_create.assert_called_once() + + +@pytest.mark.parametrize("model", ["gpt-4o-mini"]) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_url_with_format_param_openai(model, sync_mode): + from openai import AsyncOpenAI, OpenAI + + from litellm import acompletion, completion + + if sync_mode: + client = OpenAI() + else: + client = AsyncOpenAI() + + args = { + "model": model, + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "https://awsmp-logos.s3.amazonaws.com/seller-xw5kijmvmzasy/c233c9ade2ccb5491072ae232c814942.png", + "format": "image/png", + }, + }, + {"type": "text", "text": "Describe this image"}, + ], + } + ], + } + with patch.object( + client.chat.completions.with_raw_response, "create" + ) as mock_client: + try: + if sync_mode: + response = completion(**args, client=client) + else: + response = await acompletion(**args, client=client) + print(response) + except Exception as e: + print(e) + + mock_client.assert_called() + + print(mock_client.call_args.kwargs) + + json_str = json.dumps(mock_client.call_args.kwargs) + + assert "format" not in json_str + + +def test_bedrock_latency_optimized_inference(): + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + client = HTTPHandler() + with patch.object(client, "post") as mock_post: + try: + response = litellm.completion( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "Hello, how are you?"}], + performanceConfig={"latency": "optimized"}, + client=client, + ) + except Exception as e: + print(e) + + mock_post.assert_called_once() + json_data = json.loads(mock_post.call_args.kwargs["data"]) + assert json_data["performanceConfig"]["latency"] == "optimized" + + +@pytest.mark.parametrize( + ("custom_llm_provider", "model", "expected"), + [ + ("anthropic", "claude-sonnet-5", True), + ("bedrock", "us.anthropic.claude-sonnet-5-20260501-v1:0", True), + ("bedrock", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", True), + ("bedrock", "us.amazon.nova-2-lite-v1:0", False), + ("vertex_ai", "claude-sonnet-5", True), + ("vertex_ai", "gemini-3.8-flash", False), + ("azure_ai", "claude-sonnet-4-6", True), + ("azure_ai", "gpt-5.6", False), + ("openai", "gpt-5.6", False), + ("gemini", "gemini-3.8-flash", False), + ], +) +def test_is_claude_tool_target(custom_llm_provider: str, model: str, expected: bool): + assert litellm_main._is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model) is expected + + +@pytest.mark.parametrize("key", ["input_examples", "eager_input_streaming"]) +def test_drop_anthropic_only_tool_keys_strips_tool_and_function_levels(key: str): + tools = [ + {"type": "function", "name": "example_tool", key: True, "function": {"name": "example_tool", key: True}}, + "opaque_tool", + ] + + cleaned = litellm_main._drop_anthropic_only_tool_keys(tools=tools) + + assert cleaned == [ + {"type": "function", "name": "example_tool", "function": {"name": "example_tool"}}, + "opaque_tool", + ] + assert tools[0][key] is True + assert tools[0]["function"][key] is True + + +def test_completion_strips_eager_input_streaming_before_openai(respx_mock: respx.MockRouter, openai_api_response): + api_base: Final = "http://localhost:12346/v1" + mock_route: Final = respx_mock.post(url__regex=rf"{api_base}/chat/completions.*").mock( + return_value=httpx.Response(status_code=200, json=openai_api_response) + ) + + litellm.completion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Write the file"}], + tools=[ + { + "type": "function", + "function": {"name": "write_file", "parameters": {"type": "object", "properties": {}}}, + "eager_input_streaming": True, + } + ], + api_base=api_base, + api_key="fake_openai_api_key", + ) + + assert mock_route.called + sent_tool: Final = json.loads(respx_mock.calls[0].request.content)["tools"][0] + assert "eager_input_streaming" not in sent_tool + assert sent_tool["function"]["name"] == "write_file" + + +def test_custom_provider_with_extra_headers(): + + with patch.object( + litellm.llms.custom_httpx.http_handler.HTTPHandler, "post" + ) as mock_post: + response = litellm.completion( + model="custom/custom", + messages=[{"role": "user", "content": "Hello, how are you?"}], + headers={"X-Custom-Header": "custom-value"}, + api_base="https://example.com/api/v1", + ) + + mock_post.assert_called_once() + assert mock_post.call_args[1]["headers"]["X-Custom-Header"] == "custom-value" + + +def test_custom_provider_with_extra_body(): + + with patch.object( + litellm.llms.custom_httpx.http_handler.HTTPHandler, "post" + ) as mock_post: + response = litellm.completion( + model="custom/custom", + messages=[{"role": "user", "content": "Hello, how are you?"}], + extra_body={ + "X-Custom-BodyValue": "custom-value", + "X-Custom-BodyValue2": "custom-value2", + }, + api_base="https://example.com/api/v1", + ) + mock_post.assert_called_once() + + assert mock_post.call_args[1]["json"]["X-Custom-BodyValue"] == "custom-value" + assert mock_post.call_args[1]["json"] == { + "model": "custom", + "params": { + "prompt": ["Hello, how are you?"], + "max_tokens": None, + "temperature": None, + "top_p": None, + "top_k": None, + }, + "X-Custom-BodyValue": "custom-value", + "X-Custom-BodyValue2": "custom-value2", + } + + # test that extra_body is not passed if not provided + with patch.object( + litellm.llms.custom_httpx.http_handler.HTTPHandler, "post" + ) as mock_post: + response = litellm.completion( + model="custom/custom", + messages=[{"role": "user", "content": "Hello, how are you?"}], + api_base="https://example.com/api/v1", + ) + mock_post.assert_called_once() + assert mock_post.call_args[1]["json"] == { + "model": "custom", + "params": { + "prompt": ["Hello, how are you?"], + "max_tokens": None, + "temperature": None, + "top_p": None, + "top_k": None, + }, + } + + +@pytest.fixture(autouse=True) +def set_openrouter_api_key(): + original_api_key = os.environ.get("OPENROUTER_API_KEY") + os.environ["OPENROUTER_API_KEY"] = "fake-key-for-testing" + yield + if original_api_key is not None: + os.environ["OPENROUTER_API_KEY"] = original_api_key + else: + del os.environ["OPENROUTER_API_KEY"] + + +@pytest.mark.asyncio +async def test_extra_body_with_fallback( + respx_mock: respx.MockRouter, set_openrouter_api_key, monkeypatch +): + """ + test regression for https://github.com/BerriAI/litellm/issues/8425. + + This was perhaps a wider issue with the acompletion function not passing kwargs such as extra_body correctly when fallbacks are specified. + """ + + # Save original state to restore after test + original_disable_aiohttp = litellm.disable_aiohttp_transport + + try: + # since this uses respx, we need to set use_aiohttp_transport to False + # Set both the global variable and environment variable to ensure it takes effect + litellm.disable_aiohttp_transport = True + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + # Flush cache to ensure no stale aiohttp clients are used + litellm.in_memory_llm_clients_cache.flush_cache() + + # Set up test parameters + model = "openrouter/deepseek/deepseek-chat" + messages = [{"role": "user", "content": "Hello, world!"}] + extra_body = { + "provider": { + "order": ["DeepSeek"], + "allow_fallbacks": False, + "require_parameters": True, + } + } + fallbacks = [{"model": "openrouter/google/gemini-flash-1.5-8b"}] + + # Set up mock to respond to any POST request to the OpenRouter endpoint + # This ensures it works for both primary and fallback models + mock_route = respx_mock.post("https://openrouter.ai/api/v1/chat/completions") + mock_route.return_value = httpx.Response( + 200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": model, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello from mocked response!", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 12, + "total_tokens": 21, + }, + }, + ) + + response = await litellm.acompletion( + model=model, + messages=messages, + extra_body=extra_body, + fallbacks=fallbacks, + api_key="fake-openrouter-api-key", + ) + + # Verify the response + assert response is not None + assert ( + len(respx_mock.calls) > 0 + ), "Mock was not called - check if aiohttp transport is properly disabled" + + # Get the request from the mock + request: httpx.Request = respx_mock.calls[0].request + request_body = request.read() + request_body = json.loads(request_body) + + # Verify basic parameters + assert request_body["model"] == "deepseek/deepseek-chat" + assert request_body["messages"] == messages + + # Verify the extra_body parameters remain under the provider key + assert request_body["provider"]["order"] == ["DeepSeek"] + assert request_body["provider"]["allow_fallbacks"] is False + assert request_body["provider"]["require_parameters"] is True + finally: + # Restore original state to prevent test pollution + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() + + +@pytest.mark.parametrize("env_base", ["OPENAI_BASE_URL", "OPENAI_API_BASE"]) +@pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=1) +async def test_openai_env_base( + respx_mock: respx.MockRouter, env_base, openai_api_response, monkeypatch +): + "This tests OpenAI env variables are honored, including legacy OPENAI_API_BASE" + # Ensure aiohttp transport is disabled to use httpx which respx can mock + litellm.disable_aiohttp_transport = True + + expected_base_url = "http://localhost:12345/v1" + + # Assign the environment variable based on env_base, and use a fake API key. + monkeypatch.setenv(env_base, expected_base_url) + monkeypatch.setenv("OPENAI_API_KEY", "fake_openai_api_key") + + model = "gpt-4o" + messages = [{"role": "user", "content": "Hello, how are you?"}] + + # Configure respx mock to intercept the request + mock_route = respx_mock.post( + url__regex=r"http://localhost:12345/v1/chat/completions.*" + ).mock( + return_value=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": model, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello from mocked response!", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 12, + "total_tokens": 21, + }, + }, + ) + ) + + try: + response = await litellm.acompletion(model=model, messages=messages) + + # verify we had a response + assert response.choices[0].message.content == "Hello from mocked response!" + + # Verify the mock was called + assert ( + mock_route.called + ), "Mock route was not called - request may have bypassed respx" + finally: + # Clean up to avoid affecting other tests + litellm.disable_aiohttp_transport = False + + +def build_database_url(username, password, host, dbname): + username_enc = urllib.parse.quote_plus(username) + password_enc = urllib.parse.quote_plus(password) + dbname_enc = urllib.parse.quote_plus(dbname) + return f"postgresql://{username_enc}:{password_enc}@{host}/{dbname_enc}" + + +def test_build_database_url(): + url = build_database_url("user@name", "p@ss:word", "localhost", "db/name") + assert url == "postgresql://user%40name:p%40ss%3Aword@localhost/db%2Fname" + + +def test_bedrock_llama(): + litellm._turn_on_debug() + from litellm.types.utils import CallTypes + from litellm.utils import return_raw_request + + model = "bedrock/invoke/us.meta.llama4-scout-17b-instruct-v1:0" + + request = return_raw_request( + endpoint=CallTypes.completion, + kwargs={ + "model": model, + "messages": [ + {"role": "user", "content": "hi"}, + ], + }, + ) + print(request) + + assert ( + request["raw_request_body"]["prompt"] + == "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\nhi<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" + ) + + +def _mocked_openai_chat_response(model: str) -> httpx.Response: + return httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": model, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello from mocked response!", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 12, + "total_tokens": 21, + }, + }, + ) + + +def test_return_raw_request_does_not_call_provider(respx_mock: respx.MockRouter): + """Regression for #33952: return_raw_request must transform without contacting the provider. + + Previously return_raw_request invoked the real endpoint with a fake key and relied on the + provider rejecting it, which sent an unintended inference request and (in the async proxy + route) blocked the event loop on provider I/O. + """ + from litellm.types.utils import CallTypes + from litellm.utils import return_raw_request + + model = "gpt-4o" + route = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=_mocked_openai_chat_response(model) + ) + + request = return_raw_request( + endpoint=CallTypes.completion, + kwargs={ + "model": model, + "messages": [{"role": "user", "content": "hi"}], + }, + ) + + assert route.call_count == 0 + assert request.get("error") is None + assert request["raw_request_body"]["model"] == model + assert request["raw_request_body"]["messages"] == [ + {"role": "user", "content": "hi"} + ] + + +def test_completion_forwards_verbosity_in_raw_request(respx_mock: respx.MockRouter): + """Regression test: completion() must forward the verbosity param to the provider request body.""" + from litellm.types.utils import CallTypes + from litellm.utils import return_raw_request + + model = "gpt-5.2" + messages = [{"role": "user", "content": "hi"}] + respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=_mocked_openai_chat_response(model) + ) + + request = return_raw_request( + endpoint=CallTypes.completion, + kwargs={ + "model": model, + "messages": messages, + "verbosity": "high", + }, + ) + + assert request["raw_request_body"]["verbosity"] == "high" + assert request["raw_request_body"]["model"] == model + assert request["raw_request_body"]["messages"] == messages + + +@pytest.mark.asyncio +async def test_acompletion_forwards_verbosity_to_provider_request( + respx_mock: respx.MockRouter, monkeypatch +): + """Regression test: acompletion() must forward the verbosity param to the provider request body.""" + original_disable_aiohttp = litellm.disable_aiohttp_transport + try: + litellm.disable_aiohttp_transport = True + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + + model = "gpt-5.2" + messages = [{"role": "user", "content": "hi"}] + mock_route = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=_mocked_openai_chat_response(model) + ) + + response = await litellm.acompletion( + model=model, + messages=messages, + verbosity="low", + api_key="fake-openai-api-key", + ) + + assert response.choices[0].message.content == "Hello from mocked response!" + assert mock_route.called + request_body = json.loads(respx_mock.calls[0].request.read()) + assert request_body["verbosity"] == "low" + assert request_body["model"] == model + assert request_body["messages"] == messages + finally: + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() + + +def test_responses_api_bridge_check_strips_responses_prefix(): + """Test that responses_api_bridge_check strips 'responses/' prefix and sets mode.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 4096} + + model_info, model = responses_api_bridge_check( + model="responses/gpt-4-responses", + custom_llm_provider="openai", + ) + + assert model == "gpt-4-responses" + assert model_info["mode"] == "responses" + + +def test_responses_api_bridge_check_gpt_5_4_pro(): + """Test that gpt-5.4-pro routes through responses API bridge, not chat completions. + + Regression test for https://github.com/BerriAI/litellm/issues/23014 + gpt-5.4-pro is a responses-only model and must not be sent to /v1/chat/completions. + """ + from litellm.main import responses_api_bridge_check + + for model_name in ["gpt-5.4-pro", "gpt-5.4-pro-2026-03-05"]: + model_info, model = responses_api_bridge_check( + model=model_name, + custom_llm_provider="openai", + ) + assert ( + model_info.get("mode") == "responses" + ), f"{model_name} should have mode='responses', got '{model_info.get('mode')}'" + + +def test_responses_api_bridge_check_gpt_5_4_tools_plus_reasoning_routes_to_responses(): + """gpt-5.4 with both tools and reasoning_effort should route to Responses API.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort="xhigh", + ) + + assert model == "gpt-5.4" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_gpt_6_astra_tools_with_default_reasoning_routes_to_responses(): + from litellm.main import responses_api_bridge_check + + model_info, model = responses_api_bridge_check( + model="gpt-6-astra", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + ) + + assert model == "gpt-6-astra" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_gpt_5_5_tools_plus_reasoning_routes_to_responses(): + """gpt-5.5+ with both tools and reasoning_effort should route to Responses API.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.5-pro", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort="xhigh", + ) + + assert model == "gpt-5.5-pro" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_azure_gpt_5_4_tools_plus_reasoning_routes_to_responses(): + """Azure gpt-5.4 with both tools and reasoning_effort should route to Responses API.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="azure", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort="high", + ) + + assert model == "gpt-5.4" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_azure_gpt_5_4_tools_with_default_reasoning_routes_to_responses(): + """ + Azure gpt-5.4 with tools and UNSET reasoning_effort must bridge: OpenAI enables + reasoning by default for gpt-5.4+, and Chat Completions rejects function tools + whenever reasoning is on. + """ + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="azure", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + ) + + assert model == "gpt-5.4" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_gpt_5_4_tools_with_default_reasoning_routes_to_responses(): + """ + gpt-5.4 with tools and UNSET reasoning_effort must bridge: OpenAI enables reasoning + by default for gpt-5.4+, and Chat Completions rejects function tools whenever + reasoning is on ("use /v1/responses or set reasoning_effort to 'none'"). + """ + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + ) + + assert model == "gpt-5.4" + assert model_info.get("mode") == "responses" + + +@pytest.mark.parametrize( + "model_name, expected_mode", + [ + pytest.param("gpt-5.6-sol", "responses", id="above-boundary-bridges"), + pytest.param("gpt-5.1", None, id="below-boundary-stays-chat"), + ], +) +def test_responses_api_bridge_check_gpt_5_6_tools_with_default_reasoning_routes_to_responses( + monkeypatch, model_name, expected_mode +): + """ + gpt-5.6 must bridge on function tools alone. The bridge used to require an explicit + reasoning_effort, so a gpt-5.6 call carrying tools and no effort was rejected with + "Function tools with reasoning_effort are not supported for gpt-5.6-sol in + /v1/chat/completions". + + Paired with a model below the gpt-5.4 boundary, which must still stay on chat. The + gate parses the version and drops any suffix, so the family members bridge + identically and only the boundary distinguishes behaviour. + """ + import litellm + from litellm.main import responses_api_bridge_check + + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", None) + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model=model_name, + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + ) + + assert model == model_name + assert model_info.get("mode") == expected_mode + + +def test_responses_api_bridge_check_gpt_5_4_tools_with_reasoning_none_stays_chat(): + """ + Explicit reasoning_effort "none" is OpenAI's documented escape hatch that keeps + function tools servable on Chat Completions; the bridge must not fire. + """ + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort="none", + ) + + assert model == "gpt-5.4" + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_reasoning_none_with_summary_still_routes_to_responses(): + """A reasoning summary is Responses-only regardless of effort value.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="openai", + reasoning_effort="none", + reasoning_summary="detailed", + ) + + assert model == "gpt-5.4" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_gpt_5_4_custom_tools_only_stays_chat(): + """ + Chat Completions serves custom (grammar) tools natively with reasoning on; only + FUNCTION tools trigger the OpenAI rejection. Custom-only requests must stay on chat + so responses keep the native custom tool_call shape instead of the bridge's + function-shaped mapping. + """ + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "custom", "custom": {"name": "ApplyPatch", "description": "V4A patch"}}], + reasoning_effort=None, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_gpt_5_4_mixed_function_and_custom_tools_routes_to_responses(): + """One function tool in the mix is enough to make chat unservable with reasoning on.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[ + {"type": "custom", "custom": {"name": "ApplyPatch"}}, + {"type": "function", "function": {"name": "shell"}}, + ], + reasoning_effort=None, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_gpt_5_4_flat_function_tool_routes_to_responses(): + """Responses-style flat function tool defs still count as function tools.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "name": "shell", "parameters": {"type": "object"}}], + reasoning_effort=None, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +@pytest.mark.parametrize( + "custom_llm_provider, model_name, api_base", + [ + pytest.param("openai", "gpt-5.6", None, id="openai"), + pytest.param("azure_ai", "gpt-6-astra", "https://myproject.services.ai.azure.com", id="azure-ai-foundry"), + ], +) +def test_responses_api_bridge_check_function_tool_without_body_stays_chat( + monkeypatch, custom_llm_provider, model_name, api_base +): + import litellm + from litellm.main import responses_api_bridge_check + + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", None) + + model_info, model = responses_api_bridge_check( + model=model_name, + custom_llm_provider=custom_llm_provider, + tools=[{"type": "function"}], + reasoning_effort=None, + api_base=api_base, + ) + + assert model == model_name + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_dict_effort_none_stays_chat(): + """The escape hatch must honor litellm's dict form: {"effort": "none"} means reasoning off.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort={"effort": "none"}, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_dict_effort_active_routes_to_responses(): + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort={"effort": "low"}, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_dict_effort_none_with_summary_routes_to_responses(): + """A summary inside the dict form is Responses-only even when effort is none.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort={"effort": "none", "summary": "concise"}, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +@pytest.mark.parametrize("blank_api_base", [None, "", " ", "\t"]) +def test_responses_api_bridge_check_blank_api_base_is_default_openai(blank_api_base): + """ + A blank api_base (None, empty, or whitespace) resolves to the default OpenAI + endpoint downstream, which enforces the reasoning+tools constraint, so gpt-5.4+ + function-tool requests with unset reasoning_effort must still auto-bridge. + """ + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=blank_api_base, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_custom_api_base_with_unset_effort_stays_chat(): + """ + Chat-only OpenAI-compatible backends registered under the openai provider with a + custom api_base and gpt-5.4+ model names serve tools-without-reasoning fine and + have no /responses route; the unset-effort arm must not reroute them. + """ + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base="http://vllm.internal:8000/v1", + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_custom_api_base_via_global_with_unset_effort_stays_chat(monkeypatch): + """ + A custom base set through the litellm.api_base global (not the call arg) is resolved the + same way the chat handler resolves it, so the unset-effort arm must not reroute a chat-only + backend to a /responses route it lacks. Regression guard: the gate previously inspected only + the call-level api_base and bridged these requests. + """ + import litellm + from litellm.main import responses_api_bridge_check + + monkeypatch.setattr(litellm, "api_base", "http://vllm.internal:8000/v1") + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=None, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") != "responses" + + +@pytest.mark.parametrize("env_var", ["OPENAI_BASE_URL", "OPENAI_API_BASE"]) +def test_responses_api_bridge_check_custom_api_base_via_env_with_unset_effort_stays_chat(monkeypatch, env_var): + """ + A custom base set via OPENAI_BASE_URL/OPENAI_API_BASE env is resolved identically to the chat + handler, so the unset-effort arm leaves the request on chat instead of bridging it. + """ + import litellm + from litellm.main import responses_api_bridge_check + + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setenv(env_var, "http://vllm.internal:8000/v1") + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=None, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") != "responses" + + +@pytest.mark.parametrize( + "api_base", + [ + "https://southcentralus.privatelink.api.openai.com/v1", + "https://privatelink.corp.api.openai.com/v1", + "https://api.openai.com:443/v1", + "https://api.openai.com/v1/", + "HTTPS://API.OPENAI.COM/v1", + ], +) +def test_responses_api_bridge_check_openai_backed_custom_api_base_with_unset_effort_routes_to_responses(api_base): + """ + A custom api_base whose host is api.openai.com or a subdomain of it (a PrivateLink hostname, a + port-qualified or trailing-slash default) still reaches the real OpenAI backend, which rejects + function tools with reasoning on Chat Completions, so the unset-effort arm must bridge exactly as + it does for the literal default URL. Regression guard for GH #39353. + """ + from litellm.main import responses_api_bridge_check + + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=api_base, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +@pytest.mark.parametrize( + "api_base", + [ + "https://api.openai.com.evil.example/v1", + "https://notapi.openai.com/v1", + "https://gateway.example/v1?upstream=api.openai.com", + "https://openai.internal.example/api.openai.com/v1", + ], +) +def test_responses_api_bridge_check_lookalike_custom_api_base_with_unset_effort_stays_chat(api_base): + """Only the host decides: api.openai.com appearing elsewhere in the URL is still a foreign backend.""" + from litellm.main import responses_api_bridge_check + + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=api_base, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_privatelink_api_base_via_env_with_unset_effort_routes_to_responses(monkeypatch): + """A PrivateLink base set through OPENAI_BASE_URL resolves the way the chat handler's does and still bridges.""" + import litellm + from litellm.main import responses_api_bridge_check + + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setenv("OPENAI_BASE_URL", "https://southcentralus.privatelink.api.openai.com/v1") + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base=None, + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_custom_api_base_with_explicit_effort_still_routes(): + """Explicit reasoning_effort keeps its pre-existing bridging behavior on any api_base.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.6", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort="high", + api_base="http://vllm.internal:8000/v1", + ) + + assert model == "gpt-5.6" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_azure_with_api_base_and_unset_effort_routes(): + """Azure OpenAI always sets api_base and does enforce the constraint; keep bridging.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="azure", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + api_base="https://myresource.openai.azure.com", + ) + + assert model == "gpt-5.4" + assert model_info.get("mode") == "responses" + + +_FOUNDRY_API_BASE: Final = "https://myproject.services.ai.azure.com" +_FOUNDRY_FUNCTION_TOOL: Final = ({"type": "function", "function": {"name": "get_weather"}},) + + +@pytest.mark.parametrize( + "model_name, api_base, reasoning_effort", + [ + pytest.param("gpt-6-astra", _FOUNDRY_API_BASE, None, id="gpt-6-unset-effort"), + pytest.param("gpt-6-astra", _FOUNDRY_API_BASE, "low", id="gpt-6-explicit-effort"), + pytest.param("gpt-6-astra", "https://myresource.openai.azure.com", None, id="gpt-6-azure-openai-host"), + pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, "low", id="gpt-5.6-explicit-effort"), + pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, {"effort": "high"}, id="gpt-5.6-explicit-effort-dict"), + ], +) +def test_responses_api_bridge_check_azure_ai_foundry_rejected_tools_route_to_responses( + model_name, api_base, reasoning_effort +): + from litellm.main import responses_api_bridge_check + + model_info, model = responses_api_bridge_check( + model=model_name, + custom_llm_provider="azure_ai", + tools=_FOUNDRY_FUNCTION_TOOL, + reasoning_effort=reasoning_effort, + api_base=api_base, + ) + + assert model == model_name + assert model_info.get("mode") == "responses" + + +@pytest.mark.parametrize( + "model_name, api_base, reasoning_effort", + [ + pytest.param("gpt-6-astra", _FOUNDRY_API_BASE, "none", id="explicit-none-stays-chat"), + pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, None, id="gpt-5.6-unset-effort-stays-chat"), + pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, "none", id="gpt-5.6-explicit-none-stays-chat"), + pytest.param("gpt-5.5", _FOUNDRY_API_BASE, "high", id="gpt-5.5-explicit-effort-stays-chat"), + pytest.param("gpt-5.4-mini", _FOUNDRY_API_BASE, None, id="gpt-5.4-mini-unset-effort-stays-chat"), + pytest.param("gpt-5.4-mini", _FOUNDRY_API_BASE, "low", id="gpt-5.4-mini-explicit-effort-stays-chat"), + pytest.param("gpt-6-astra", "https://myproject.models.ai.azure.com", None, id="serverless-host-stays-chat"), + pytest.param("Mistral-large-2411", _FOUNDRY_API_BASE, None, id="non-gpt-5-model-stays-chat"), + pytest.param("claude-opus-4-1", _FOUNDRY_API_BASE, None, id="claude-on-foundry-stays-chat"), + ], +) +def test_responses_api_bridge_check_azure_ai_without_foundry_responses_route_stays_chat( + model_name, api_base, reasoning_effort +): + from litellm.main import responses_api_bridge_check + + model_info, model = responses_api_bridge_check( + model=model_name, + custom_llm_provider="azure_ai", + tools=_FOUNDRY_FUNCTION_TOOL, + reasoning_effort=reasoning_effort, + api_base=api_base, + ) + + assert model == model_name + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_older_gpt_5_tools_without_reasoning_stays_chat(): + """Pre-5.4 GPT-5 names keep the old boundary: tools alone never bridge.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.1", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort=None, + ) + + assert model == "gpt-5.1" + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_gpt_5_4_reasoning_summary_without_tools_routes_to_responses(): + """gpt-5.4+ with reasoning_effort + reasoningSummary but no tools should bridge (AI SDK).""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="openai", + tools=None, + reasoning_effort="medium", + reasoning_summary="auto", + ) + + assert model == "gpt-5.4" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_gpt_5_reasoning_summary_routes_to_responses(): + """Bare ``gpt-5`` with reasoning_effort + reasoningSummary should bridge (not 5.4+).""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5", + custom_llm_provider="openai", + tools=None, + reasoning_effort="medium", + reasoning_summary="auto", + ) + + assert model == "gpt-5" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_gpt_5_tools_without_summary_stays_chat(): + """gpt-5 with tools + reasoning_effort but no summary should stay on chat.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-5", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + reasoning_effort="medium", + reasoning_summary=None, + ) + + assert model == "gpt-5" + assert model_info.get("mode") != "responses" + + +@patch("litellm.completion_extras.responses_api_bridge.completion") +def test_gpt_5_4_responses_bridge_preserves_reasoning_summary_dict( + mock_responses_completion, +): + """When routed to Responses, preserve reasoning_effort summary dict.""" + mock_responses_completion.return_value = MagicMock() + + import litellm + + litellm.completion( + model="gpt-5.4", + messages=[{"role": "user", "content": "What is the capital of France?"}], + tools=[ + { + "type": "function", + "function": { + "name": "get_capital", + "description": "Get the capital of a country", + "parameters": { + "type": "object", + "properties": {"country": {"type": "string"}}, + }, + }, + } + ], + reasoning_effort={"effort": "xhigh", "summary": "detailed"}, + api_key="fake-key", + ) + + assert mock_responses_completion.called is True + optional_params = mock_responses_completion.call_args.kwargs["optional_params"] + assert optional_params["reasoning_effort"] == { + "effort": "xhigh", + "summary": "detailed", + } + + +@pytest.mark.parametrize("reasoning_effort", ["high", {"effort": "high"}]) +def test_responses_bridge_preserves_reasoning_effort_with_drop_params( + reasoning_effort, + restore_model_registry, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + response_body: Final = { + "id": "resp_test", + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "test-responses-bridge", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": None, + "temperature": None, + "tool_choice": "auto", + "tools": [], + "top_p": None, + "max_output_tokens": None, + "previous_response_id": None, + "reasoning": None, + "truncation": None, + "user": None, + } + response_route: Final = respx_mock.post("https://api.perplexity.ai/v1/responses").respond(json=response_body) + model: Final = "perplexity/test-responses-bridge" + litellm.register_model( + { + model: { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_reasoning": False, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + }, + persist_across_reloads=False, + ) + + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=reasoning_effort, + drop_params=True, + api_key="fake-key", + api_base="https://api.perplexity.ai", + ) + + request_body: Final = json.loads(response_route.calls[0].request.content) + assert request_body["reasoning"] == {"effort": "high"} + + +_FOUNDRY_RESPONSES_FUNCTION_CALL_BODY: Final = { + "id": "resp_foundry", + "object": "response", + "created_at": 1789852145, + "status": "completed", + "model": "gpt-6-astra", + "output": [ + { + "id": "fc_1", + "type": "function_call", + "status": "completed", + "arguments": '{"city":"Paris"}', + "call_id": "call_1", + "name": "get_weather", + } + ], + "parallel_tool_calls": True, + "usage": { + "input_tokens": 53, + "output_tokens": 18, + "total_tokens": 71, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": {}, + "temperature": 1.0, + "tool_choice": "auto", + "tools": [], + "top_p": 1.0, + "max_output_tokens": 200, + "previous_response_id": None, + "reasoning": {"effort": "medium", "summary": None}, + "truncation": "disabled", + "user": None, +} + + +def test_completion_bridges_azure_ai_foundry_gpt_5_4_plus_function_tools_to_responses( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + responses_route: Final = respx_mock.post(f"{_FOUNDRY_API_BASE}/openai/v1/responses").respond( + json=_FOUNDRY_RESPONSES_FUNCTION_CALL_BODY + ) + + response: Final = litellm.completion( + model="azure_ai/gpt-6-astra", + messages=[{"role": "user", "content": "What is the weather in Paris? Use the tool."}], + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, + } + ], + max_tokens=200, + api_base=_FOUNDRY_API_BASE, + api_key="fake-foundry-key", + ) + + assert [str(call.request.url) for call in respx_mock.calls] == [f"{_FOUNDRY_API_BASE}/openai/v1/responses"] + request: Final = responses_route.calls[0].request + request_body: Final = json.loads(request.content) + assert request_body["tools"][0]["type"] == "function" + assert request_body["tools"][0]["name"] == "get_weather" + assert request.headers["api-key"] == "fake-foundry-key" + assert response.choices[0].finish_reason == "tool_calls" + assert response.choices[0].message.tool_calls[0].function.name == "get_weather" + + +@pytest.mark.parametrize( + "model, model_info, expected_model_param, expected_base_model_param", + [ + ("gemini/gemini-3.1-pro", None, "gemini-3.1-pro", None), + ( + "gemini/gemini-3.1-pro", + {"base_model": "gemini-3.1-pro-preview"}, + "gemini-3.1-pro", + "gemini-3.1-pro-preview", + ), + ], +) +def test_completion_optional_params_base_model( + model: str, + model_info: dict | None, + expected_model_param: str, + expected_base_model_param: str | None, +): + """``model_info.base_model`` must reach ``get_optional_params`` as ``base_model`` + (an additive capability hint), without overwriting ``model`` with the label. + + Regression for #29618: overwriting ``model`` with a friendly ``base_model`` + label made Bedrock drop ``tools``/``tool_choice`` under ``drop_params``.""" + with patch("litellm.main.get_optional_params") as mock_get_optional_params: + mock_get_optional_params.return_value = MagicMock() + + import litellm + + kwargs = { + "model": model, + "messages": [{"role": "user", "content": "What is the capital of France?"}], + "api_key": "fake-key", + "mock_response": "Hey, how's it going?", + } + if model_info is not None: + kwargs["model_info"] = model_info + + litellm.completion(**kwargs) + + assert mock_get_optional_params.called is True + call_kwargs = mock_get_optional_params.call_args.kwargs + assert call_kwargs["model"] == expected_model_param + assert call_kwargs["base_model"] == expected_base_model_param + + +@patch("litellm.completion_extras.responses_api_bridge.completion") +def test_gpt_5_4_responses_bridge_merges_reasoning_summary_kwarg_without_tools( + mock_responses_completion, +): + """reasoningSummary without tools should route and merge into reasoning_effort dict.""" + mock_responses_completion.return_value = MagicMock() + + import litellm + + litellm.completion( + model="gpt-5.4", + messages=[{"role": "user", "content": "ok"}], + reasoning_effort="medium", + reasoningSummary="auto", + api_key="fake-key", + ) + + assert mock_responses_completion.called is True + optional_params = mock_responses_completion.call_args.kwargs["optional_params"] + assert optional_params["reasoning_effort"] == { + "effort": "medium", + "summary": "auto", + } + assert "reasoningSummary" not in optional_params + assert "reasoning_summary" not in optional_params + + +@patch("litellm.completion_extras.responses_api_bridge.completion") +def test_responses_bridge_preserves_reasoning_summary_without_effort( + mock_responses_completion, +): + """Reasoning summary should survive responses routing even without effort.""" + mock_responses_completion.return_value = MagicMock() + + import litellm + + with patch.object(litellm, "route_all_chat_openai_to_responses", True): + litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "ok"}], + reasoningSummary="auto", + api_key="fake-key", + ) + + assert mock_responses_completion.called is True + optional_params = mock_responses_completion.call_args.kwargs["optional_params"] + assert optional_params["reasoning_effort"] == {"summary": "auto"} + assert "reasoningSummary" not in optional_params + assert "reasoning_summary" not in optional_params + + +@patch("litellm.completion_extras.responses_api_bridge.completion") +def test_gpt_5_responses_bridge_tools_and_reasoning_summary( + mock_responses_completion, +): + """Bare gpt-5 with tools + reasoningSummary should bridge (OpenCode-style).""" + mock_responses_completion.return_value = MagicMock() + + import litellm + + litellm.completion( + model="gpt-5", + messages=[{"role": "user", "content": "ok"}], + tools=[ + { + "type": "function", + "function": { + "name": "apply_patch", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + tool_choice="auto", + reasoning_effort="medium", + reasoningSummary="auto", + stream=True, + api_key="fake-key", + ) + + assert mock_responses_completion.called is True + optional_params = mock_responses_completion.call_args.kwargs["optional_params"] + assert optional_params.get("reasoning_effort") == { + "effort": "medium", + "summary": "auto", + } + + +def test_responses_api_bridge_check_handles_exception(): + """Test that responses_api_bridge_check handles exceptions and still processes responses/ models.""" + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.side_effect = Exception("Model not found") + + model_info, model = responses_api_bridge_check( + model="responses/custom-model", custom_llm_provider="custom" + ) + + assert model == "custom-model" + assert model_info["mode"] == "responses" + + +def test_responses_api_bridge_check_global_flag_routes_openai(): + """When route_all_chat_openai_to_responses is True, any OpenAI model routes to responses.""" + from litellm.main import responses_api_bridge_check + + with patch.object(litellm, "route_all_chat_openai_to_responses", True): + model_info, model = responses_api_bridge_check( + model="gpt-4o", + custom_llm_provider="openai", + ) + + assert model == "gpt-4o" + assert model_info.get("mode") == "responses" + + +def test_responses_api_bridge_check_global_flag_does_not_affect_azure(): + """route_all_chat_openai_to_responses should not affect Azure models.""" + from litellm.main import responses_api_bridge_check + + with patch.object(litellm, "route_all_chat_openai_to_responses", True): + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 4096} + model_info, model = responses_api_bridge_check( + model="gpt-4o", + custom_llm_provider="azure", + ) + + assert model_info.get("mode") != "responses" + + +def test_responses_api_bridge_check_global_flag_default_false(): + """By default, route_all_chat_openai_to_responses is False and doesn't affect routing.""" + from litellm.main import responses_api_bridge_check + + with patch.object(litellm, "route_all_chat_openai_to_responses", False): + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 4096} + model_info, model = responses_api_bridge_check( + model="gpt-4o", + custom_llm_provider="openai", + ) + + assert model_info.get("mode") != "responses" + + +@pytest.mark.asyncio +async def test_async_mock_delay(): + """Use asyncio await for mock delay on acompletion""" + import time + + from litellm import acompletion + + start_time = time.time() + result = await acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hey, how's it going?"}], + mock_delay=0.01, + mock_response="Hello world", + ) + end_time = time.time() + delay = end_time - start_time + assert delay >= 0.01 + + +def test_stream_chunk_builder_keeps_tool_calls_carried_only_by_a_later_choice_of_a_multi_choice_chunk(): + from litellm import stream_chunk_builder + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + def chunk(choices: list[StreamingChoices]) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-multi-choice", + created=1751934860, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=choices, + ) + + chunks = [ + chunk( + [ + StreamingChoices(index=0, delta=Delta(role="assistant", content="hello")), + StreamingChoices( + index=1, + delta=Delta( + role="assistant", + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call_1", + index=0, + type="function", + function=Function(name="lookup_fruit", arguments='{"fruit":'), + ) + ], + ), + ), + ] + ), + chunk( + [ + StreamingChoices(index=0, delta=Delta(content=" world"), finish_reason="stop"), + StreamingChoices( + index=1, + delta=Delta( + tool_calls=[ChatCompletionDeltaToolCall(index=0, function=Function(arguments='"kiwi"}'))] + ), + finish_reason="tool_calls", + ), + ] + ), + ] + + response = stream_chunk_builder(chunks=chunks) + + tool_calls = response.choices[0].message.tool_calls + assert tool_calls is not None + assert [(call.id, call.function.name, call.function.arguments) for call in tool_calls] == [ + ("call_1", "lookup_fruit", '{"fruit":"kiwi"}') + ] + + +def test_stream_chunk_builder_thinking_blocks(): + from litellm import stream_chunk_builder + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + chunks = [ + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_content="I need to summar", + thinking_blocks=[ + { + "type": "thinking", + "thinking": "I need to summar", + "signature": None, + } + ], + provider_specific_fields={ + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "I need to summar", + "signature": None, + } + ] + }, + content="", + role="assistant", + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_content="ize the previous agent's thinking process into a", + thinking_blocks=[ + { + "type": "thinking", + "thinking": "ize the previous agent's thinking process into a", + "signature": None, + } + ], + provider_specific_fields={ + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "ize the previous agent's thinking process into a", + "signature": None, + } + ] + }, + content="", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_content=" short description. Based on the input data provide", + thinking_blocks=[ + { + "type": "thinking", + "thinking": " short description. Based on the input data provide", + "signature": None, + } + ], + provider_specific_fields={ + "thinking_blocks": [ + { + "type": "thinking", + "thinking": " short description. Based on the input data provide", + "signature": None, + } + ] + }, + content="", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_content="d, it seems the agent was planning to refine their search", + thinking_blocks=[ + { + "type": "thinking", + "thinking": "d, it seems the agent was planning to refine their search", + "signature": None, + } + ], + provider_specific_fields={ + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "d, it seems the agent was planning to refine their search", + "signature": None, + } + ] + }, + content="", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_content=" to focus more on technical aspects of home automation and home", + thinking_blocks=[ + { + "type": "thinking", + "thinking": " to focus more on technical aspects of home automation and home", + "signature": None, + } + ], + provider_specific_fields={ + "thinking_blocks": [ + { + "type": "thinking", + "thinking": " to focus more on technical aspects of home automation and home", + "signature": None, + } + ] + }, + content="", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_content=" energy system management.\n\nI'll create a brief", + thinking_blocks=[ + { + "type": "thinking", + "thinking": " energy system management.\n\nI'll create a brief", + "signature": None, + } + ], + provider_specific_fields={ + "thinking_blocks": [ + { + "type": "thinking", + "thinking": " energy system management.\n\nI'll create a brief", + "signature": None, + } + ] + }, + content="", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_content=" summary of what the agent was doing.", + thinking_blocks=[ + { + "type": "thinking", + "thinking": " summary of what the agent was doing.", + "signature": None, + } + ], + provider_specific_fields={ + "thinking_blocks": [ + { + "type": "thinking", + "thinking": " summary of what the agent was doing.", + "signature": None, + } + ] + }, + content="", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + reasoning_content="", + thinking_blocks=[ + { + "type": "thinking", + "thinking": "", + "signature": "ErUBCkYIBRgCIkAKBSMkB2+MBF643wiWxlERsGXVdlhbPx9lnTIbygzjFIeZ5uhTV+HNWDon9vQV4hmXvAKwQfwS8vkNFB366l05Egzt2U18IpRrZRyQn1UaDDdYvKHYP8Ps1IbWjSIw8eSYOU9gtqNcwR6D0wY7iOPx2GliDEatLI5rSs96CByoTIoADL2M5bX8KP0jEpbHKh0ccYryigdH/3J8EiFt/BmGUceVASP5l9r22dFWiBgC", + } + ], + provider_specific_fields={ + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "", + "signature": "ErUBCkYIBRgCIkAKBSMkB2+MBF643wiWxlERsGXVdlhbPx9lnTIbygzjFIeZ5uhTV+HNWDon9vQV4hmXvAKwQfwS8vkNFB366l05Egzt2U18IpRrZRyQn1UaDDdYvKHYP8Ps1IbWjSIw8eSYOU9gtqNcwR6D0wY7iOPx2GliDEatLI5rSs96CByoTIoADL2M5bX8KP0jEpbHKh0ccYryigdH/3J8EiFt/BmGUceVASP5l9r22dFWiBgC", + } + ] + }, + content="", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=1, + delta=Delta( + provider_specific_fields=None, + content='{"a', + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=1, + delta=Delta( + provider_specific_fields=None, + content='gent_doing"', + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=1, + delta=Delta( + provider_specific_fields=None, + content=': "Re', + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=1, + delta=Delta( + provider_specific_fields=None, + content="searching", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=1, + delta=Delta( + provider_specific_fields=None, + content=" technic", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=1, + delta=Delta( + provider_specific_fields=None, + content="al aspect", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=1, + delta=Delta( + provider_specific_fields=None, + content="s of home au", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=1, + delta=Delta( + provider_specific_fields=None, + content='tomation"}', + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + citations=None, + ), + ModelResponseStream( + id="chatcmpl-e8febeb7-cf7d-4947-9417-59ae5e6989f9", + created=1751934860, + model="claude-3-7-sonnet-latest", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason="tool_calls", + index=0, + delta=Delta( + provider_specific_fields=None, + content=None, + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + ), + ] + + response = stream_chunk_builder(chunks=chunks) + print(response) + + assert response is not None + assert response.choices[0].message.content is not None + assert response.choices[0].message.thinking_blocks is not None + + +from litellm.llms.openai.openai import OpenAIChatCompletion + + +def throw_retryable_error(*_, **__): + raise RuntimeError("BOOM") + + +@pytest.mark.asyncio +async def test_retrying() -> None: + litellm.num_retries = 10 + with ( + patch.object( + OpenAIChatCompletion, + "make_openai_chat_completion_request", + side_effect=throw_retryable_error, + ) as mock_request, + pytest.raises(litellm.InternalServerError, match="LiteLLM Retried: 10 times"), + ): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Hello"}], + ) + + +def test_anthropic_disable_url_suffix_env_var(): + """Test that LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX prevents /v1/messages suffix.""" + import os + from unittest.mock import MagicMock, patch + + from litellm import completion + + # Test with environment variable disabled (default behavior) + with patch.dict(os.environ, {"ANTHROPIC_API_BASE": "https://api.example.com"}): + actual_api_base = None + + with patch("litellm.main.anthropic_chat_completions") as mock_anthropic: + + def capture_completion(**kwargs): + nonlocal actual_api_base + actual_api_base = kwargs.get("api_base") + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + return mock_response + + mock_anthropic.completion = capture_completion + + # This should append /v1/messages + completion( + model="anthropic/claude-3-sonnet", + messages=[{"role": "user", "content": "test"}], + api_key="test-key", + ) + + # Verify the api_base has /v1/messages appended + assert actual_api_base.endswith("/v1/messages") + assert actual_api_base == "https://api.example.com/v1/messages" + + # Test with environment variable enabled + with patch.dict( + os.environ, + { + "ANTHROPIC_API_BASE": "https://api.example.com/custom/path", + "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX": "true", + }, + ): + actual_api_base = None + + with patch("litellm.main.anthropic_chat_completions") as mock_anthropic: + + def capture_completion(**kwargs): + nonlocal actual_api_base + actual_api_base = kwargs.get("api_base") + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + return mock_response + + mock_anthropic.completion = capture_completion + + # This should NOT append /v1/messages + completion( + model="anthropic/claude-3-sonnet", + messages=[{"role": "user", "content": "test"}], + api_key="test-key", + ) + + # Verify the api_base does not have /v1/messages appended + assert actual_api_base == "https://api.example.com/custom/path" + assert not actual_api_base.endswith("/v1/messages") + + +def test_anthropic_text_disable_url_suffix_env_var(): + """Test that LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX prevents /v1/complete suffix for anthropic_text.""" + import os + from unittest.mock import MagicMock, patch + + from litellm import completion + + # Test with environment variable disabled (default behavior) + with patch.dict(os.environ, {"ANTHROPIC_API_BASE": "https://api.example.com"}): + actual_api_base = None + + with patch("litellm.main.base_llm_http_handler") as mock_handler: + + def capture_completion(**kwargs): + nonlocal actual_api_base + actual_api_base = kwargs.get("api_base") + return MagicMock() + + mock_handler.completion = capture_completion + + # This should append /v1/complete + completion( + model="anthropic_text/claude-instant-1", + messages=[{"role": "user", "content": "test"}], + api_key="test-key", + ) + + # Verify the api_base has /v1/complete appended + assert actual_api_base.endswith("/v1/complete") + assert actual_api_base == "https://api.example.com/v1/complete" + + # Test with environment variable enabled + with patch.dict( + os.environ, + { + "ANTHROPIC_API_BASE": "https://api.example.com/custom/complete", + "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX": "true", + }, + ): + actual_api_base = None + + with patch("litellm.main.base_llm_http_handler") as mock_handler: + + def capture_completion(**kwargs): + nonlocal actual_api_base + actual_api_base = kwargs.get("api_base") + return MagicMock() + + mock_handler.completion = capture_completion + + # This should NOT append /v1/complete + completion( + model="anthropic_text/claude-instant-1", + messages=[{"role": "user", "content": "test"}], + api_key="test-key", + ) + + # Verify the api_base does not have /v1/complete appended + assert actual_api_base == "https://api.example.com/custom/complete" + assert not actual_api_base.endswith("/v1/complete") + + +def test_image_edit_merges_headers_and_extra_headers(): + from litellm.images.main import base_llm_http_handler + + combined_headers = { + "x-test-header-one": "value-1", + "x-test-header-two": "value-2", + } + + mock_image_edit_config = MagicMock() + mock_image_edit_config.get_supported_openai_params.return_value = set() + mock_image_edit_config.map_openai_params.side_effect = lambda **kwargs: dict( + kwargs["image_edit_optional_params"] + ) + + with ( + patch( + "litellm.images.main.ProviderConfigManager.get_provider_image_edit_config", + return_value=mock_image_edit_config, + ) as mock_config, + patch.object( + base_llm_http_handler, + "image_edit_handler", + return_value="ok", + ) as mock_handler, + ): + response = litellm.image_edit( + image=MagicMock(name="image"), + prompt="test", + model="azure/gpt-image-1", + headers={"x-test-header-one": "value-1"}, + extra_headers={ + "x-test-header-two": "value-2", + }, + ) + + assert response == "ok" + mock_config.assert_called_once() + + handler_kwargs = mock_handler.call_args.kwargs + assert handler_kwargs["extra_headers"] == combined_headers + assert "extra_headers" not in handler_kwargs["image_edit_optional_request_params"] + + +@pytest.mark.parametrize("metadata_key", ("metadata", "litellm_metadata")) +@pytest.mark.parametrize("input_tokens", (51234, 0)) +def test_mock_completion_usage_reports_admission_input_tokens(metadata_key: str, input_tokens: int): + response = litellm.completion( + model="anthropic/claude-sonnet-5", + messages=[{"role": "user", "content": "hello"}], + mock_response="ok", + api_key="mock", + **{metadata_key: {"user_api_key_budget_reservation": {"reserved_cost": 1.0, "input_tokens": input_tokens}}}, + ) + + assert response.usage.prompt_tokens == input_tokens + assert response.usage.total_tokens == input_tokens + response.usage.completion_tokens + + +def test_mock_completion_usage_falls_back_to_default_without_admission_count(): + response = litellm.completion( + model="anthropic/claude-sonnet-5", + messages=[{"role": "user", "content": "hello"}], + mock_response="ok", + api_key="mock", + metadata={"user_api_key_budget_reservation": {"reserved_cost": 1.0}}, + ) + + assert response.usage.prompt_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT + + +_AZURE_AI_CUSTOM_PRICED_DEPLOYMENT: Final = { + "model_name": "azure-ai-custom-priced", + "litellm_params": { + "model": "azure_ai/gpt-5.6", + "api_key": "mock", + "api_base": "https://example.services.ai.azure.com", + "mock_response": "ok", + "input_cost_per_token": 3e-6, + "output_cost_per_token": 7e-6, + "cache_read_input_token_cost": 1e-7, + "cache_creation_input_token_cost": 5e-7, + }, + "model_info": {"id": "azure-ai-custom-priced-deployment-id"}, +} + + +def _expected_custom_price(response: litellm.ModelResponse) -> float: + params: Final = _AZURE_AI_CUSTOM_PRICED_DEPLOYMENT["litellm_params"] + return ( + response.usage.prompt_tokens * params["input_cost_per_token"] + + response.usage.completion_tokens * params["output_cost_per_token"] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", (False, True)) +async def test_mock_completion_prices_azure_ai_router_deployment_with_custom_pricing(use_async: bool): + router: Final = litellm.Router(model_list=[_AZURE_AI_CUSTOM_PRICED_DEPLOYMENT]) + messages: Final = [{"role": "user", "content": "hello"}] + + response: Final = ( + await router.acompletion(model="azure-ai-custom-priced", messages=messages) + if use_async + else router.completion(model="azure-ai-custom-priced", messages=messages) + ) + + assert response._hidden_params["response_cost"] == pytest.approx(_expected_custom_price(response)) + assert response._hidden_params["custom_llm_provider"] == "azure_ai" + + +@pytest.mark.parametrize( + ("model", "expected_provider"), + (("anthropic/claude-sonnet-5", "anthropic"), ("no-such-provider-model", None)), +) +def test_mock_completion_infers_provider_when_called_directly_without_one(model: str, expected_provider: str | None): + response: Final = litellm.mock_completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + mock_response="ok", + ) + + assert response.choices[0].message.content == "ok" + assert response._hidden_params.get("custom_llm_provider") == expected_provider + + +_ADMISSION_INPUT_TOKENS: Final = 51234 + + +def _admission_metadata(input_tokens: int) -> dict[str, object]: # mutable-ok: logging writes into metadata + return {"user_api_key_budget_reservation": {"reserved_cost": 1.0, "input_tokens": input_tokens}} + + +_ADMISSION_METADATA: Final = _admission_metadata(_ADMISSION_INPUT_TOKENS) +_MOCK_STREAM_MESSAGES: Final = [{"role": "user", "content": "hello " * 200}] +_STREAM_CHUNK_BUILDER_TOKEN_COUNTER: Final = "litellm.litellm_core_utils.streaming_chunk_builder_utils.token_counter" + + +def _prompt_token_counter_calls(token_counter: MagicMock) -> list[object]: + return [call for call in token_counter.call_args_list if call.kwargs.get("messages") is not None] + + +def _client_usage_chunks(chunks: list[ModelResponseStream]) -> list[Usage]: + return [chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None] + + +@pytest.mark.parametrize("n", (None, 2)) +def test_mock_completion_stream_usage_reports_admission_input_tokens_without_tokenizer_fallback(n: int | None): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + n=n, + stream_options={"include_usage": True}, + metadata=_ADMISSION_METADATA, + ) + ) + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == _ADMISSION_INPUT_TOKENS + assert usage_chunks[0].completion_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT + assert usage_chunks[0].total_tokens == _ADMISSION_INPUT_TOKENS + usage_chunks[0].completion_tokens + assert _prompt_token_counter_calls(token_counter) == [] + assert all(chunk.choices for chunk in chunks[:-1]) + assert {chunk.id for chunk in chunks} == {chunks[0].id} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("n", (None, 2)) +async def test_mock_acompletion_stream_usage_reports_admission_input_tokens_without_tokenizer_fallback( + n: int | None, +): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response: Final = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + n=n, + stream_options={"include_usage": True}, + litellm_metadata=_ADMISSION_METADATA, + ) + chunks: Final = [chunk async for chunk in response] + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == _ADMISSION_INPUT_TOKENS + assert usage_chunks[0].total_tokens == _ADMISSION_INPUT_TOKENS + usage_chunks[0].completion_tokens + assert _prompt_token_counter_calls(token_counter) == [] + assert all(chunk.choices for chunk in chunks[:-1]) + assert {chunk.id for chunk in chunks} == {chunks[0].id} + + +def test_mock_completion_stream_without_include_usage_hides_usage_chunk_but_logs_admission_count(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + metadata=_ADMISSION_METADATA, + ) + ) + + assert _client_usage_chunks(chunks) == [] + assert all(len(chunk.choices) == 1 for chunk in chunks) + assert chunks[-1]._hidden_params["usage"].prompt_tokens == _ADMISSION_INPUT_TOKENS + assert _prompt_token_counter_calls(token_counter) == [] + + +def test_mock_completion_stream_with_empty_stream_options_completes_and_logs_admission_count(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={}, + metadata=_ADMISSION_METADATA, + ) + ) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok" + assert _client_usage_chunks(chunks) == [] + assert _prompt_token_counter_calls(token_counter) == [] + + +@pytest.mark.asyncio +async def test_mock_acompletion_stream_with_empty_stream_options_completes_and_logs_admission_count(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response: Final = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={}, + litellm_metadata=_ADMISSION_METADATA, + ) + chunks: Final = [chunk async for chunk in response] + + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok" + assert _client_usage_chunks(chunks) == [] + assert _prompt_token_counter_calls(token_counter) == [] + + +def test_mock_completion_stream_without_admission_count_falls_back_to_tokenizer(): + expected_prompt_tokens: Final = litellm.token_counter(model="openai/gpt-5.4-mini", messages=_MOCK_STREAM_MESSAGES) + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + metadata={"user_api_key_budget_reservation": {"reserved_cost": 1.0}}, + ) + ) + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == expected_prompt_tokens + assert usage_chunks[0].total_tokens == expected_prompt_tokens + usage_chunks[0].completion_tokens + assert len(_prompt_token_counter_calls(token_counter)) >= 1 + + +@pytest.mark.asyncio +async def test_mock_acompletion_stream_without_admission_count_falls_back_to_tokenizer(): + expected_prompt_tokens: Final = litellm.token_counter(model="openai/gpt-5.4-mini", messages=_MOCK_STREAM_MESSAGES) + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response: Final = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + ) + chunks: Final = [chunk async for chunk in response] + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == expected_prompt_tokens + assert len(_prompt_token_counter_calls(token_counter)) >= 1 + + +def _usage_triple(usage: Usage) -> tuple[int, int, int]: + return (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) + + +@pytest.mark.parametrize("input_tokens", (_ADMISSION_INPUT_TOKENS, 0)) +def test_mock_completion_stream_and_non_stream_report_the_same_admission_usage(input_tokens: int): + metadata: Final = _admission_metadata(input_tokens) + non_stream: Final = litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + metadata=metadata, + ) + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + metadata=metadata, + ) + ) + + assert _usage_triple(non_stream.usage) == _usage_triple(_client_usage_chunks(chunks)[0]) + assert non_stream.usage.prompt_tokens == input_tokens + assert _prompt_token_counter_calls(token_counter) == [] + + +@pytest.mark.asyncio +async def test_mock_acompletion_stream_reports_zero_admission_input_tokens_without_tokenizer_fallback(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response: Final = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=[{"role": "user", "content": ""}], + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + litellm_metadata=_admission_metadata(0), + ) + chunks: Final = [chunk async for chunk in response] + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert _usage_triple(usage_chunks[0]) == (0, usage_chunks[0].completion_tokens, usage_chunks[0].completion_tokens) + assert _prompt_token_counter_calls(token_counter) == [] + + +def test_mock_text_completion_stream_and_non_stream_report_the_same_zero_admission_usage(): + metadata: Final = _admission_metadata(0) + non_stream: Final = litellm.text_completion( + model="openai/gpt-5.4-mini", prompt="", mock_response="ok", api_key="mock", metadata=metadata + ) + chunks: Final = list( + litellm.text_completion( + model="openai/gpt-5.4-mini", + prompt="", + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + metadata=metadata, + ) + ) + + stream_usages: Final = tuple(chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None) + assert len(stream_usages) == 1 + assert _usage_triple(non_stream.usage) == _usage_triple(stream_usages[0]) + assert non_stream.usage.prompt_tokens == 0 + + +def test_mock_completion_stream_with_model_response(): + """Test that mock_completion correctly handles stream=True with a ModelResponse as mock_response.""" + from litellm import completion + from litellm.types.utils import Choices, Message, ModelResponse, Usage + + # Create a ModelResponse object + mock_model_response = ModelResponse( + id="chatcmpl-test-123", + created=1234567890, + model="gpt-4o-mini", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="This is a test response", + role="assistant", + ), + ) + ], + usage=Usage( + prompt_tokens=10, + completion_tokens=20, + total_tokens=30, + ), + ) + + # Call completion with stream=True and mock_response as ModelResponse + response = completion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + mock_response=mock_model_response, + ) + + # Verify that the response is a stream + assert response is not None + + # Collect all chunks from the stream + chunks = [] + for chunk in response: + chunks.append(chunk) + print(f"Chunk: {chunk}") + + # Verify we got chunks + assert len(chunks) > 0 + + # Verify the content is streamed correctly + accumulated_content = "" + for chunk in chunks: + if ( + hasattr(chunk.choices[0].delta, "content") + and chunk.choices[0].delta.content + ): + accumulated_content += chunk.choices[0].delta.content + + assert "This is a test response" in accumulated_content or len(chunks) > 0 + + +@pytest.mark.asyncio +async def test_async_mock_completion_stream_with_model_response(): + """Test that async mock_completion correctly handles stream=True with a ModelResponse as mock_response.""" + from litellm import acompletion + from litellm.types.utils import Choices, Message, ModelResponse, Usage + + # Create a ModelResponse object + mock_model_response = ModelResponse( + id="chatcmpl-test-456", + created=1234567890, + model="gpt-4o-mini", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="This is an async test response", + role="assistant", + ), + ) + ], + usage=Usage( + prompt_tokens=15, + completion_tokens=25, + total_tokens=40, + ), + ) + + # Call acompletion with stream=True and mock_response as ModelResponse + response = await acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Hello async"}], + stream=True, + mock_response=mock_model_response, + ) + + # Verify that the response is a stream + assert response is not None + + # Collect all chunks from the stream + chunks = [] + async for chunk in response: + chunks.append(chunk) + print(f"Async Chunk: {chunk}") + + # Verify we got chunks + assert len(chunks) > 0 + + # Verify the content is streamed correctly + accumulated_content = "" + for chunk in chunks: + if ( + hasattr(chunk.choices[0].delta, "content") + and chunk.choices[0].delta.content + ): + accumulated_content += chunk.choices[0].delta.content + + assert "This is an async test response" in accumulated_content or len(chunks) > 0 + + +class TestCallTypesOCR: + """Test that OCR call types are properly defined in CallTypes enum. + + Fixes https://github.com/BerriAI/litellm/issues/17381 + """ + + def test_ocr_call_type_exists(self): + """Test that CallTypes.ocr exists and has correct value.""" + from litellm.types.utils import CallTypes + + assert hasattr(CallTypes, "ocr") + assert CallTypes.ocr.value == "ocr" + + def test_aocr_call_type_exists(self): + """Test that CallTypes.aocr exists and has correct value.""" + from litellm.types.utils import CallTypes + + assert hasattr(CallTypes, "aocr") + assert CallTypes.aocr.value == "aocr" + + def test_ocr_call_type_from_string(self): + """Test that CallTypes can be constructed from 'ocr' string.""" + from litellm.types.utils import CallTypes + + call_type = CallTypes("ocr") + assert call_type == CallTypes.ocr + + def test_aocr_call_type_from_string(self): + """Test that CallTypes can be constructed from 'aocr' string. + + This is the actual use case that was failing - the OCR endpoint + uses route_type='aocr' and guardrails try to instantiate + CallTypes('aocr'). + """ + from litellm.types.utils import CallTypes + + call_type = CallTypes("aocr") + assert call_type == CallTypes.aocr + + +def test_stream_chunk_builder_text_completion_combines_text_and_usage(): + from litellm.main import stream_chunk_builder_text_completion + from litellm.types.utils import TextCompletionResponse + + chunks = [ + TextCompletionResponse( + id="cmpl-1", + object="text_completion", + created=1, + model="gpt-3.5-turbo-instruct", + choices=[{"text": "Hello", "index": 0, "logprobs": None, "finish_reason": None}], + ), + TextCompletionResponse( + id="cmpl-1", + object="text_completion", + created=1, + model="gpt-3.5-turbo-instruct", + choices=[{"text": " world", "index": 0, "logprobs": None, "finish_reason": "stop"}], + ), + ] + + response = stream_chunk_builder_text_completion( + chunks=chunks, messages=[{"role": "user", "content": "say hello"}] + ) + + assert response.choices[0].text == "Hello world" + assert response.choices[0].finish_reason == "stop" + assert response.usage.prompt_tokens > 0 + assert response.usage.completion_tokens > 0 + assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens + + +def test_completion_forwards_store_and_prompt_cache_key_to_openai(): + """ + Regression test for https://github.com/BerriAI/litellm/issues/33184 + + store and prompt_cache_key are documented OpenAI chat completion params that + were accepted as supported but silently dropped before the provider request + was built, because they were not named parameters of completion() and + get_optional_params() the way safety_identifier is. + """ + from openai import OpenAI + + client = OpenAI(api_key="fake-api-key") + + with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + try: + litellm.completion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + store=False, + prompt_cache_key="test-cache-key", + client=client, + ) + except Exception as e: + print(e) + + mock_client.assert_called_once() + request_body = mock_client.call_args.kwargs + assert request_body["store"] is False + assert request_body["prompt_cache_key"] == "test-cache-key" + + +@pytest.mark.asyncio +async def test_acompletion_forwards_store_and_prompt_cache_key_to_openai(): + """ + Async variant of the store/prompt_cache_key forwarding regression test for + https://github.com/BerriAI/litellm/issues/33184 + """ + from openai import AsyncOpenAI + + client = AsyncOpenAI(api_key="fake-api-key") + + with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + try: + await litellm.acompletion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + store=False, + prompt_cache_key="test-cache-key", + client=client, + ) + except Exception as e: + print(e) + + mock_client.assert_called_once() + request_body = mock_client.call_args.kwargs + assert request_body["store"] is False + assert request_body["prompt_cache_key"] == "test-cache-key" + + +def test_completion_omits_store_and_prompt_cache_key_when_not_passed(): + """ + When store and prompt_cache_key are not passed, they must not appear in the + outbound request body (guards against always forwarding None defaults). + """ + from openai import OpenAI + + client = OpenAI(api_key="fake-api-key") + + with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + try: + litellm.completion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + client=client, + ) + except Exception as e: + print(e) + + mock_client.assert_called_once() + request_body = mock_client.call_args.kwargs + assert "store" not in request_body + assert "prompt_cache_key" not in request_body + + +def test_completion_forwards_store_and_prompt_cache_key_to_mcp_gateway(): + """ + Regression test for the MCP gateway early-return in completion(): store and + prompt_cache_key are named params, so they no longer travel via **kwargs and + must be forwarded explicitly like safety_identifier and service_tier. + """ + with patch.object( + import_module("litellm.responses.mcp.chat_completions_handler"), "acompletion_with_mcp" + ) as mock_mcp: + result = litellm.completion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy"}], + store=False, + prompt_cache_key="test-cache-key", + ) + + result.close() + mock_mcp.assert_called_once() + call_kwargs = mock_mcp.call_args.kwargs + assert call_kwargs["store"] is False + assert call_kwargs["prompt_cache_key"] == "test-cache-key" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "aws_credential_kwargs", + [ + { + "aws_session_name": "litellm-gcp", + "aws_role_name": "arn:aws:iam::123456789012:role/litellm-bedrock-role", + "aws_web_identity_token": "oidc/google/108963886734710037768", + }, + { + "aws_access_key_id": "AKIASTATICKEYFORTEST", + "aws_secret_access_key": "static-secret-key", + "aws_session_token": "static-session-token", + }, + ], + ids=["web_identity", "static_keys"], +) +async def test_acompletion_forwards_aws_credentials_through_responses_bridge( + respx_mock: respx.MockRouter, monkeypatch, aws_credential_kwargs: dict +): + from botocore.credentials import Credentials + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + original_disable_aiohttp = litellm.disable_aiohttp_transport + try: + litellm.disable_aiohttp_transport = True + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + get_credentials_mock = MagicMock(return_value=Credentials("fake-key", "fake-secret")) + monkeypatch.setattr(BaseAWSLLM, "get_credentials", get_credentials_mock) + + respx_mock.post("https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses").respond( + json={ + "id": "resp_123", + "object": "response", + "created_at": 1760144904, + "status": "completed", + "model": "openai.gpt-5.4", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + } + ) + + response = await litellm.acompletion( + model="bedrock_mantle/openai.gpt-5.4", + messages=[{"role": "user", "content": "hi"}], + api_base="https://bedrock-mantle.us-east-2.api.aws/v1", + aws_region_name="us-east-2", + num_retries=0, + **aws_credential_kwargs, + ) + + assert response.choices[0].message.content == "ok" + credential_kwargs = get_credentials_mock.call_args.kwargs + assert credential_kwargs["aws_region_name"] == "us-east-2" + for key, value in aws_credential_kwargs.items(): + assert credential_kwargs[key] == value + authorization = respx_mock.calls.last.request.headers["Authorization"] + assert authorization.startswith("AWS4-HMAC-SHA256") + assert "fake-key" in authorization + finally: + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() + + +_GEMINI_RESPONSE_BODY = { + "candidates": [{"content": {"parts": [{"text": "hello"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 2, "candidatesTokenCount": 1, "totalTokenCount": 3}, +} + + +def _gemini_client_returning_a_reply(): + """An injected HTTP client whose post() answers like generativelanguage does.""" + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + client = HTTPHandler() + request = httpx.Request("POST", "https://generativelanguage.googleapis.com/") + post = MagicMock(return_value=httpx.Response(200, json=_GEMINI_RESPONSE_BODY, request=request)) + return client, post + + +@pytest.fixture +def restore_model_registry(): + """litellm.model_cost and the provider name sets are module-global. + + register_model merges into the existing entry in place, hence the deep copy. + """ + model_cost = copy.deepcopy(litellm.model_cost) + openai_models = set(litellm.open_ai_chat_completion_models) + yield + litellm.model_cost.clear() + litellm.model_cost.update(model_cost) + litellm.open_ai_chat_completion_models.clear() + litellm.open_ai_chat_completion_models.update(openai_models) + + +def test_openai_model_name_does_not_outrank_explicit_provider(): + """`gemini/gpt-4o` goes to Google, not to litellm's OpenAI handler. + + completion() checks `model in litellm.open_ai_chat_completion_models` ahead of + the gemini branch, so the call used to reach the OpenAI handler carrying + VertexGeminiConfig, whose transform_request raises NotImplementedError. + """ + assert "gpt-4o" in litellm.open_ai_chat_completion_models + client, post = _gemini_client_returning_a_reply() + + with patch.object(client, "post", new=post): + response = litellm.completion( + model="gemini/gpt-4o", + messages=[{"role": "user", "content": "hello"}], + api_key="test-api-key", + client=client, + ) + + assert "generativelanguage.googleapis.com" in post.call_args.kwargs["url"] + assert "models/gpt-4o" in post.call_args.kwargs["url"] + assert response.choices[0].message.content == "hello" + + +def test_mislabelled_pricing_entry_does_not_reroute_provider(restore_model_registry): + """register_model is the other way into the same failure. + + An entry claiming litellm_provider "openai" adds its name to + open_ai_chat_completion_models, so one mislabelled price reroutes every later + call to that model in the process. + """ + litellm.register_model( + { + "gemini-2.5-pro": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 4e-06, + } + } + ) + assert "gemini-2.5-pro" in litellm.open_ai_chat_completion_models + client, post = _gemini_client_returning_a_reply() + + with patch.object(client, "post", new=post): + response = litellm.completion( + model="gemini/gemini-2.5-pro", + messages=[{"role": "user", "content": "hello"}], + api_key="test-api-key", + client=client, + ) + + assert "generativelanguage.googleapis.com" in post.call_args.kwargs["url"] + assert response.choices[0].message.content == "hello" + + +def test_openai_model_without_a_provider_still_routes_to_openai(): + from openai import OpenAI + + client = OpenAI(api_key="fake-key") + raw_response = client.chat.completions.with_raw_response + with patch.object(raw_response, "create") as mock_create, contextlib.suppress(Exception): + litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "hello"}], + client=client, + ) + + mock_create.assert_called() + + +def _openai_chat_create_kwargs(client, **completion_kwargs): + with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with contextlib.suppress(Exception): + litellm.completion( + messages=[{"role": "system", "content": "sys"}, {"role": "user", "content": "hi"}], + cache_control_injection_points=[{"location": "message", "role": "system"}], + client=client, + **completion_kwargs, + ) + + mock_client.assert_called_once() + return mock_client.call_args.kwargs + + +@pytest.fixture +def _no_openai_api_base_override(monkeypatch): + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", None) + + +@pytest.mark.usefixtures("_no_openai_api_base_override") +def test_completion_custom_api_base_sends_no_prompt_cache_breakpoint_for_gpt_5_6(): + from openai import OpenAI + + client = OpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") + request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", api_base="http://127.0.0.1:9/v1") + + assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} + assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) + assert "prompt_cache_options" not in json.dumps(request_body) + + +@pytest.mark.usefixtures("_no_openai_api_base_override") +def test_completion_custom_base_url_sends_no_prompt_cache_breakpoint_for_gpt_5_6(): + from openai import OpenAI + + client = OpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") + request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", base_url="http://127.0.0.1:9/v1") + + assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} + assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) + assert "prompt_cache_options" not in json.dumps(request_body) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_no_openai_api_base_override") +async def test_acompletion_custom_base_url_sends_no_prompt_cache_breakpoint_for_gpt_5_6(): + from openai import AsyncOpenAI + + client = AsyncOpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") + with patch.object(client.chat.completions.with_raw_response, "create") as mock_create: + with contextlib.suppress(Exception): + await litellm.acompletion( + model="gpt-5.6", + messages=[{"role": "system", "content": "sys"}, {"role": "user", "content": "hi"}], + cache_control_injection_points=[{"location": "message", "role": "system"}], + client=client, + base_url="http://127.0.0.1:9/v1", + ) + + mock_create.assert_called_once() + request_body = mock_create.call_args.kwargs + + assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} + assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) + assert "prompt_cache_options" not in json.dumps(request_body) + + +@pytest.mark.usefixtures("_no_openai_api_base_override") +def test_completion_default_api_base_sends_prompt_cache_breakpoint_for_gpt_5_6(): + from openai import OpenAI + + client = OpenAI(api_key="fake-api-key") + request_body = _openai_chat_create_kwargs(client, model="gpt-5.6") + + assert request_body["messages"][0]["content"] == [ + {"type": "text", "text": "sys", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] + assert request_body["extra_body"]["prompt_cache_options"] == {"mode": "explicit"} + + +_SUBSCRIPTION_OAUTH_CREDENTIAL = "Bearer sk-ant-oat01-fake-subscription-token-for-testing-0123456789" + + +def _scoped_headers_for_oauth_request(): + from litellm.types.utils import ProviderSpecificHeader + + return [ + ProviderSpecificHeader( + custom_llm_provider="anthropic,bedrock,vertex_ai", + extra_headers={"anthropic-version": "2023-06-01"}, + ), + ProviderSpecificHeader( + custom_llm_provider="anthropic", + extra_headers={"authorization": _SUBSCRIPTION_OAUTH_CREDENTIAL}, + ), + ] + + +def _run_anthropic_hop_with_shared_headers(shared_headers): + litellm.completion( + model="anthropic/claude-3-5-sonnet-20240620", + messages=[{"role": "user", "content": "Say OK"}], + extra_headers=shared_headers, + provider_specific_header=_scoped_headers_for_oauth_request(), + api_key="sk-fake-anthropic-key", + mock_response="OK", + ) + + +def test_completion_does_not_mutate_caller_supplied_headers(): + shared_headers = {"x-tenant": "acme"} + + _run_anthropic_hop_with_shared_headers(shared_headers) + + assert shared_headers == {"x-tenant": "acme"} + + +def test_anthropic_oauth_credential_does_not_persist_into_next_provider_hop(): + shared_headers = {"x-tenant": "acme"} + + _run_anthropic_hop_with_shared_headers(shared_headers) + + leaked = [name for name, value in shared_headers.items() if value == _SUBSCRIPTION_OAUTH_CREDENTIAL] + assert leaked == [] + assert "anthropic-version" not in shared_headers + + +STREAM_COST_MODEL = "gpt-4o" +STREAMED_USAGE = {"prompt_tokens": 137, "completion_tokens": 42, "total_tokens": 179} + + +def _text_chunk(content, finish_reason=None, usage=None): + chunk = { + "id": "chatcmpl-stream-cost", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": STREAM_COST_MODEL, + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": content}, + "finish_reason": finish_reason, + } + ], + } + if usage is not None: + chunk["usage"] = usage + return chunk + + +def _priced_at(prompt_tokens, completion_tokens): + prices = litellm.model_cost[STREAM_COST_MODEL] + return ( + prompt_tokens * prices["input_cost_per_token"] + + completion_tokens * prices["output_cost_per_token"] + ) + + +@pytest.fixture +def local_cost_map(monkeypatch): + """The prices these tests assert are the checked-in ones. Setting the environment + variable alone does not reload the map, so pin the map itself. + + Prices are read through two separate lru_caches, so pinning ``model_cost`` is not + enough on its own: an entry warmed against the network-fetched map keeps its old + prices and billing reads those while the assertions read the pinned map. + ``_invalidate_model_cost_lowercase_map`` clears both caches, where + ``get_model_info.cache_clear`` reaches only one. Invalidate on the way in and out + so entries never leak across tests in either direction.""" + from litellm.utils import _invalidate_model_cost_lowercase_map + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + _invalidate_model_cost_lowercase_map() + yield + _invalidate_model_cost_lowercase_map() + + +def test_a_streamed_response_bills_the_usage_the_provider_reported(local_cost_map): + rebuilt = litellm.stream_chunk_builder( + chunks=[ + _text_chunk("Hello"), + _text_chunk(" there"), + _text_chunk(None, finish_reason="stop", usage=STREAMED_USAGE), + ], + messages=[{"role": "user", "content": "hi"}], + ) + + assert rebuilt.choices[0].message.content == "Hello there" + assert rebuilt.usage.prompt_tokens == STREAMED_USAGE["prompt_tokens"] + assert rebuilt.usage.completion_tokens == STREAMED_USAGE["completion_tokens"] + + cost = litellm.completion_cost(completion_response=rebuilt, model=STREAM_COST_MODEL) + + assert cost == pytest.approx(_priced_at(137, 42)) + + +def test_streaming_and_not_streaming_bill_the_same_usage_the_same(local_cost_map): + rebuilt = litellm.stream_chunk_builder( + chunks=[ + _text_chunk("Hello"), + _text_chunk(" there"), + _text_chunk(None, finish_reason="stop", usage=STREAMED_USAGE), + ], + messages=[{"role": "user", "content": "hi"}], + ) + whole = litellm.ModelResponse( + id="chatcmpl-stream-cost", + model=STREAM_COST_MODEL, + object="chat.completion", + created=1700000000, + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello there"}, + "finish_reason": "stop", + } + ], + usage=STREAMED_USAGE, + ) + + assert litellm.completion_cost( + completion_response=rebuilt, model=STREAM_COST_MODEL + ) == pytest.approx(litellm.completion_cost(completion_response=whole, model=STREAM_COST_MODEL)) + + +def test_a_stream_that_reported_no_usage_is_still_billed(local_cost_map): + rebuilt = litellm.stream_chunk_builder( + chunks=[ + _text_chunk("Hello"), + _text_chunk(" there"), + _text_chunk(None, finish_reason="stop"), + ], + messages=[{"role": "user", "content": "hi"}], + ) + + assert rebuilt.usage.prompt_tokens > 0 + assert rebuilt.usage.completion_tokens > 0 + + cost = litellm.completion_cost(completion_response=rebuilt, model=STREAM_COST_MODEL) + + assert cost > 0 + assert cost == pytest.approx( + _priced_at(rebuilt.usage.prompt_tokens, rebuilt.usage.completion_tokens) + ) + + +@pytest.mark.asyncio +async def test_acompletion_resolves_provider_from_api_base(): + response = await litellm.acompletion( + model="deepseek-chat", + api_base="https://api.deepseek.com/v1", + api_key="fake-key", + messages=[{"role": "user", "content": "hi"}], + mock_response="resolved", + ) + + assert response.choices[0].message.content == "resolved" + + +@dataclass(frozen=True, slots=True) +class _RecordedSpeechSuccess: + call_type: str | None + spend_metadata: Mapping[str, object] + response_cost: float | None + logged_response_cost: float | None + + +def _record_speech_success(payload: dict[str, object]) -> _RecordedSpeechSuccess: + call_type: Final = payload.get("call_type") + response_cost: Final = payload.get("response_cost") + logging_payload: Final = payload.get("standard_logging_object") + logged_cost: Final = logging_payload.get("response_cost") if isinstance(logging_payload, dict) else None + return _RecordedSpeechSuccess( + call_type=call_type if isinstance(call_type, str) else None, + spend_metadata=get_litellm_metadata_from_kwargs(payload), + response_cost=response_cost if isinstance(response_cost, float) else None, + logged_response_cost=logged_cost if isinstance(logged_cost, float) else None, + ) + + +class _SuccessEventRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: list[_RecordedSpeechSuccess] = [] # mutable-ok: test recorder of success-callback events + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.events.append(_record_speech_success(kwargs)) + + +async def _wait_for_success_event(recorder: _SuccessEventRecorder, call_type: str) -> _RecordedSpeechSuccess: + for _ in range(100): + if (event := next((e for e in recorder.events if e.call_type == call_type), None)) is not None: + return event + await asyncio.sleep(0.05) + pytest.fail(f"no {call_type} success event; got {[e.call_type for e in recorder.events]}") + + +def _gemini_tts_generate_content_response() -> dict[str, object]: + return { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "audio/L16;codec=pcm;rate=24000", + "data": base64.b64encode(b"pcm-audio-bytes").decode(), + } + } + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 60, + "totalTokenCount": 65, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 5}], + "candidatesTokensDetails": [{"modality": "AUDIO", "tokenCount": 60}], + }, + "modelVersion": "gemini-2.5-flash-preview-tts", + } + + +@pytest.mark.asyncio +async def test_aspeech_gemini_bridge_keeps_proxy_metadata_for_spend_tracking( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + recorder: Final = _SuccessEventRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + mock_route: Final = respx_mock.post( + url__regex=r"https://generativelanguage\.googleapis\.com/v1beta/models/gemini-2\.5-flash-preview-tts:generateContent.*" + ).mock(return_value=httpx.Response(200, json=_gemini_tts_generate_content_response())) + + await litellm.aspeech( + model="gemini/gemini-2.5-flash-preview-tts", + input="spend tracking check", + voice="Kore", + api_key="fake-gemini-key", + metadata={"user_api_key": "hashed-virtual-key", "user_api_key_user_id": "user-1"}, + ) + + assert mock_route.called + assert mock_route.calls.last.request.headers["x-goog-api-key"] == "fake-gemini-key" + speech_event: Final = await _wait_for_success_event(recorder, call_type="aspeech") + assert speech_event.spend_metadata["user_api_key"] == "hashed-virtual-key" + assert speech_event.spend_metadata["user_api_key_user_id"] == "user-1" + expected_prompt_cost, expected_completion_cost = litellm.cost_per_token( + model="gemini/gemini-2.5-flash-preview-tts", + usage_object=Usage(prompt_tokens=5, completion_tokens=60, total_tokens=65), + ) + expected_cost: Final = expected_prompt_cost + expected_completion_cost + assert expected_cost > 0 + assert speech_event.response_cost == pytest.approx(expected_cost) + assert speech_event.logged_response_cost == pytest.approx(expected_cost) + + +def _stream_builder_text_chunk(model: str, content: str, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-cost", + created=1724900000, + model=model, + object="chat.completion.chunk", + choices=[StreamingChoices(finish_reason=finish_reason, index=0, delta=Delta(content=content, role="assistant"))], + ) + + +def test_stream_chunk_builder_sets_hidden_response_cost_for_known_model(): + chunks: Final = [ + _stream_builder_text_chunk("gpt-4o", "Hello "), + _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), + ] + + response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) + + assert response is not None + prompt_cost, completion_cost = litellm.cost_per_token(model="gpt-4o", usage_object=response.usage) + expected_cost: Final = prompt_cost + completion_cost + assert expected_cost > 0 + assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) + + +def test_stream_chunk_builder_unknown_model_leaves_response_cost_unset(): + chunks: Final = [ + _stream_builder_text_chunk("totally-unknown-model-xyz", "Hello "), + _stream_builder_text_chunk("totally-unknown-model-xyz", "world.", finish_reason="stop"), + ] + + response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) + + assert response is not None + assert response._hidden_params.get("response_cost") is None + assert response.choices[0].message.content == "Hello world." + + +def test_stream_chunk_builder_prices_proxy_alias_via_model_map(): + chunks: Final = [ + _stream_builder_text_chunk("claude-opus-5", "Hello "), + _stream_builder_text_chunk("claude-opus-5", "world.", finish_reason="stop"), + ] + for chunk in chunks: + chunk._hidden_params = {"custom_llm_provider": "openai"} + + response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) + + assert response is not None + assert response._hidden_params["custom_llm_provider"] == "openai" + prompt_cost, completion_cost = litellm.cost_per_token(model="claude-opus-5", usage_object=response.usage) + expected_cost: Final = prompt_cost + completion_cost + assert expected_cost > 0 + assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) + + +def _stream_builder_logging_obj(model: str = "gpt-4o", custom_llm_provider: str = "openai") -> LiteLLMLogging: + logging_obj: Final = LiteLLMLogging( + model=model, + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="test-function-id", + ) + logging_obj.update_environment_variables( + model=model, + user=None, + optional_params={}, + litellm_params={"custom_llm_provider": custom_llm_provider}, + custom_llm_provider=custom_llm_provider, + ) + return logging_obj + + +def test_stream_chunk_builder_stamps_streaming_usage_cost_by_default(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) + chunks: Final = [ + _stream_builder_text_chunk("gpt-4o", "Hello "), + _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=_stream_builder_logging_obj() + ) + + assert response is not None + usage_cost: Final = getattr(response.usage, "cost", None) + assert usage_cost is not None + assert usage_cost > 0 + assert response._hidden_params["response_cost"] == pytest.approx(usage_cost) + + +def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable(): + import time as time_module + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + logging_obj: Final = LiteLLMLogging( + model="us.anthropic.claude-opus-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=time_module.time(), + litellm_call_id="stream-builder-alias-unpriceable", + function_id="1", + ) + logging_obj.model_call_details["custom_llm_provider"] = "bedrock" + logging_obj.optional_params = {} + usage_chunk: Final = _stream_builder_text_chunk("bedrock-claude-opus-5", "") + usage_chunk.usage = Usage(prompt_tokens=40, completion_tokens=5, total_tokens=45) + chunks: Final = [ + _stream_builder_text_chunk("bedrock-claude-opus-5", "Hello ", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj + ) + + assert response is not None + assert getattr(response.usage, "cost", None) is None + assert response._hidden_params.get("response_cost") is None + + +def test_stream_chunk_builder_keeps_provider_reported_usage_cost(): + usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") + usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5) + chunks: Final = [ + _stream_builder_text_chunk("gpt-4o", "Hello "), + _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=_stream_builder_logging_obj() + ) + + assert response is not None + assert getattr(response.usage, "cost", None) == pytest.approx(0.5) + assert response._hidden_params["response_cost"] == pytest.approx(0.5) + + +def test_stream_chunk_builder_prices_alias_from_openai_sdk_usage_chunk(): + from openai.types.completion_usage import CompletionUsage + + usage_chunk: Final = _stream_builder_text_chunk("mantle-claude", "") + usage_chunk.usage = CompletionUsage(prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704) + assert type(usage_chunk.usage) is CompletionUsage + chunks: Final = [ + _stream_builder_text_chunk("mantle-claude", "Hello "), + _stream_builder_text_chunk("mantle-claude", "world.", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) + + assert response is not None + assert response.usage.prompt_tokens == 20 + assert response.usage.completion_tokens == 60 + assert getattr(response.usage, "cost", None) == pytest.approx(0.000704) + assert response._hidden_params["response_cost"] == pytest.approx(0.000704) + + +def test_stream_chunk_builder_leaves_xai_reported_cost_to_the_calculator(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "cost_margin_config", {"xai": 0.5}) + usage_chunk: Final = _stream_builder_text_chunk("grok-4", "") + usage_chunk.usage = Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7, cost=0.42) + chunks: Final = [ + _stream_builder_text_chunk("grok-4", "Hello "), + _stream_builder_text_chunk("grok-4", "world.", finish_reason="stop"), + usage_chunk, + ] + logging_obj: Final = _stream_builder_logging_obj(model="grok-4", custom_llm_provider="xai") + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj + ) + + assert response is not None + assert getattr(response.usage, "cost", None) == pytest.approx(0.42) + assert response._hidden_params.get("response_cost") is None + assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.63) + + +def test_speech_mistral_dispatches_and_decodes_audio(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + audio_bytes: Final = b"ID3-fake-mp3-bytes" + mock_route: Final = respx_mock.post("https://api.mistral.ai/v1/audio/speech").mock( + return_value=httpx.Response(200, json={"audio_data": base64.b64encode(audio_bytes).decode()}) + ) + + response: Final = litellm.speech( + model="mistral/voxtral-mini-tts-2603", + input="hello from litellm", + voice="en_paul_neutral", + response_format="wav", + speed=2, + instructions="sound cheerful", + ) + + assert mock_route.called + request_body: Final = json.loads(mock_route.calls.last.request.content) + assert request_body == { + "model": "voxtral-mini-tts-2603", + "input": "hello from litellm", + "voice_id": "en_paul_neutral", + "response_format": "wav", + } + assert mock_route.calls.last.request.headers["authorization"] == "Bearer sk-mistral-test" + assert response.content == audio_bytes + + +def test_speech_mistral_routes_to_configured_api_base(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + audio_bytes: Final = b"ID3-gateway-bytes" + gateway_route: Final = respx_mock.post("https://mistral.gateway.internal/v1/audio/speech").mock( + return_value=httpx.Response(200, json={"audio_data": base64.b64encode(audio_bytes).decode()}) + ) + + response: Final = litellm.speech( + model="mistral/voxtral-mini-tts-2603", + input="hello from litellm", + voice="en_paul_neutral", + api_base="https://mistral.gateway.internal", + ) + + assert gateway_route.called + assert response.content == audio_bytes + + +FOUNDRY_HOST: Final = "https://my-project.services.ai.azure.com" + + +def test_azure_ai_transcription_on_a_foundry_host_uses_the_azure_openai_deployment_route( + respx_mock: respx.MockRouter, +): + route: Final = respx_mock.post( + url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/whisper-1/audio/transcriptions\?api-version=.+" + ).mock(return_value=httpx.Response(200, json={"text": "hello"})) + + response: Final = litellm.transcription( + model="azure_ai/whisper-1", + file=("tone.wav", b"RIFF\x00\x00\x00\x00WAVE", "audio/wav"), + api_base=FOUNDRY_HOST, + api_key="fake-key", + ) + + assert route.called + assert response.text == "hello" + + +def test_azure_ai_speech_on_a_foundry_host_uses_the_azure_openai_deployment_route( + respx_mock: respx.MockRouter, +): + route: Final = respx_mock.post( + url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/tts-1/audio/speech\?api-version=.+" + ).mock(return_value=httpx.Response(200, content=b"mp3-bytes")) + + response: Final = litellm.speech( + model="azure_ai/tts-1", + input="hello", + voice="alloy", + api_base=FOUNDRY_HOST, + api_key="fake-key", + ) + + assert route.called + assert response.content == b"mp3-bytes" + + +FORWARDED_CLIENT_HEADERS: Final = {"x-forwarded-for": "10.0.0.1", "x-amzn-trace-id": "Root=1-lit7694"} + + +def _chat_completion_json() -> Mapping[str, object]: + return { + "id": "chatcmpl-lit7694", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.4", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + + +def _chat_completion_sse() -> bytes: + chunk: Final = { + "id": "chatcmpl-lit7694", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-5.4", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + } + return f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode() + + +@pytest.mark.parametrize("stream", [False, True]) +def test_bridged_responses_with_openai_http_handler_keeps_forwarded_headers_out_of_the_body( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, stream: bool +): + monkeypatch.setenv("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER", "true") + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response(200, content=_chat_completion_sse(), headers={"content-type": "text/event-stream"}) + if stream + else httpx.Response(200, json=_chat_completion_json()) + ) + + response: Final = litellm.responses( + model="openai/gpt-5.4", + input="Reply with the single word ok", + stream=stream, + use_chat_completions_api=True, + headers=dict(FORWARDED_CLIENT_HEADERS), + api_key="sk-test", + ) + if stream: + list(response) + + assert route.called + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert "extra_headers" not in body + assert body["model"] == "gpt-5.4" + assert {k: request.headers[k] for k in FORWARDED_CLIENT_HEADERS} == FORWARDED_CLIENT_HEADERS + + +@pytest.mark.parametrize("http2_on", [True, False]) +def test_aiohttp_openai_warns_only_when_http2_enabled( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, http2_on: bool +): + from litellm.main import base_llm_aiohttp_handler + + monkeypatch.setattr(litellm, "http2", http2_on) + monkeypatch.delenv("LITELLM_HTTP2", raising=False) + + handler_completion: Final = MagicMock(return_value=MagicMock()) + monkeypatch.setattr(base_llm_aiohttp_handler, "completion", handler_completion) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + litellm.completion( + model="aiohttp_openai/gpt-4o", + messages=[{"role": "user", "content": "hi"}], + api_key="sk-test", + ) + + assert handler_completion.called + warned: Final = "aiohttp_openai/ always uses aiohttp" in caplog.text + assert warned is http2_on + + +@pytest.mark.parametrize("tool_choice", [{"type": "bogus"}, {"name": "lookup_fruit"}, {"type": "file_search"}]) +def test_completion_rejects_untranslatable_tool_choice_with_a_400(tool_choice): + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.completion( + model="anthropic/claude-haiku-4-5", + messages=[{"role": "user", "content": "Which fruit is red?"}], + tools=[{"type": "function", "function": {"name": "lookup_fruit", "parameters": {"type": "object"}}}], + tool_choice=tool_choice, + api_key="sk-unused", + ) + assert exc_info.value.status_code == 400 + assert f"tool_choice={tool_choice}" in str(exc_info.value) diff --git a/tests/test_litellm/test_main_module_header.py b/tests/unit/test_main_module_header.py similarity index 100% rename from tests/test_litellm/test_main_module_header.py rename to tests/unit/test_main_module_header.py diff --git a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py b/tests/unit/test_mistral_medium_3_5_model_metadata.py similarity index 100% rename from tests/test_litellm/test_mistral_medium_3_5_model_metadata.py rename to tests/unit/test_mistral_medium_3_5_model_metadata.py diff --git a/tests/test_litellm/test_mistral_small_4_0_model_metadata.py b/tests/unit/test_mistral_small_4_0_model_metadata.py similarity index 100% rename from tests/test_litellm/test_mistral_small_4_0_model_metadata.py rename to tests/unit/test_mistral_small_4_0_model_metadata.py diff --git a/tests/test_litellm/test_mistral_zai_glm_5_2_model_metadata.py b/tests/unit/test_mistral_zai_glm_5_2_model_metadata.py similarity index 100% rename from tests/test_litellm/test_mistral_zai_glm_5_2_model_metadata.py rename to tests/unit/test_mistral_zai_glm_5_2_model_metadata.py diff --git a/tests/test_litellm/test_model_block_unblock.py b/tests/unit/test_model_block_unblock.py similarity index 100% rename from tests/test_litellm/test_model_block_unblock.py rename to tests/unit/test_model_block_unblock.py diff --git a/tests/test_litellm/test_model_cost_aliases.py b/tests/unit/test_model_cost_aliases.py similarity index 100% rename from tests/test_litellm/test_model_cost_aliases.py rename to tests/unit/test_model_cost_aliases.py diff --git a/tests/test_litellm/test_model_param_helper.py b/tests/unit/test_model_param_helper.py similarity index 100% rename from tests/test_litellm/test_model_param_helper.py rename to tests/unit/test_model_param_helper.py diff --git a/tests/test_litellm/test_model_prices_schema.py b/tests/unit/test_model_prices_schema.py similarity index 100% rename from tests/test_litellm/test_model_prices_schema.py rename to tests/unit/test_model_prices_schema.py diff --git a/tests/test_litellm/test_model_response_normalization.py b/tests/unit/test_model_response_normalization.py similarity index 100% rename from tests/test_litellm/test_model_response_normalization.py rename to tests/unit/test_model_response_normalization.py diff --git a/tests/test_litellm/test_muse_spark_1_1_model_metadata.py b/tests/unit/test_muse_spark_1_1_model_metadata.py similarity index 100% rename from tests/test_litellm/test_muse_spark_1_1_model_metadata.py rename to tests/unit/test_muse_spark_1_1_model_metadata.py diff --git a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py b/tests/unit/test_muse_spark_1_2_model_metadata.py similarity index 100% rename from tests/test_litellm/test_muse_spark_1_2_model_metadata.py rename to tests/unit/test_muse_spark_1_2_model_metadata.py diff --git a/tests/test_litellm/test_muse_spark_1_3_model_metadata.py b/tests/unit/test_muse_spark_1_3_model_metadata.py similarity index 100% rename from tests/test_litellm/test_muse_spark_1_3_model_metadata.py rename to tests/unit/test_muse_spark_1_3_model_metadata.py diff --git a/tests/test_litellm/test_mutation_report.py b/tests/unit/test_mutation_report.py similarity index 100% rename from tests/test_litellm/test_mutation_report.py rename to tests/unit/test_mutation_report.py diff --git a/tests/test_litellm/test_nested_drop_params.py b/tests/unit/test_nested_drop_params.py similarity index 100% rename from tests/test_litellm/test_nested_drop_params.py rename to tests/unit/test_nested_drop_params.py diff --git a/tests/test_litellm/test_non_chat_routes_open_llm_spans.py b/tests/unit/test_non_chat_routes_open_llm_spans.py similarity index 100% rename from tests/test_litellm/test_non_chat_routes_open_llm_spans.py rename to tests/unit/test_non_chat_routes_open_llm_spans.py diff --git a/tests/test_litellm/test_openai_embedding_encoding_format_default.py b/tests/unit/test_openai_embedding_encoding_format_default.py similarity index 100% rename from tests/test_litellm/test_openai_embedding_encoding_format_default.py rename to tests/unit/test_openai_embedding_encoding_format_default.py diff --git a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py b/tests/unit/test_openai_service_tier_long_context_pricing.py similarity index 100% rename from tests/test_litellm/test_openai_service_tier_long_context_pricing.py rename to tests/unit/test_openai_service_tier_long_context_pricing.py diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/unit/test_pre_commit_lint.py similarity index 100% rename from tests/test_litellm/test_pre_commit_lint.py rename to tests/unit/test_pre_commit_lint.py diff --git a/tests/test_litellm/test_prisma_generate_if_needed.py b/tests/unit/test_prisma_generate_if_needed.py similarity index 100% rename from tests/test_litellm/test_prisma_generate_if_needed.py rename to tests/unit/test_prisma_generate_if_needed.py diff --git a/tests/test_litellm/test_process_helpers.py b/tests/unit/test_process_helpers.py similarity index 100% rename from tests/test_litellm/test_process_helpers.py rename to tests/unit/test_process_helpers.py diff --git a/tests/test_litellm/test_project_alias_tracking.py b/tests/unit/test_project_alias_tracking.py similarity index 100% rename from tests/test_litellm/test_project_alias_tracking.py rename to tests/unit/test_project_alias_tracking.py diff --git a/tests/test_litellm/test_project_tags_pydantic.py b/tests/unit/test_project_tags_pydantic.py similarity index 100% rename from tests/test_litellm/test_project_tags_pydantic.py rename to tests/unit/test_project_tags_pydantic.py diff --git a/tests/test_litellm/test_proxy_auth.py b/tests/unit/test_proxy_auth.py similarity index 100% rename from tests/test_litellm/test_proxy_auth.py rename to tests/unit/test_proxy_auth.py diff --git a/tests/test_litellm/test_rag_openai_ingestion.py b/tests/unit/test_rag_openai_ingestion.py similarity index 100% rename from tests/test_litellm/test_rag_openai_ingestion.py rename to tests/unit/test_rag_openai_ingestion.py diff --git a/tests/test_litellm/test_rate_limit_error_unification.py b/tests/unit/test_rate_limit_error_unification.py similarity index 100% rename from tests/test_litellm/test_rate_limit_error_unification.py rename to tests/unit/test_rate_limit_error_unification.py diff --git a/tests/test_litellm/test_read_rc_version.py b/tests/unit/test_read_rc_version.py similarity index 100% rename from tests/test_litellm/test_read_rc_version.py rename to tests/unit/test_read_rc_version.py diff --git a/tests/test_litellm/test_redact_string_in_error_paths.py b/tests/unit/test_redact_string_in_error_paths.py similarity index 100% rename from tests/test_litellm/test_redact_string_in_error_paths.py rename to tests/unit/test_redact_string_in_error_paths.py diff --git a/tests/test_litellm/test_redis.py b/tests/unit/test_redis.py similarity index 100% rename from tests/test_litellm/test_redis.py rename to tests/unit/test_redis.py diff --git a/tests/test_litellm/test_redis_credential_provider.py b/tests/unit/test_redis_credential_provider.py similarity index 100% rename from tests/test_litellm/test_redis_credential_provider.py rename to tests/unit/test_redis_credential_provider.py diff --git a/tests/test_litellm/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py similarity index 100% rename from tests/test_litellm/test_register_model_custom_pricing.py rename to tests/unit/test_register_model_custom_pricing.py diff --git a/tests/test_litellm/test_register_model_zero_cost_persistence.py b/tests/unit/test_register_model_zero_cost_persistence.py similarity index 100% rename from tests/test_litellm/test_register_model_zero_cost_persistence.py rename to tests/unit/test_register_model_zero_cost_persistence.py diff --git a/tests/test_litellm/test_replicate_model_key_format.py b/tests/unit/test_replicate_model_key_format.py similarity index 100% rename from tests/test_litellm/test_replicate_model_key_format.py rename to tests/unit/test_replicate_model_key_format.py diff --git a/tests/test_litellm/test_responses_api_bridge_non_stream.py b/tests/unit/test_responses_api_bridge_non_stream.py similarity index 100% rename from tests/test_litellm/test_responses_api_bridge_non_stream.py rename to tests/unit/test_responses_api_bridge_non_stream.py diff --git a/tests/test_litellm/test_responses_id_security.py b/tests/unit/test_responses_id_security.py similarity index 94% rename from tests/test_litellm/test_responses_id_security.py rename to tests/unit/test_responses_id_security.py index a6081670172..704a52fc202 100644 --- a/tests/test_litellm/test_responses_id_security.py +++ b/tests/unit/test_responses_id_security.py @@ -4,7 +4,7 @@ Tests for ResponsesIDSecurity hook. Tests the security hook that prevents user B from seeing response from user A. """ -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch import pytest from fastapi import HTTPException @@ -113,63 +113,6 @@ class TestDecryptResponseId: assert team_id is None -class TestEncryptResponseId: - """Test _encrypt_response_id function""" - - @pytest.mark.skip( - reason="Flaky on CI; disabling temporarily until responses_id_security is fixed" - ) - def test_encrypt_response_id_success( - self, responses_id_security, mock_user_api_key_dict - ): - """Test encrypting a response ID with user information""" - mock_response = ResponsesAPIResponse( - id="resp_123", created_at=1234567890, output=[], status="completed" - ) - - with patch( - "litellm.proxy.hooks.responses_id_security.encrypt_value_helper" - ) as mock_encrypt: - mock_encrypt.return_value = "encrypted_base64_value" - - with patch.object( - responses_id_security, "_get_signing_key", return_value="test-key" - ): - result = responses_id_security._encrypt_response_id( - mock_response, mock_user_api_key_dict - ) - - assert result.id == "resp_encrypted_base64_value" - assert result.id.startswith("resp_") - mock_encrypt.assert_called_once() - - @pytest.mark.skip( - reason="Flaky on CI; disabling temporarily until responses_id_security is fixed" - ) - def test_encrypt_response_id_maintains_prefix( - self, responses_id_security, mock_user_api_key_dict - ): - """Test that encrypted response ID maintains 'resp_' prefix""" - mock_response = ResponsesAPIResponse( - id="resp_456", created_at=1234567890, output=[], status="in_progress" - ) - - with patch( - "litellm.proxy.common_utils.encrypt_decrypt_utils._get_salt_key", - return_value="test-salt-key", - ): - with patch.object( - responses_id_security, "_get_signing_key", return_value="test-key" - ): - result = responses_id_security._encrypt_response_id( - mock_response, mock_user_api_key_dict - ) - - assert result.id.startswith("resp_") - # The encrypted ID should be different from the original - assert result.id != "resp_456" - - class TestCheckUserAccessToResponseId: """Test check_user_access_to_response_id function""" @@ -857,7 +800,6 @@ class TestAsyncPostCallSuccessHook: assert result == mock_response - _FABRICATED_PROVIDER_RESPONSE_ID = "resp_fabricatedprovideridaaaaaaaaaaaaaaaa" _FABRICATED_UNMANAGED_ID = "resp_fabricatedunmanagedidbbbbbbbbbbbbbbbb" _UNIT_TEST_SALT_KEY = "lit6837-unit-test-salt-key" diff --git a/tests/test_litellm/test_responses_streaming_container_ownership.py b/tests/unit/test_responses_streaming_container_ownership.py similarity index 100% rename from tests/test_litellm/test_responses_streaming_container_ownership.py rename to tests/unit/test_responses_streaming_container_ownership.py diff --git a/tests/test_litellm/test_retrieve_batch_bedrock_dispatch.py b/tests/unit/test_retrieve_batch_bedrock_dispatch.py similarity index 100% rename from tests/test_litellm/test_retrieve_batch_bedrock_dispatch.py rename to tests/unit/test_retrieve_batch_bedrock_dispatch.py diff --git a/tests/test_litellm/test_router.py b/tests/unit/test_router/test_router.py similarity index 100% rename from tests/test_litellm/test_router.py rename to tests/unit/test_router/test_router.py diff --git a/tests/test_litellm/test_router_block_helpers.py b/tests/unit/test_router_block_helpers.py similarity index 100% rename from tests/test_litellm/test_router_block_helpers.py rename to tests/unit/test_router_block_helpers.py diff --git a/tests/test_litellm/test_router_exception_redaction.py b/tests/unit/test_router_exception_redaction.py similarity index 100% rename from tests/test_litellm/test_router_exception_redaction.py rename to tests/unit/test_router_exception_redaction.py diff --git a/tests/test_litellm/test_router_google_genai.py b/tests/unit/test_router_google_genai.py similarity index 100% rename from tests/test_litellm/test_router_google_genai.py rename to tests/unit/test_router_google_genai.py diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py similarity index 100% rename from tests/test_litellm/test_router_model_cost_isolation.py rename to tests/unit/test_router_model_cost_isolation.py diff --git a/tests/test_litellm/test_router_order_fallback.py b/tests/unit/test_router_order_fallback.py similarity index 100% rename from tests/test_litellm/test_router_order_fallback.py rename to tests/unit/test_router_order_fallback.py diff --git a/tests/test_litellm/test_router_per_deployment_num_retries.py b/tests/unit/test_router_per_deployment_num_retries.py similarity index 100% rename from tests/test_litellm/test_router_per_deployment_num_retries.py rename to tests/unit/test_router_per_deployment_num_retries.py diff --git a/tests/test_litellm/test_router_redis_init.py b/tests/unit/test_router_redis_init.py similarity index 100% rename from tests/test_litellm/test_router_redis_init.py rename to tests/unit/test_router_redis_init.py diff --git a/tests/test_litellm/test_router_retry_backoff_headers.py b/tests/unit/test_router_retry_backoff_headers.py similarity index 100% rename from tests/test_litellm/test_router_retry_backoff_headers.py rename to tests/unit/test_router_retry_backoff_headers.py diff --git a/tests/test_litellm/test_router_retry_non_retryable_errors.py b/tests/unit/test_router_retry_non_retryable_errors.py similarity index 100% rename from tests/test_litellm/test_router_retry_non_retryable_errors.py rename to tests/unit/test_router_retry_non_retryable_errors.py diff --git a/tests/test_litellm/test_router_retry_policy_update.py b/tests/unit/test_router_retry_policy_update.py similarity index 100% rename from tests/test_litellm/test_router_retry_policy_update.py rename to tests/unit/test_router_retry_policy_update.py diff --git a/tests/test_litellm/test_router_silent_experiment.py b/tests/unit/test_router_silent_experiment.py similarity index 92% rename from tests/test_litellm/test_router_silent_experiment.py rename to tests/unit/test_router_silent_experiment.py index d62962da275..ab65e09e133 100644 --- a/tests/test_litellm/test_router_silent_experiment.py +++ b/tests/unit/test_router_silent_experiment.py @@ -388,47 +388,6 @@ async def test_shadow_of_a_shadow_is_not_launched(recording_logger): assert model_groups == ["shadow-a"] -def test_silent_experiment_completion_direct(): - """ - Test _silent_experiment_completion directly (for router code coverage). - Mocks router.completion to avoid real API call. - """ - model_list = [ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake-key"}, - }, - ] - router = Router(model_list=model_list) - messages = [{"role": "user", "content": "hi"}] - with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None): - router._silent_experiment_completion( - silent_model="gpt-3.5-turbo", - messages=messages, - ) - - -@pytest.mark.asyncio -async def test_silent_experiment_acompletion_direct(): - """ - Test _silent_experiment_acompletion directly (for router code coverage). - Mocks router.acompletion to avoid real API call. - """ - model_list = [ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake-key"}, - }, - ] - router = Router(model_list=model_list) - messages = [{"role": "user", "content": "hi"}] - with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None): - await router._silent_experiment_acompletion( - silent_model="gpt-3.5-turbo", - messages=messages, - ) - - @pytest.mark.asyncio async def test_run_silent_experiment_drains_stream_so_callbacks_fire(recording_logger): router = Router(model_list=_streaming_model_list(None)) @@ -602,3 +561,44 @@ def test_router_silent_experiment_completion(): assert silent_call[1]["model"] == "openai/gpt-4" # Verify model_group is set to the silent model name for correct metric attribution assert silent_call[1]["metadata"]["model_group"] == "silent-model" + + +SILENT_EXPERIMENT_RUNNERS: Final = ( + pytest.param(lambda router, **kwargs: router._silent_experiment_completion(**kwargs), id="sync"), + pytest.param(lambda router, **kwargs: asyncio.run(router._silent_experiment_acompletion(**kwargs)), id="async"), +) + + +@pytest.mark.parametrize("run_silent_experiment", SILENT_EXPERIMENT_RUNNERS) +def test_silent_experiment_sends_shadow_request_attributed_to_the_silent_model(run_silent_experiment): + router = Router(model_list=_streaming_model_list(["shadow-a"])) + primary_metadata: Final = {"model_group": "primary-model"} + with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None) as acompletion: + run_silent_experiment( + router, + silent_model="shadow-a", + messages=[{"role": "user", "content": "hi"}], + metadata=primary_metadata, + ) + + acompletion.assert_awaited_once() + shadow_call: Final = acompletion.await_args.kwargs + assert shadow_call["model"] == "shadow-a" + assert shadow_call["messages"] == [{"role": "user", "content": "hi"}] + assert shadow_call["metadata"]["model_group"] == "shadow-a" + assert shadow_call["metadata"]["is_silent_experiment"] is True + assert primary_metadata == {"model_group": "primary-model"} + + +@pytest.mark.parametrize("run_silent_experiment", SILENT_EXPERIMENT_RUNNERS) +def test_silent_experiment_does_not_launch_from_a_shadow_request(run_silent_experiment): + router = Router(model_list=_streaming_model_list(["shadow-a"])) + with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None) as acompletion: + run_silent_experiment( + router, + silent_model="shadow-a", + messages=[{"role": "user", "content": "hi"}], + metadata={"is_silent_experiment": True}, + ) + + acompletion.assert_not_awaited() diff --git a/tests/test_litellm/test_router_streaming_fallback_metadata.py b/tests/unit/test_router_streaming_fallback_metadata.py similarity index 100% rename from tests/test_litellm/test_router_streaming_fallback_metadata.py rename to tests/unit/test_router_streaming_fallback_metadata.py diff --git a/tests/test_litellm/test_router_weighted_failover.py b/tests/unit/test_router_weighted_failover.py similarity index 100% rename from tests/test_litellm/test_router_weighted_failover.py rename to tests/unit/test_router_weighted_failover.py diff --git a/tests/test_litellm/test_ruff_strict_gate.py b/tests/unit/test_ruff_strict_gate.py similarity index 100% rename from tests/test_litellm/test_ruff_strict_gate.py rename to tests/unit/test_ruff_strict_gate.py diff --git a/tests/test_litellm/test_sambanova_model_metadata.py b/tests/unit/test_sambanova_model_metadata.py similarity index 100% rename from tests/test_litellm/test_sambanova_model_metadata.py rename to tests/unit/test_sambanova_model_metadata.py diff --git a/tests/test_litellm/test_secret_redaction.py b/tests/unit/test_secret_redaction.py similarity index 100% rename from tests/test_litellm/test_secret_redaction.py rename to tests/unit/test_secret_redaction.py diff --git a/tests/test_litellm/test_select_ui_test_scope.py b/tests/unit/test_select_ui_test_scope.py similarity index 100% rename from tests/test_litellm/test_select_ui_test_scope.py rename to tests/unit/test_select_ui_test_scope.py diff --git a/tests/test_litellm/test_service_logger.py b/tests/unit/test_service_logger.py similarity index 100% rename from tests/test_litellm/test_service_logger.py rename to tests/unit/test_service_logger.py diff --git a/tests/test_litellm/test_setup_wizard.py b/tests/unit/test_setup_wizard.py similarity index 100% rename from tests/test_litellm/test_setup_wizard.py rename to tests/unit/test_setup_wizard.py diff --git a/tests/test_litellm/test_shared_session_integration.py b/tests/unit/test_shared_session_integration.py similarity index 100% rename from tests/test_litellm/test_shared_session_integration.py rename to tests/unit/test_shared_session_integration.py diff --git a/tests/test_litellm/test_ssl_verify_unit.py b/tests/unit/test_ssl_verify_unit.py similarity index 83% rename from tests/test_litellm/test_ssl_verify_unit.py rename to tests/unit/test_ssl_verify_unit.py index c39362c01a2..f47cdf3e6cd 100644 --- a/tests/test_litellm/test_ssl_verify_unit.py +++ b/tests/unit/test_ssl_verify_unit.py @@ -50,41 +50,6 @@ class TestBaseAWSLLMSSLVerify: # Result depends on environment, just verify it doesn't crash assert result is not None or result is None # Can be None, True, False, or path - @patch("boto3.client") - def test_get_credentials_propagates_ssl_verify(self, mock_boto_client): - """Test that get_credentials propagates ssl_verify to boto3 clients.""" - base_llm = BaseAWSLLM() - - # Mock the boto3 client - mock_sts_client = Mock() - mock_sts_client.assume_role.return_value = { - "Credentials": { - "AccessKeyId": "test_key", - "SecretAccessKey": "test_secret", - "SessionToken": "test_token", - "Expiration": "2026-01-20T00:00:00Z", - } - } - mock_boto_client.return_value = mock_sts_client - - # Call get_credentials with ssl_verify parameter - cert_path = "/path/to/cert.pem" - try: - base_llm.get_credentials( - aws_access_key_id="test_key", - aws_secret_access_key="test_secret", - aws_region_name="us-east-1", - ssl_verify=cert_path, - ) - except Exception: - # May fail due to missing credentials, but we're checking the call - pass - - # Verify boto3.client was called with verify parameter - # Note: This test verifies the parameter is accepted, actual propagation - # is tested in integration tests - assert True # If we got here without error, parameter was accepted - class TestAimGuardrailSSLVerify: """Test SSL verification parameter handling in AimGuardrail.""" diff --git a/tests/test_litellm/test_stream_chunk_builder_annotations.py b/tests/unit/test_stream_chunk_builder_annotations.py similarity index 100% rename from tests/test_litellm/test_stream_chunk_builder_annotations.py rename to tests/unit/test_stream_chunk_builder_annotations.py diff --git a/tests/test_litellm/test_stream_chunk_builder_citations.py b/tests/unit/test_stream_chunk_builder_citations.py similarity index 100% rename from tests/test_litellm/test_stream_chunk_builder_citations.py rename to tests/unit/test_stream_chunk_builder_citations.py diff --git a/tests/test_litellm/test_stream_chunk_builder_images.py b/tests/unit/test_stream_chunk_builder_images.py similarity index 100% rename from tests/test_litellm/test_stream_chunk_builder_images.py rename to tests/unit/test_stream_chunk_builder_images.py diff --git a/tests/test_litellm/test_streaming_connection_cleanup.py b/tests/unit/test_streaming_connection_cleanup.py similarity index 100% rename from tests/test_litellm/test_streaming_connection_cleanup.py rename to tests/unit/test_streaming_connection_cleanup.py diff --git a/tests/test_litellm/test_sync_together_ai_models.py b/tests/unit/test_sync_together_ai_models.py similarity index 100% rename from tests/test_litellm/test_sync_together_ai_models.py rename to tests/unit/test_sync_together_ai_models.py diff --git a/tests/test_litellm/test_system_message_format_bug.py b/tests/unit/test_system_message_format_bug.py similarity index 100% rename from tests/test_litellm/test_system_message_format_bug.py rename to tests/unit/test_system_message_format_bug.py diff --git a/tests/test_litellm/test_test_quality_gate.py b/tests/unit/test_test_quality_gate.py similarity index 100% rename from tests/test_litellm/test_test_quality_gate.py rename to tests/unit/test_test_quality_gate.py diff --git a/tests/test_litellm/test_thinking_enabled.py b/tests/unit/test_thinking_enabled.py similarity index 100% rename from tests/test_litellm/test_thinking_enabled.py rename to tests/unit/test_thinking_enabled.py diff --git a/tests/test_litellm/test_together_ai_model_metadata.py b/tests/unit/test_together_ai_model_metadata.py similarity index 100% rename from tests/test_litellm/test_together_ai_model_metadata.py rename to tests/unit/test_together_ai_model_metadata.py diff --git a/tests/test_litellm/test_type_check_gate.py b/tests/unit/test_type_check_gate.py similarity index 100% rename from tests/test_litellm/test_type_check_gate.py rename to tests/unit/test_type_check_gate.py diff --git a/tests/test_litellm/test_type_discipline_gate.py b/tests/unit/test_type_discipline_gate.py similarity index 100% rename from tests/test_litellm/test_type_discipline_gate.py rename to tests/unit/test_type_discipline_gate.py diff --git a/tests/test_litellm/test_typesafe_model_metadata.py b/tests/unit/test_typesafe_model_metadata.py similarity index 100% rename from tests/test_litellm/test_typesafe_model_metadata.py rename to tests/unit/test_typesafe_model_metadata.py diff --git a/tests/test_litellm/test_unit_shard_missing_paths.py b/tests/unit/test_unit_shard_missing_paths.py similarity index 97% rename from tests/test_litellm/test_unit_shard_missing_paths.py rename to tests/unit/test_unit_shard_missing_paths.py index b91c2cff764..4fa9c5bd3c1 100644 --- a/tests/test_litellm/test_unit_shard_missing_paths.py +++ b/tests/unit/test_unit_shard_missing_paths.py @@ -36,6 +36,7 @@ def _run_shard(tmp_path: Path, test_path: str, workers: str) -> subprocess.Compl **os.environ, **_SHARD_ENV, "PATH": f"{shim_dir}{os.pathsep}{os.environ['PATH']}", + "GITHUB_OUTPUT": str(tmp_path / "github_output"), "TEST_PATH": test_path, "WORKERS": workers, }, diff --git a/tests/test_litellm/test_unit_shard_per_test_timeout.py b/tests/unit/test_unit_shard_per_test_timeout.py similarity index 100% rename from tests/test_litellm/test_unit_shard_per_test_timeout.py rename to tests/unit/test_unit_shard_per_test_timeout.py diff --git a/tests/test_litellm/test_utils.py b/tests/unit/test_utils.py similarity index 100% rename from tests/test_litellm/test_utils.py rename to tests/unit/test_utils.py diff --git a/tests/test_litellm/test_utils_module_docstring.py b/tests/unit/test_utils_module_docstring.py similarity index 100% rename from tests/test_litellm/test_utils_module_docstring.py rename to tests/unit/test_utils_module_docstring.py diff --git a/tests/test_litellm/test_uuid_helper.py b/tests/unit/test_uuid_helper.py similarity index 100% rename from tests/test_litellm/test_uuid_helper.py rename to tests/unit/test_uuid_helper.py diff --git a/tests/test_litellm/test_vcr_safe_body_matcher.py b/tests/unit/test_vcr_safe_body_matcher.py similarity index 98% rename from tests/test_litellm/test_vcr_safe_body_matcher.py rename to tests/unit/test_vcr_safe_body_matcher.py index 712ecf09911..cf4e4a1c276 100644 --- a/tests/test_litellm/test_vcr_safe_body_matcher.py +++ b/tests/unit/test_vcr_safe_body_matcher.py @@ -52,14 +52,6 @@ def test_safe_body_matcher_accepts_str_bytes_equivalent(): _safe_body_matcher(_req("hello"), _req(b"hello")) -def test_safe_body_matcher_handles_jsonl_without_crashing(): - jsonl = ( - b'{"recordId": "request-1", "modelInput": {}}\n' - b'{"recordId": "request-2", "modelInput": {}}\n' - ) - _safe_body_matcher(_req(jsonl), _req(jsonl)) - - def test_safe_body_matcher_rejects_different_jsonl_bodies(): a = b'{"recordId": "request-1"}\n{"recordId": "request-2"}\n' b = b'{"recordId": "request-1"}\n{"recordId": "request-3"}\n' diff --git a/tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py b/tests/unit/test_vertex_ai_xai_grok_prompt_caching_metadata.py similarity index 100% rename from tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py rename to tests/unit/test_vertex_ai_xai_grok_prompt_caching_metadata.py diff --git a/tests/test_litellm/test_video_generation.py b/tests/unit/test_video_generation.py similarity index 100% rename from tests/test_litellm/test_video_generation.py rename to tests/unit/test_video_generation.py diff --git a/tests/test_litellm/test_with_dashboard_node.py b/tests/unit/test_with_dashboard_node.py similarity index 100% rename from tests/test_litellm/test_with_dashboard_node.py rename to tests/unit/test_with_dashboard_node.py diff --git a/tests/test_litellm/test_xai_grok_4_3_model_metadata.py b/tests/unit/test_xai_grok_4_3_model_metadata.py similarity index 100% rename from tests/test_litellm/test_xai_grok_4_3_model_metadata.py rename to tests/unit/test_xai_grok_4_3_model_metadata.py diff --git a/tests/test_litellm/test_xai_responses_auto_routing.py b/tests/unit/test_xai_responses_auto_routing.py similarity index 100% rename from tests/test_litellm/test_xai_responses_auto_routing.py rename to tests/unit/test_xai_responses_auto_routing.py diff --git a/tests/test_litellm/types/test_completion.py b/tests/unit/types/test_completion.py similarity index 99% rename from tests/test_litellm/types/test_completion.py rename to tests/unit/types/test_completion.py index cd51913c5dd..4971a0c7e0a 100644 --- a/tests/test_litellm/types/test_completion.py +++ b/tests/unit/types/test_completion.py @@ -5,7 +5,7 @@ This test suite validates the CompletionRequest model and its compatibility with OpenAI ChatCompletion API message formats. Usage: - pytest tests/test_litellm/types/test_completion.py -v + pytest tests/unit/types/test_completion.py -v """ import dataclasses diff --git a/tests/test_litellm/types/test_guardrails_case_normalization.py b/tests/unit/types/test_guardrails_case_normalization.py similarity index 100% rename from tests/test_litellm/types/test_guardrails_case_normalization.py rename to tests/unit/types/test_guardrails_case_normalization.py diff --git a/tests/test_litellm/types/test_mcp.py b/tests/unit/types/test_mcp.py similarity index 100% rename from tests/test_litellm/types/test_mcp.py rename to tests/unit/types/test_mcp.py diff --git a/tests/test_litellm/types/test_presidio_entity_expansion.py b/tests/unit/types/test_presidio_entity_expansion.py similarity index 100% rename from tests/test_litellm/types/test_presidio_entity_expansion.py rename to tests/unit/types/test_presidio_entity_expansion.py diff --git a/tests/test_litellm/types/test_prometheus_label_value_sanitize.py b/tests/unit/types/test_prometheus_label_value_sanitize.py similarity index 100% rename from tests/test_litellm/types/test_prometheus_label_value_sanitize.py rename to tests/unit/types/test_prometheus_label_value_sanitize.py diff --git a/tests/test_litellm/types/test_prometheus_latency_buckets.py b/tests/unit/types/test_prometheus_latency_buckets.py similarity index 100% rename from tests/test_litellm/types/test_prometheus_latency_buckets.py rename to tests/unit/types/test_prometheus_latency_buckets.py diff --git a/tests/test_litellm/types/test_router.py b/tests/unit/types/test_router.py similarity index 100% rename from tests/test_litellm/types/test_router.py rename to tests/unit/types/test_router.py diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/unit/types/test_types_utils.py similarity index 100% rename from tests/test_litellm/types/test_types_utils.py rename to tests/unit/types/test_types_utils.py diff --git a/tests/test_litellm/types/test_uk_pii_entities.py b/tests/unit/types/test_uk_pii_entities.py similarity index 100% rename from tests/test_litellm/types/test_uk_pii_entities.py rename to tests/unit/types/test_uk_pii_entities.py diff --git a/tests/test_litellm/files/__init__.py b/tests/unit/vector_stores/__init__.py similarity index 100% rename from tests/test_litellm/files/__init__.py rename to tests/unit/vector_stores/__init__.py diff --git a/tests/test_litellm/vector_stores/test_main.py b/tests/unit/vector_stores/test_main.py similarity index 100% rename from tests/test_litellm/vector_stores/test_main.py rename to tests/unit/vector_stores/test_main.py diff --git a/tests/test_litellm/vector_stores/test_vector_store_create_provider_logic.py b/tests/unit/vector_stores/test_vector_store_create_provider_logic.py similarity index 100% rename from tests/test_litellm/vector_stores/test_vector_store_create_provider_logic.py rename to tests/unit/vector_stores/test_vector_store_create_provider_logic.py diff --git a/tests/test_litellm/vector_stores/test_vector_store_registry.py b/tests/unit/vector_stores/test_vector_store_registry.py similarity index 100% rename from tests/test_litellm/vector_stores/test_vector_store_registry.py rename to tests/unit/vector_stores/test_vector_store_registry.py