From 0ebec123fd465676f2fb62a1f8355631dec8b438 Mon Sep 17 00:00:00 2001
From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 13 Oct 2026 22:35:23 +0000
Subject: [PATCH] chore(typing): clear 670 basedpyright Any errors by
validating provider payloads at the boundary (#45388)
* chore(typing): validate provider payloads at the boundary to clear 464 basedpyright Any errors
Replaces untyped reads of decoded JSON in 48 provider, guardrail and proxy
modules with one pydantic validation at the point the payload enters the
module, so every value read afterwards carries a real type.
reportAny 6111 -> 5663 and reportExplicitAny 1416 -> 1400. No basedpyright,
ruff-strict or type-discipline rule is higher repo-wide or in any single
file, and no cast, ignore or suppression marker was added.
* chore(typing): replace 47 Any annotations that every caller already satisfies with object
Each annotation here is a parameter or return type whose value is only
stored, compared, logged or forwarded, so object describes it exactly and
no function body changes. A site was kept only if no basedpyright,
ruff-strict or type-discipline rule rose in any file that references it.
reportExplicitAny 1400 -> 1353 and reportAny 5663 -> 5655.
* chore(typing): type upstream responses and loose annotations across providers, proxy and integrations
Second pass of the Any reduction. Decoded upstream JSON is validated once at
the provider boundary and read through typed views afterwards, annotations
that every caller already satisfies move from Any to a concrete type, and the
nine search transformations validate each result item on its own so error
text no longer carries the item index.
No cast, no type: ignore, no suppression marker and no budget file edit
* ci(unit): run the brave, google_pse, nlp_cloud and opensandbox tests in the providers shard
---
.github/workflows/test-unit.yml | 4 +
.../providers/bedrock_agentcore/handler.py | 2 +-
litellm/batches/batch_utils.py | 2 +-
.../harness/handlers/deepagents_handler.py | 2 +-
litellm/integrations/argilla.py | 2 +-
.../arize/arize_phoenix_client.py | 4 +-
.../clickhouse/clickhouse_spend_logger.py | 2 +-
.../integrations/cloudzero/cz_stream_api.py | 2 +-
litellm/integrations/custom_guardrail.py | 2 +-
.../integrations/dotprompt/prompt_manager.py | 4 +-
litellm/integrations/focus/focus_logger.py | 2 +-
.../integrations/vantage/vantage_logger.py | 4 +-
litellm/llms/a2a/chat/transformation.py | 6 +-
litellm/llms/a2a/common_utils.py | 2 +-
.../aiml/image_generation/transformation.py | 33 +-
.../anthropic/count_tokens/token_counter.py | 8 +-
.../llms/anthropic/files/transformation.py | 9 +-
.../pass_through/adapters/transformation.py | 2 +-
.../anthropic/count_tokens/handler.py | 9 +-
.../anthropic/count_tokens/token_counter.py | 20 +-
.../image_edit/flux2_transformation.py | 2 +-
.../base_llm/managed_resources/isolation.py | 5 +-
.../amazon_qwen2_transformation.py | 21 +-
.../amazon_qwen3_transformation.py | 20 +-
.../base_invoke_transformation.py | 6 +-
litellm/llms/brave/search/transformation.py | 28 +-
litellm/llms/cohere/embed/transformation.py | 11 +-
litellm/llms/cometapi/embed/transformation.py | 10 +-
.../llms/dashscope/embed/transformation.py | 2 +-
.../llms/databricks/chat/transformation.py | 4 +-
.../llms/dataforseo/search/transformation.py | 24 +-
.../llms/deepagents/harness/transformation.py | 17 +-
.../llms/duckduckgo/search/transformation.py | 23 +-
litellm/llms/exa_ai/search/transformation.py | 31 +-
.../llms/fal_ai/image_edit/transformation.py | 9 +-
.../image_generation/bria_transformation.py | 9 +-
.../imagen4_transformation.py | 14 +-
.../fal_ai/image_generation/transformation.py | 10 +-
litellm/llms/fastcrw/search/transformation.py | 24 +-
.../llms/firecrawl/search/transformation.py | 63 ++--
.../llms/fireworks_ai/chat/transformation.py | 16 +-
.../gemini/image_generation/transformation.py | 29 +-
.../llms/google_pse/search/transformation.py | 24 +-
.../llms/infinity/embedding/transformation.py | 10 +-
litellm/llms/langgraph/chat/transformation.py | 4 +-
litellm/llms/lemonade/chat/transformation.py | 21 +-
litellm/llms/linkup/search/transformation.py | 40 ++-
.../mistral/audio_speech/transformation.py | 9 +-
litellm/llms/nlp_cloud/chat/transformation.py | 3 +-
litellm/llms/oci/chat/transformation.py | 2 +-
litellm/llms/ollama/common_utils.py | 6 +-
.../llms/ollama/completion/transformation.py | 5 +-
.../openai/responses/count_tokens/handler.py | 4 +-
.../responses/count_tokens/token_counter.py | 14 +-
.../responses/count_tokens/transformation.py | 2 +-
.../llms/openrouter/chat/transformation.py | 16 +-
.../opensandbox/sandbox/transformation.py | 9 +-
.../llms/ovhcloud/embedding/transformation.py | 10 +-
.../llms/perplexity/chat/transformation.py | 4 +-
.../perplexity/embedding/transformation.py | 17 +-
.../llms/perplexity/search/transformation.py | 24 +-
litellm/llms/replicate/chat/handler.py | 8 +-
litellm/llms/replicate/chat/transformation.py | 7 +-
.../sambanova/embedding/transformation.py | 10 +-
litellm/llms/sap/chat/transformation.py | 3 +-
.../llms/searchapi/search/transformation.py | 24 +-
litellm/llms/searxng/search/transformation.py | 24 +-
litellm/llms/serper/search/transformation.py | 24 +-
litellm/llms/tavily/search/transformation.py | 24 +-
.../llms/vertex_ai/files/transformation.py | 2 +-
litellm/llms/vertex_ai/fine_tuning/handler.py | 5 +-
.../vertex_imagen_transformation.py | 17 +-
.../vertex_gemini_transformation.py | 23 +-
.../vertex_imagen_transformation.py | 21 +-
.../text_to_speech/text_to_speech_handler.py | 13 +-
.../text_to_speech/transformation.py | 31 +-
.../vertex_embeddings/embedding_handler.py | 12 +-
.../llms/voyage/embedding/transformation.py | 10 +-
.../embedding/transformation_contextual.py | 9 +-
.../embedding/transformation_multimodal.py | 9 +-
.../llms/watsonx/completion/transformation.py | 4 +-
litellm/llms/xai/common_utils.py | 9 +-
.../guardrail_translation/handler.py | 6 +-
.../mcp_server/mcp_server_manager.py | 2 +-
.../mcp_server/openapi_to_mcp_generator.py | 2 +-
.../_experimental/mcp_server/tool_registry.py | 6 +-
.../proxy/agent_endpoints/a2a_endpoints.py | 2 +-
.../proxy/anthropic_endpoints/endpoints.py | 2 +-
litellm/proxy/auth/ip_address_utils.py | 2 +-
litellm/proxy/auth/model_checks.py | 6 +-
litellm/proxy/auth/trusted_proxy_utils.py | 6 +-
litellm/proxy/caching_routes.py | 4 +-
litellm/proxy/client/cli/commands/auth.py | 26 +-
litellm/proxy/client/cli/commands/http.py | 45 ++-
litellm/proxy/client/cli/commands/keys.py | 27 +-
litellm/proxy/client/cli/commands/models.py | 2 +-
litellm/proxy/client/cli/main.py | 13 +-
litellm/proxy/common_request_processing.py | 14 +-
.../proxy/common_utils/http_parsing_utils.py | 2 +-
.../common_utils/openapi_schema_compat.py | 4 +-
litellm/proxy/guardrails/_content_utils.py | 4 +-
.../guardrail_hooks/dynamoai/dynamoai.py | 6 +-
.../generic_guardrail_api.py | 4 +-
.../ibm_guardrails/ibm_detector.py | 4 +-
.../llm_shield_proxy/llm_shield_proxy.py | 11 +-
.../guardrails/guardrail_hooks/presidio.py | 1 +
.../semantic_guard/route_loader.py | 6 +-
.../shared_health_check_manager.py | 4 +-
litellm/proxy/litellm_pre_call_utils.py | 2 +-
litellm/proxy/mcp_tools.py | 7 +-
.../file_content_streaming_handler.py | 2 +-
.../anthropic_passthrough_logging_handler.py | 7 +-
litellm/proxy/plugin_routes.py | 4 +-
.../policy_engine/attachment_registry.py | 2 +-
litellm/proxy/policy_engine/init_policies.py | 8 +-
.../proxy/policy_engine/policy_registry.py | 15 +-
litellm/proxy/proxy_server.py | 8 +-
litellm/proxy/read_model_list.py | 4 +-
.../response_polling/background_streaming.py | 17 +-
.../proxy/response_polling/polling_handler.py | 25 +-
litellm/proxy/types_utils/utils.py | 12 +-
litellm/proxy/vector_store_endpoints/utils.py | 6 +-
litellm/proxy/video_endpoints/utils.py | 2 +-
litellm/rag/rag_query.py | 2 +-
litellm/router.py | 4 +-
.../secret_managers/google_secret_manager.py | 14 +-
litellm/utils.py | 4 +-
.../test_bedrock_agentcore_a2a.py | 50 +++
.../integrations/focus/test_focus_logger.py | 92 ++++++
tests/unit/integrations/test_argilla.py | 110 +++++++
tests/unit/integrations/vantage/__init__.py | 0
.../vantage/test_vantage_logger.py | 92 ++++++
.../a2a/chat/test_a2a_chat_transformation.py | 105 +++++++
...st_aiml_image_generation_transformation.py | 133 +++++++-
.../test_anthropic_files_transformation.py | 151 +++++++++
.../unit/llms/azure_ai/anthropic/__init__.py | 0
.../anthropic/count_tokens/__init__.py | 0
.../count_tokens/test_token_counter.py | 188 ++++++++++++
.../test_amazon_qwen2_transformation.py | 100 +++++-
.../test_amazon_qwen3_transformation.py | 85 ++++-
.../test_base_invoke_transformation.py | 21 ++
tests/unit/llms/brave/__init__.py | 0
tests/unit/llms/brave/search/__init__.py | 0
.../test_brave_search_transformation.py | 123 ++++++++
.../llms/cohere/embed/test_transformation.py | 139 +++++++++
tests/unit/llms/cometapi/embed/__init__.py | 0
.../test_cometapi_embed_transformation.py | 89 ++++++
.../test_dataforseo_search_transformation.py | 118 +++++++
.../test_duckduckgo_search_transformation.py | 133 ++++++++
.../llms/exa_ai/search/test_transformation.py | 109 +++++++
.../test_fal_ai_image_edit_transformation.py | 68 ++++
.../test_bria_transformation.py | 93 ++++++
..._fal_ai_image_generation_transformation.py | 97 ++++++
.../test_imagen4_transformation.py | 112 +++++++
.../fastcrw/search/test_transformation.py | 102 ++++++
.../test_firecrawl_search_transformation.py | 160 ++++++++++
.../test_fireworks_ai_chat_transformation.py | 118 ++++++-
..._gemini_image_generation_transformation.py | 116 +++++++
tests/unit/llms/google_pse/__init__.py | 0
tests/unit/llms/google_pse/search/__init__.py | 0
.../test_google_pse_search_transformation.py | 110 +++++++
.../unit/llms/infinity/embedding/__init__.py | 0
.../test_infinity_embedding_transformation.py | 89 ++++++
.../chat/test_langgraph_transformation.py | 36 +++
tests/unit/llms/lemonade/chat/__init__.py | 0
.../chat/test_lemonade_chat_transformation.py | 172 +++++++++++
.../test_linkup_search_transformation.py | 100 ++++++
...est_mistral_audio_speech_transformation.py | 112 +++++++
tests/unit/llms/nlp_cloud/__init__.py | 0
tests/unit/llms/nlp_cloud/chat/__init__.py | 0
.../test_nlp_cloud_chat_transformation.py | 47 +++
.../test_ollama_completion_transformation.py | 53 ++++
.../count_tokens/test_token_counter.py | 187 +++++++++++
.../test_openrouter_chat_transformation.py | 85 ++++-
tests/unit/llms/opensandbox/__init__.py | 0
.../unit/llms/opensandbox/sandbox/__init__.py | 0
.../sandbox/test_transformation.py | 153 +++++++++
...test_ovhcloud_embeddings_transformation.py | 89 ++++++
.../test_perplexity_chat_transformation.py | 167 ++++++++++
...est_perplexity_embedding_transformation.py | 100 ++++++
.../test_perplexity_search_transformation.py | 116 +++++++
.../unit/llms/replicate/chat/test_handler.py | 206 +++++++++++++
.../replicate/chat/test_transformation.py | 103 ++++++-
.../unit/llms/sambanova/embedding/__init__.py | 0
...test_sambanova_embedding_transformation.py | 89 ++++++
.../test_searchapi_search_transformation.py | 94 +++++-
.../test_searxng_search_transformation.py | 105 +++++++
.../test_serper_search_transformation.py | 83 +++++
.../test_tavily_search_transformation.py | 91 ++++++
.../vertex_ai/fine_tuning/test_handler.py | 106 +++++++
...est_vertex_ai_image_edit_transformation.py | 186 +++++++++++
...rtex_ai_image_generation_transformation.py | 290 +++++++++++++++++-
.../test_text_to_speech_handler.py | 155 ++++++++++
.../text_to_speech/test_transformation.py | 168 ++++++++++
.../vertex_ai/vertex_embeddings/__init__.py | 0
.../test_embedding_handler.py | 289 +++++++++++++++++
.../test_transformation_contextual.py | 89 ++++++
.../test_transformation_multimodal.py | 91 ++++++
.../test_voyage_embedding_transformation.py | 91 ++++++
.../unit/llms/watsonx/completion/__init__.py | 0
.../test_watsonx_completion_transformation.py | 60 ++++
tests/unit/llms/xai/test_xai_common_utils.py | 124 ++++++++
.../agent_endpoints/test_a2a_endpoints.py | 59 ++++
.../proxy/client/cli/test_auth_commands.py | 91 ++++++
.../proxy/client/cli/test_global_options.py | 103 +++++++
.../proxy/client/cli/test_keys_commands.py | 91 +++++-
tests/unit/proxy/client/test_http_commands.py | 119 +++++++
.../guardrail_hooks/dynamoai/test_dynamoai.py | 85 +++++
.../ibm_guardrails/__init__.py | 0
.../ibm_guardrails/test_ibm_detector.py | 124 ++++++++
...t_anthropic_passthrough_logging_handler.py | 46 +++
.../policy_engine/test_policy_registry.py | 131 ++++++++
.../test_response_polling_handler.py | 45 +++
tests/unit/proxy/test_caching_routes.py | 63 ++++
tests/unit/proxy/test_plugin_routes.py | 90 ++++++
tests/unit/proxy/types_utils/test_utils.py | 250 +++++++++++++++
.../test_vector_store_utils.py | 94 ++++++
.../test_google_secret_manager.py | 243 +++++++++++++++
218 files changed, 9258 insertions(+), 484 deletions(-)
create mode 100644 tests/unit/integrations/focus/test_focus_logger.py
create mode 100644 tests/unit/integrations/test_argilla.py
create mode 100644 tests/unit/integrations/vantage/__init__.py
create mode 100644 tests/unit/integrations/vantage/test_vantage_logger.py
create mode 100644 tests/unit/llms/azure_ai/anthropic/__init__.py
create mode 100644 tests/unit/llms/azure_ai/anthropic/count_tokens/__init__.py
create mode 100644 tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py
create mode 100644 tests/unit/llms/brave/__init__.py
create mode 100644 tests/unit/llms/brave/search/__init__.py
create mode 100644 tests/unit/llms/brave/search/test_brave_search_transformation.py
create mode 100644 tests/unit/llms/cometapi/embed/__init__.py
create mode 100644 tests/unit/llms/cometapi/embed/test_cometapi_embed_transformation.py
create mode 100644 tests/unit/llms/fal_ai/image_generation/test_bria_transformation.py
create mode 100644 tests/unit/llms/fal_ai/image_generation/test_fal_ai_image_generation_transformation.py
create mode 100644 tests/unit/llms/fal_ai/image_generation/test_imagen4_transformation.py
create mode 100644 tests/unit/llms/firecrawl/search/test_firecrawl_search_transformation.py
create mode 100644 tests/unit/llms/google_pse/__init__.py
create mode 100644 tests/unit/llms/google_pse/search/__init__.py
create mode 100644 tests/unit/llms/google_pse/search/test_google_pse_search_transformation.py
create mode 100644 tests/unit/llms/infinity/embedding/__init__.py
create mode 100644 tests/unit/llms/infinity/embedding/test_infinity_embedding_transformation.py
create mode 100644 tests/unit/llms/lemonade/chat/__init__.py
create mode 100644 tests/unit/llms/lemonade/chat/test_lemonade_chat_transformation.py
create mode 100644 tests/unit/llms/nlp_cloud/__init__.py
create mode 100644 tests/unit/llms/nlp_cloud/chat/__init__.py
create mode 100644 tests/unit/llms/nlp_cloud/chat/test_nlp_cloud_chat_transformation.py
create mode 100644 tests/unit/llms/openai/responses/count_tokens/test_token_counter.py
create mode 100644 tests/unit/llms/opensandbox/__init__.py
create mode 100644 tests/unit/llms/opensandbox/sandbox/__init__.py
create mode 100644 tests/unit/llms/opensandbox/sandbox/test_transformation.py
create mode 100644 tests/unit/llms/replicate/chat/test_handler.py
create mode 100644 tests/unit/llms/sambanova/embedding/__init__.py
create mode 100644 tests/unit/llms/sambanova/embedding/test_sambanova_embedding_transformation.py
create mode 100644 tests/unit/llms/vertex_ai/text_to_speech/test_text_to_speech_handler.py
create mode 100644 tests/unit/llms/vertex_ai/vertex_embeddings/__init__.py
create mode 100644 tests/unit/llms/vertex_ai/vertex_embeddings/test_embedding_handler.py
create mode 100644 tests/unit/llms/voyage/embedding/test_transformation_multimodal.py
create mode 100644 tests/unit/llms/voyage/embedding/test_voyage_embedding_transformation.py
create mode 100644 tests/unit/llms/watsonx/completion/__init__.py
create mode 100644 tests/unit/llms/watsonx/completion/test_watsonx_completion_transformation.py
create mode 100644 tests/unit/llms/xai/test_xai_common_utils.py
create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py
create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py
create mode 100644 tests/unit/proxy/types_utils/test_utils.py
create mode 100644 tests/unit/proxy/vector_store_endpoints/test_vector_store_utils.py
diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml
index e9236ee7787..0b448569095 100644
--- a/.github/workflows/test-unit.yml
+++ b/.github/workflows/test-unit.yml
@@ -233,6 +233,7 @@ jobs:
tests/unit/llms/base_llm
tests/unit/llms/baseten
tests/unit/llms/black_forest_labs
+ tests/unit/llms/brave
tests/unit/llms/bytez
tests/unit/llms/cerebras
tests/unit/llms/chat
@@ -268,6 +269,7 @@ jobs:
tests/unit/llms/gdc
tests/unit/llms/gemini
tests/unit/llms/gigachat
+ tests/unit/llms/google_pse
tests/unit/llms/github_copilot
tests/unit/llms/gradient_ai
tests/unit/llms/groq
@@ -300,6 +302,7 @@ jobs:
tests/unit/llms/nebius
tests/unit/llms/neosantara
tests/unit/llms/nimble
+ tests/unit/llms/nlp_cloud
tests/unit/llms/novita
tests/unit/llms/nscale
tests/unit/llms/nvidia_nim
@@ -310,6 +313,7 @@ jobs:
tests/unit/llms/oobabooga
tests/unit/llms/openai_like
tests/unit/llms/opencode
+ tests/unit/llms/opensandbox
tests/unit/llms/openrouter
tests/unit/llms/ovhcloud
tests/unit/llms/parallel_ai
diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py
index aecc887fbfe..e36d455a20f 100644
--- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py
+++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py
@@ -117,7 +117,7 @@ class BedrockAgentCoreA2AHandler:
response.raise_for_status()
# Check content type — AgentCore may return JSON instead of SSE
- content_type: Final = response.headers.get("content-type", "").lower()
+ content_type: Final = ", ".join(response.headers.get_list("content-type")).lower()
if "application/json" in content_type:
# Single JSON response fallback (not SSE)
diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py
index 6573f13be75..416391d194b 100644
--- a/litellm/batches/batch_utils.py
+++ b/litellm/batches/batch_utils.py
@@ -768,7 +768,7 @@ def _get_response_from_batch_job_output_file(
def _batch_response_was_successful(
- batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai"
+ batch_job_output_file: Mapping[str, object], custom_llm_provider: str = "openai"
) -> bool:
"""
Check if the batch job response was successful
diff --git a/litellm/harness/handlers/deepagents_handler.py b/litellm/harness/handlers/deepagents_handler.py
index e6ca7c161c7..5215562a1c5 100644
--- a/litellm/harness/handlers/deepagents_handler.py
+++ b/litellm/harness/handlers/deepagents_handler.py
@@ -158,7 +158,7 @@ class DeepAgentsHandler(BaseHarnessHandler):
agent, deps = self._require_agent()
snapshot = await agent.aget_state(self._run_config(ctx, None))
messages = (snapshot.values or MappingProxyType({})).get("messages") or ()
- converted: list[dict[str, Any]] = deps.convert_to_openai_messages( # mutable-ok: LangChain returns a list
+ converted: list[dict[str, object]] = deps.convert_to_openai_messages( # mutable-ok: LangChain returns a list
messages
)
return converted
diff --git a/litellm/integrations/argilla.py b/litellm/integrations/argilla.py
index 664ef8efda1..f873f274ab6 100644
--- a/litellm/integrations/argilla.py
+++ b/litellm/integrations/argilla.py
@@ -50,7 +50,7 @@ class ArgillaLogger(CustomBatchLogger):
if litellm.argilla_transformation_object is None:
raise Exception("'litellm.argilla_transformation_object' is required, to log your payload to Argilla.")
self.validate_argilla_transformation_object(litellm.argilla_transformation_object)
- self.argilla_transformation_object = litellm.argilla_transformation_object
+ self.argilla_transformation_object: Mapping[str, str] = litellm.argilla_transformation_object
self.default_credentials = self.get_credentials_from_env(
argilla_api_key=argilla_api_key,
argilla_dataset_name=argilla_dataset_name,
diff --git a/litellm/integrations/arize/arize_phoenix_client.py b/litellm/integrations/arize/arize_phoenix_client.py
index 18d35fee34b..b63cacede03 100644
--- a/litellm/integrations/arize/arize_phoenix_client.py
+++ b/litellm/integrations/arize/arize_phoenix_client.py
@@ -3,7 +3,7 @@ Arize Phoenix API client for fetching prompt versions from Arize Phoenix.
"""
import urllib.parse
-from typing import Any, Final
+from typing import Final
from litellm.llms.custom_httpx.http_handler import HTTPHandler
@@ -53,7 +53,7 @@ class ArizePhoenixClient:
# Initialize HTTPHandler
self.http_handler = HTTPHandler(disable_default_headers=True)
- def get_prompt_version(self, prompt_version_id: str) -> dict[str, Any] | None:
+ def get_prompt_version(self, prompt_version_id: str) -> dict[str, object] | None:
"""
Fetch a prompt version from Arize Phoenix.
diff --git a/litellm/integrations/clickhouse/clickhouse_spend_logger.py b/litellm/integrations/clickhouse/clickhouse_spend_logger.py
index 585ea6c2afc..21c5e828152 100644
--- a/litellm/integrations/clickhouse/clickhouse_spend_logger.py
+++ b/litellm/integrations/clickhouse/clickhouse_spend_logger.py
@@ -80,7 +80,7 @@ def _json(value: object) -> str:
return value if isinstance(value, str) else json.dumps(value, default=str)
-def _json_mapping(value: Mapping[str, Any]) -> str:
+def _json_mapping(value: Mapping[str, object]) -> str:
return _json(dict(value))
diff --git a/litellm/integrations/cloudzero/cz_stream_api.py b/litellm/integrations/cloudzero/cz_stream_api.py
index 81d08e6a4f4..1e0413e8576 100644
--- a/litellm/integrations/cloudzero/cz_stream_api.py
+++ b/litellm/integrations/cloudzero/cz_stream_api.py
@@ -209,7 +209,7 @@ class CloudZeroStreamer:
return payload
- def _convert_cbf_to_api_format(self, row: dict[str, object]) -> dict[str, Any] | None:
+ def _convert_cbf_to_api_format(self, row: dict[str, object]) -> dict[str, object] | None:
"""Convert CBF row to CloudZero API format - keeping CBF field names as CloudZero expects them."""
try:
# CloudZero expects CBF format field names directly, not converted names
diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py
index 45b20211ac5..09ca4857789 100644
--- a/litellm/integrations/custom_guardrail.py
+++ b/litellm/integrations/custom_guardrail.py
@@ -511,7 +511,7 @@ class CustomGuardrail(CustomLogger):
sticky_session_routing=self.sticky_session_routing,
)
- def _get_session_id_from_request_data(self, request_data: dict[str, Any]) -> str | None:
+ def _get_session_id_from_request_data(self, request_data: dict[str, object]) -> str | None:
"""Extract session_id from request data."""
return get_session_id_from_request_data(request_data)
diff --git a/litellm/integrations/dotprompt/prompt_manager.py b/litellm/integrations/dotprompt/prompt_manager.py
index 47c19b12b65..0eda266740c 100644
--- a/litellm/integrations/dotprompt/prompt_manager.py
+++ b/litellm/integrations/dotprompt/prompt_manager.py
@@ -70,7 +70,7 @@ class PromptManager:
self,
prompt_id: str | None = None,
prompt_directory: str | None = None,
- prompt_data: dict[str, dict[str, Any]] | None = None,
+ prompt_data: dict[str, dict[str, object]] | None = None,
prompt_file: str | None = None,
):
self.prompt_directory = Path(prompt_directory) if prompt_directory else None
@@ -372,6 +372,6 @@ class PromptManager:
}
return result
- def load_prompts_from_json_data(self, prompt_data: dict[str, dict[str, Any]]) -> None:
+ def load_prompts_from_json_data(self, prompt_data: dict[str, dict[str, object]]) -> None:
"""Load additional prompts from JSON data (merges with existing prompts)."""
self._load_prompts_from_json(prompt_data)
diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py
index dce51b8190b..449c7a2f9a3 100644
--- a/litellm/integrations/focus/focus_logger.py
+++ b/litellm/integrations/focus/focus_logger.py
@@ -115,7 +115,7 @@ class FocusLogger(CustomLogger):
pod_lock_manager: PodLockManager | None = None
if proxy_logging_obj is not None:
- writer: Final = getattr(proxy_logging_obj, "db_spend_update_writer", None)
+ writer: Final[object] = getattr(proxy_logging_obj, "db_spend_update_writer", None)
if writer is not None:
pod_lock_manager = getattr(writer, "pod_lock_manager", None)
diff --git a/litellm/integrations/vantage/vantage_logger.py b/litellm/integrations/vantage/vantage_logger.py
index 48a492fdd72..79beac78fcc 100644
--- a/litellm/integrations/vantage/vantage_logger.py
+++ b/litellm/integrations/vantage/vantage_logger.py
@@ -16,6 +16,8 @@ from litellm.integrations.focus.focus_logger import FocusLogger
if TYPE_CHECKING:
from apscheduler.schedulers.asyncio import AsyncIOScheduler
+
+ from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
else:
AsyncIOScheduler = Any
@@ -91,7 +93,7 @@ class VantageLogger(FocusLogger):
"""
from litellm.proxy.proxy_server import proxy_logging_obj
- pod_lock_manager = None
+ pod_lock_manager: PodLockManager | None = None
if proxy_logging_obj is not None:
writer: Final[object] = getattr(proxy_logging_obj, "db_spend_update_writer", None)
if writer is not None:
diff --git a/litellm/llms/a2a/chat/transformation.py b/litellm/llms/a2a/chat/transformation.py
index 6171b341ca4..be3b5fea601 100644
--- a/litellm/llms/a2a/chat/transformation.py
+++ b/litellm/llms/a2a/chat/transformation.py
@@ -7,6 +7,7 @@ from collections.abc import AsyncIterator, Iterator, Mapping
from typing import TYPE_CHECKING, Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.llms.azure_ai.common_utils import AZURE_ENTRA_LITELLM_PARAM_KEYS, get_azure_ai_agent_entra_token
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
@@ -30,6 +31,7 @@ if TYPE_CHECKING:
_REGISTRY_PARAMS_KEPT_OUT_OF_OPTIONAL_PARAMS: Final = (
frozenset({"api_key", "api_base", "headers", "model"}) | AZURE_ENTRA_LITELLM_PARAM_KEYS
)
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
def _card_declares_no_streaming(agent_card_params: Mapping[str, object]) -> bool:
@@ -325,7 +327,7 @@ class A2AConfig(BaseConfig):
# Check for JSON-RPC error
if "error" in response_json:
- error: Final = response_json["error"]
+ error: Final = _JSON_OBJECT.validate_python(response_json["error"])
raise A2AError(
status_code=raw_response.status_code,
message=f"A2A error: {error.get('message', 'Unknown error')}",
@@ -333,7 +335,7 @@ class A2AConfig(BaseConfig):
)
# Extract text from A2A response
- text: Final = extract_text_from_a2a_response(response_json)
+ text: Final = extract_text_from_a2a_response(_JSON_OBJECT.validate_python(response_json))
# Populate model response
model_response.choices = [
diff --git a/litellm/llms/a2a/common_utils.py b/litellm/llms/a2a/common_utils.py
index 030c5bc222e..1818b406823 100644
--- a/litellm/llms/a2a/common_utils.py
+++ b/litellm/llms/a2a/common_utils.py
@@ -22,7 +22,7 @@ class A2AError(BaseLLMException):
self,
status_code: int,
message: str,
- headers: dict[str, Any] = {},
+ headers: dict[str, str] = {},
):
super().__init__(
status_code=status_code,
diff --git a/litellm/llms/aiml/image_generation/transformation.py b/litellm/llms/aiml/image_generation/transformation.py
index b585d35a2ac..931038689cc 100644
--- a/litellm/llms/aiml/image_generation/transformation.py
+++ b/litellm/llms/aiml/image_generation/transformation.py
@@ -1,6 +1,8 @@
+from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
@@ -21,6 +23,10 @@ if TYPE_CHECKING:
else:
LiteLLMLoggingObj = Any
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object]]] = TypeAdapter(
+ Mapping[str, object] | str | Sequence[object], config=ConfigDict(hide_input_in_errors=True)
+)
OPENAI_STYLE_IMAGE_MODEL_PREFIXES: Final[tuple[str, ...]] = ("openai/",)
@@ -198,29 +204,33 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig):
if "data" in response_data and isinstance(response_data["data"], list):
# Handle OpenAI-like format: {"data": [{"url": "...", "width": 1024, "height": 768, "content_type": "image/jpeg"}]}
- for image in response_data["data"]:
+ for raw_image in _JSON_CONTAINER.validate_python(response_data["data"]):
+ image = raw_image if isinstance(raw_image, str) else _JSON_CONTAINER.validate_python(raw_image)
if "url" in image:
+ url_image = _JSON_OBJECT.validate_python(image)
model_response.data.append(
ImageObject(
b64_json=None,
- url=image["url"],
- revised_prompt=image.get("revised_prompt"),
+ url=url_image["url"],
+ revised_prompt=url_image.get("revised_prompt"),
)
)
elif "b64_json" in image or "image_base64" in image:
+ b64_image = _JSON_OBJECT.validate_python(image)
model_response.data.append(
ImageObject(
- b64_json=image.get("b64_json") or image.get("image_base64"),
+ b64_json=b64_image.get("b64_json") or b64_image.get("image_base64"),
url=None,
- revised_prompt=image.get("revised_prompt"),
+ revised_prompt=b64_image.get("revised_prompt"),
)
)
elif "output" in response_data and "choices" in response_data["output"]:
- for choice in response_data["output"]["choices"]:
+ for raw_choice in _JSON_CONTAINER.validate_python(response_data["output"]["choices"]):
+ choice = raw_choice if isinstance(raw_choice, str) else _JSON_CONTAINER.validate_python(raw_choice)
if "image_base64" in choice:
model_response.data.append(
ImageObject(
- b64_json=choice["image_base64"],
+ b64_json=_JSON_OBJECT.validate_python(choice)["image_base64"],
url=None,
)
)
@@ -228,23 +238,24 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig):
model_response.data.append(
ImageObject(
b64_json=None,
- url=choice["url"],
+ url=_JSON_OBJECT.validate_python(choice)["url"],
)
)
elif "images" in response_data:
# Handle alternative format: {"images": [{"url": "...", "width": 1024, "height": 768, "content_type": "image/jpeg"}]}
- for image in response_data["images"]:
+ for raw_image in _JSON_CONTAINER.validate_python(response_data["images"]):
+ image = raw_image if isinstance(raw_image, str) else _JSON_CONTAINER.validate_python(raw_image)
if "url" in image:
model_response.data.append(
ImageObject(
b64_json=None,
- url=image["url"],
+ url=_JSON_OBJECT.validate_python(image)["url"],
)
)
elif "image_base64" in image:
model_response.data.append(
ImageObject(
- b64_json=image["image_base64"],
+ b64_json=_JSON_OBJECT.validate_python(image)["image_base64"],
url=None,
)
)
diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py
index 401c09b580b..a45d6c1184f 100644
--- a/litellm/llms/anthropic/count_tokens/token_counter.py
+++ b/litellm/llms/anthropic/count_tokens/token_counter.py
@@ -4,6 +4,8 @@ Anthropic Token Counter implementation using the CountTokens API.
from typing import Any, Final
+from pydantic import JsonValue
+
from litellm._logging import verbose_logger
from litellm.exceptions import AuthenticationError
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
@@ -26,12 +28,12 @@ class AnthropicTokenCounter(BaseTokenCounter):
async def count_tokens(
self,
model_to_use: str,
- messages: list[dict[str, Any]] | None,
+ messages: list[dict[str, JsonValue]] | None,
contents: list[dict[str, object]] | None,
deployment: dict[str, Any] | None = None,
request_model: str = "",
- tools: list[dict[str, Any]] | None = None,
- system: Any | None = None,
+ tools: list[dict[str, JsonValue]] | None = None,
+ system: JsonValue = None,
) -> TokenCountResponse | None:
"""
Count tokens using Anthropic's CountTokens API.
diff --git a/litellm/llms/anthropic/files/transformation.py b/litellm/llms/anthropic/files/transformation.py
index 082902e8897..8433235dde4 100644
--- a/litellm/llms/anthropic/files/transformation.py
+++ b/litellm/llms/anthropic/files/transformation.py
@@ -14,7 +14,7 @@ Anthropic Files API endpoints:
import calendar
import time
-from collections.abc import Mapping
+from collections.abc import Iterable, Mapping
from typing import Final, cast
import httpx
@@ -49,6 +49,7 @@ ANTHROPIC_FILES_BETA_HEADER: Final = "files-api-2025-04-14"
ANTHROPIC_MESSAGE_BATCH_ID_PREFIX: Final = "msgbatch_"
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_ITEMS: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
class AnthropicFilesConfig(BaseFilesConfig):
@@ -284,9 +285,9 @@ class AnthropicFilesConfig(BaseFilesConfig):
"last_id": "..."
}
"""
- response_json: Final = raw_response.json()
- files_data: Final = response_json.get("data", [])
- return [self._parse_anthropic_file(f) for f in files_data]
+ response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
+ files_data: Final = _JSON_ITEMS.validate_python(response_json.get("data", []))
+ return [self._parse_anthropic_file(_JSON_OBJECT.validate_python(f)) for f in files_data]
def transform_file_content_request(
self,
diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py
index 2ee23852894..d2a0849ffef 100644
--- a/litellm/llms/anthropic/pass_through/adapters/transformation.py
+++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py
@@ -1347,7 +1347,7 @@ class LiteLLMAnthropicMessagesAdapter:
self,
choices: list[Choices],
tool_name_mapping: dict[str, str] | None = None,
- ) -> list[dict[str, Any]]:
+ ) -> list[dict[str, object]]:
new_content: Final[list[dict[str, Any]]] = []
for choice, compaction_blocks in (
(choice, _compaction_blocks(_optional_attr(choice.message, "provider_specific_fields")))
diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py
index 3270fb3534a..0e0b68df449 100644
--- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py
+++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py
@@ -4,9 +4,10 @@ Azure AI Anthropic CountTokens API handler.
Uses httpx for HTTP requests with Azure authentication.
"""
-from typing import Any, Final
+from typing import Final
import httpx
+from pydantic import JsonValue
import litellm
from litellm._logging import verbose_logger
@@ -28,14 +29,14 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
async def handle_count_tokens_request(
self,
model: str,
- messages: list[dict[str, Any]],
+ messages: list[dict[str, JsonValue]],
api_key: str,
api_base: str,
litellm_params: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
- tools: list[dict[str, Any]] | None = None,
+ tools: list[dict[str, JsonValue]] | None = None,
system: object = None,
- ) -> dict[str, Any]:
+ ) -> dict[str, object]:
"""
Handle a CountTokens request using httpx with Azure authentication.
diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py
index c31d2c427bb..1d6d59c0592 100644
--- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py
+++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py
@@ -5,6 +5,8 @@ Azure AI Anthropic Token Counter implementation using the CountTokens API.
import os
from typing import Any, Final
+from pydantic import JsonValue
+
from litellm._logging import verbose_logger
from litellm.llms.azure_ai.anthropic.count_tokens.handler import (
AzureAIAnthropicCountTokensHandler,
@@ -28,11 +30,11 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
async def count_tokens(
self,
model_to_use: str,
- messages: list[dict[str, object]] | None,
+ messages: list[dict[str, JsonValue]] | None,
contents: list[dict[str, object]] | None,
deployment: dict[str, Any] | None = None,
request_model: str = "",
- tools: list[dict[str, object]] | None = None,
+ tools: list[dict[str, JsonValue]] | None = None,
system: object | None = None,
) -> TokenCountResponse | None:
"""
@@ -86,12 +88,14 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
)
if result is not None:
- return TokenCountResponse(
- total_tokens=result.get("input_tokens", 0),
- request_model=request_model,
- model_used=model_to_use,
- tokenizer_type="azure_ai_anthropic_api",
- original_response=result,
+ return TokenCountResponse.model_validate(
+ {
+ "total_tokens": result.get("input_tokens", 0),
+ "request_model": request_model,
+ "model_used": model_to_use,
+ "tokenizer_type": "azure_ai_anthropic_api",
+ "original_response": result,
+ }
)
except AnthropicError as e:
verbose_logger.warning(
diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py
index f91a87ba0f4..9bd476efdff 100644
--- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py
+++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py
@@ -113,7 +113,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
for index, reference_image in enumerate(images, start=1)
}
)
- request_body: Final[dict[str, Any]] = {
+ request_body: Final[dict[str, object]] = {
"prompt": prompt,
"model": model,
**reference_images,
diff --git a/litellm/llms/base_llm/managed_resources/isolation.py b/litellm/llms/base_llm/managed_resources/isolation.py
index 547bed21247..09ca23a8ded 100644
--- a/litellm/llms/base_llm/managed_resources/isolation.py
+++ b/litellm/llms/base_llm/managed_resources/isolation.py
@@ -10,6 +10,7 @@ they created. Callers with no admin role and no identifying ids at all
are denied so an empty user_id can never select an unscoped query.
"""
+from collections.abc import Mapping, Sequence
from typing import Any, Final
from litellm.types.proxy.auth.user_api_key_auth import (
@@ -46,7 +47,7 @@ def resolve_resource_owner_id(
return None
-def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, Any]:
+def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, object]:
"""Build the OpenAI-style paginated list response shape used by managed
file/batch/vector-store listings. ``first_id`` and ``last_id`` are
sourced from each item's ``.id`` attribute."""
@@ -61,7 +62,7 @@ def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, Any]:
def build_owner_filter(
user_api_key_dict: UserAPIKeyAuth,
-) -> dict[str, Any] | None:
+) -> dict[str, str | Sequence[Mapping[str, str]]] | None:
"""Return a Prisma `where` fragment that scopes a managed-resource listing
to records the caller is allowed to see.
diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py
index 67364ccfda0..9ee5b264935 100644
--- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py
+++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py
@@ -7,9 +7,11 @@ The main difference is in the response format: Qwen2 uses "text" field while Qwe
Qwen2 + Invoke API Tutorial: https://docs.aws.amazon.com/bedrock/latest/userguide/invoke-imported-model.html
"""
+from collections.abc import Mapping
from typing import TYPE_CHECKING, Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformation import (
AmazonQwen3Config,
@@ -23,6 +25,9 @@ from litellm.types.utils import ModelResponse, Usage
if TYPE_CHECKING:
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_TEXT: Final = TypeAdapter(str, config=ConfigDict(strict=True, hide_input_in_errors=True))
+
class AmazonQwen2Config(AmazonQwen3Config):
"""
@@ -54,10 +59,10 @@ class AmazonQwen2Config(AmazonQwen3Config):
Qwen2 uses "text" field, but we also support "generation" field for compatibility.
"""
try:
- response_data: Final = raw_response.json()
+ response_data: Final = _JSON_OBJECT.validate_python(raw_response.json())
# Extract the generated text - Qwen2 uses "text" field, but also support "generation" for compatibility
- generated_text = response_data.get("generation", "") or response_data.get("text", "")
+ generated_text = _TEXT.validate_python(response_data.get("generation", "") or response_data.get("text", ""))
# Clean up the response (remove assistant start token if present)
generated_text = generated_text.removeprefix("<|im_start|>assistant\n")
@@ -71,14 +76,16 @@ class AmazonQwen2Config(AmazonQwen3Config):
# Set usage information if available in response
if "usage" in response_data:
- usage_data: Final = response_data["usage"]
+ usage_data: Final = _JSON_OBJECT.validate_python(response_data["usage"])
setattr(
model_response,
"usage",
- Usage(
- prompt_tokens=usage_data.get("prompt_tokens", 0),
- completion_tokens=usage_data.get("completion_tokens", 0),
- total_tokens=usage_data.get("total_tokens", 0),
+ Usage.model_validate(
+ {
+ "prompt_tokens": usage_data.get("prompt_tokens", 0),
+ "completion_tokens": usage_data.get("completion_tokens", 0),
+ "total_tokens": usage_data.get("total_tokens", 0),
+ }
),
)
diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py
index 04248a785dd..1e31f0292c0 100644
--- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py
+++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py
@@ -10,6 +10,7 @@ from collections.abc import Mapping
from typing import TYPE_CHECKING, Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.llms.base_llm.chat.transformation import BaseConfig
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
@@ -22,6 +23,9 @@ from litellm.types.utils import ModelResponse, Usage
if TYPE_CHECKING:
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_TEXT: Final = TypeAdapter(str, config=ConfigDict(strict=True, hide_input_in_errors=True))
+
class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig):
"""
@@ -179,10 +183,10 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig):
Transform Qwen3 Bedrock response to OpenAI format
"""
try:
- response_data: Final = raw_response.json()
+ response_data: Final = _JSON_OBJECT.validate_python(raw_response.json())
# Extract the generated text - Qwen3 uses "generation" field
- generated_text = response_data.get("generation", "")
+ generated_text = _TEXT.validate_python(response_data.get("generation", ""))
# Clean up the response (remove assistant start token if present)
generated_text = generated_text.removeprefix("<|im_start|>assistant\n")
@@ -196,14 +200,16 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig):
# Set usage information if available in response
if "usage" in response_data:
- usage_data: Final = response_data["usage"]
+ usage_data: Final = _JSON_OBJECT.validate_python(response_data["usage"])
setattr(
model_response,
"usage",
- Usage(
- prompt_tokens=usage_data.get("prompt_tokens", 0),
- completion_tokens=usage_data.get("completion_tokens", 0),
- total_tokens=usage_data.get("total_tokens", 0),
+ Usage.model_validate(
+ {
+ "prompt_tokens": usage_data.get("prompt_tokens", 0),
+ "completion_tokens": usage_data.get("completion_tokens", 0),
+ "total_tokens": usage_data.get("total_tokens", 0),
+ }
),
)
diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
index af833d93084..2356efe0f48 100644
--- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
+++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
@@ -51,6 +51,10 @@ _GUARDRAIL_CONFIG_EXPECTED_FORMAT: Final = (
"{'guardrailIdentifier': str, 'guardrailVersion': str, 'trace': 'enabled'|'disabled'|'enabled_full'}"
)
+_INVOKE_PROVIDERS: Final = TypeAdapter(tuple[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, ...]).validate_python(
+ get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL)
+)
+
def _bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]":
try:
@@ -568,7 +572,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
if provider is not None:
return provider
- for provider in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL):
+ for provider in _INVOKE_PROVIDERS:
if provider in model:
return provider
return None
diff --git a/litellm/llms/brave/search/transformation.py b/litellm/llms/brave/search/transformation.py
index 1f78f8d154c..459bc3193cb 100644
--- a/litellm/llms/brave/search/transformation.py
+++ b/litellm/llms/brave/search/transformation.py
@@ -6,15 +6,20 @@ Documentation: https://api-dashboard.search.brave.com/app/documentation/web-sear
from __future__ import annotations
import re
+from collections.abc import Iterable, Mapping
from datetime import datetime, timezone
from typing import Final, Literal, TypedDict
import httpx
from dateutil import parser
+from pydantic import ConfigDict, TypeAdapter
_ISO_YMD: Final = re.compile(r"^\s*\d{4}[-/]\d{1,2}[-/]\d{1,2}\s*$")
_UNIX_TIMESTAMP: Final = re.compile(r"^\s*-?\d+(\.\d+)?\s*$")
BRAVE_SECTIONS: Final = ["web", "discussions", "faqs", "faq", "news", "videos"]
+_JSON_VALUE: Final = TypeAdapter(object)
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.search.transformation import (
@@ -26,7 +31,7 @@ from litellm.secret_managers.main import get_secret_str
def to_yyyy_mm_dd(
- s: str | float | None,
+ s: object,
*,
dayfirst: bool = False,
yearfirst: bool = False,
@@ -246,7 +251,7 @@ class BraveSearchConfig(BaseSearchConfig):
"""
Transform Brave Search API response to LiteLLM unified SearchResponse format.
"""
- response_json: Final = raw_response.json()
+ response_body: Final = _JSON_VALUE.validate_python(raw_response.json())
# Transform results to SearchResult objects
results: Final[list[SearchResult]] = []
@@ -254,9 +259,11 @@ class BraveSearchConfig(BaseSearchConfig):
query_params: Final = raw_response.request.url.params if raw_response.request else {}
sections_to_process: Final = self._sections_from_params(dict(query_params))
max_results: Final = max(1, min(int(query_params.get("count", 20)), 20))
+ response_json: Final = _JSON_OBJECT.validate_python(response_body)
for section in sections_to_process:
- for result in response_json.get(section, {}).get("results", []):
+ section_json = _JSON_OBJECT.validate_python(response_json.get(section, {}))
+ for raw_result in _JSON_VALUES.validate_python(section_json.get("results", [])):
# Because the `max_results`/`count` parameters do not affect
# the number of "discussion", "faq", "news", or "videos"
# results, we need to manually limit the number of results
@@ -264,18 +271,21 @@ class BraveSearchConfig(BaseSearchConfig):
if len(results) >= max_results:
break
+ result = _JSON_OBJECT.validate_python(raw_result)
title = result.get("title", "")
url = result.get("url", "")
snippet = result.get("description", "")
date = to_yyyy_mm_dd(result.get("page_age") or result.get("age"))
last_updated = to_yyyy_mm_dd(result.get("fetched_content_timestamp", ""))
- search_result = SearchResult(
- title=title,
- url=url,
- snippet=snippet,
- date=date,
- last_updated=last_updated,
+ search_result = SearchResult.model_validate(
+ {
+ "title": title,
+ "url": url,
+ "snippet": snippet,
+ "date": date,
+ "last_updated": last_updated,
+ }
)
results.append(search_result)
diff --git a/litellm/llms/cohere/embed/transformation.py b/litellm/llms/cohere/embed/transformation.py
index bac899c4142..73b58254a3d 100644
--- a/litellm/llms/cohere/embed/transformation.py
+++ b/litellm/llms/cohere/embed/transformation.py
@@ -10,10 +10,11 @@ Convers
Docs - https://docs.cohere.com/v2/reference/embed
"""
-from collections.abc import Sized
+from collections.abc import Iterable, Mapping, Sized
from typing import Final, Protocol, cast
import httpx
+from pydantic import ConfigDict, TypeAdapter
import litellm
from litellm import COHERE_DEFAULT_EMBEDDING_INPUT_TYPE
@@ -30,6 +31,9 @@ from litellm.utils import is_base64_encoded
from ..common_utils import CohereError
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_EMBEDDINGS_BY_TYPE: Final = TypeAdapter(Mapping[str, Iterable[object]], config=ConfigDict(hide_input_in_errors=True))
+
class _SupportsEncode(Protocol):
"""Tokenizer handle: the embedding usage path only encodes text to measure its token length."""
@@ -198,7 +202,8 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig):
'usage'
}
"""
- embeddings: Final = response_json["embeddings"]
+ payload: Final = _JSON_OBJECT.validate_python(response_json)
+ embeddings: Final = _EMBEDDINGS_BY_TYPE.validate_python(payload["embeddings"])
output_data: Final = []
for k, embedding_list in embeddings.items():
for idx, embedding in enumerate(embedding_list):
@@ -213,7 +218,7 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig):
setattr(
model_response,
"usage",
- self._calculate_usage(input, encoding, response_json.get("meta", {})),
+ self._calculate_usage(input, encoding, dict(_JSON_OBJECT.validate_python(payload.get("meta", {})))),
)
return model_response
diff --git a/litellm/llms/cometapi/embed/transformation.py b/litellm/llms/cometapi/embed/transformation.py
index adcecd5c6c9..60f781be261 100644
--- a/litellm/llms/cometapi/embed/transformation.py
+++ b/litellm/llms/cometapi/embed/transformation.py
@@ -2,9 +2,11 @@
CometAPI Embedding API support - OpenAI compatible
"""
+from collections.abc import Mapping
from typing import Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
@@ -15,6 +17,9 @@ from litellm.types.utils import EmbeddingResponse, Usage
from ..common_utils import CometAPIException
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_INT: Final = TypeAdapter(int)
+
class CometAPIEmbeddingConfig(BaseEmbeddingConfig):
"""
@@ -136,9 +141,10 @@ class CometAPIEmbeddingConfig(BaseEmbeddingConfig):
model_response.data = raw_response_json.get("data")
model_response.object = raw_response_json.get("object")
+ usage_payload: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {}))
usage: Final = Usage(
- prompt_tokens=raw_response_json.get("usage", {}).get("prompt_tokens", 0),
- total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0),
+ prompt_tokens=_INT.validate_python(usage_payload.get("prompt_tokens", 0) or 0),
+ total_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0),
)
model_response.usage = usage
diff --git a/litellm/llms/dashscope/embed/transformation.py b/litellm/llms/dashscope/embed/transformation.py
index 63ee984a65c..e5a987ca1ad 100644
--- a/litellm/llms/dashscope/embed/transformation.py
+++ b/litellm/llms/dashscope/embed/transformation.py
@@ -147,7 +147,7 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig):
)
if "error" in response_json:
- error: Final = response_json["error"]
+ error: Final[object] = response_json["error"]
message: Final = error.get("message", str(error)) if isinstance(error, dict) else str(error)
raise DashScopeError(
status_code=raw_response.status_code,
diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py
index 94a74d703df..4b6f178fa1a 100644
--- a/litellm/llms/databricks/chat/transformation.py
+++ b/litellm/llms/databricks/chat/transformation.py
@@ -475,7 +475,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
# Move message-level cache_control into a content block when content is a string.
if "cache_control" in _message and isinstance(_message.get("content"), str):
_message = self._move_cache_control_into_string_content_block(_message)
- _sanitize_empty_content(cast(dict[str, Any], _message))
+ _sanitize_empty_content(cast(dict[str, object], _message))
if _is_bare_assistant_message(_message):
continue
new_messages.append(_message)
@@ -504,7 +504,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
"""
content: Final = message.get("content")
# Create new message with cache_control moved into content block
- transformed_message: Final = cast(dict[str, Any], message.copy())
+ transformed_message: Final = cast(dict[str, object], message.copy())
cache_control: Final = transformed_message.pop("cache_control")
transformed_message["content"] = [
{
diff --git a/litellm/llms/dataforseo/search/transformation.py b/litellm/llms/dataforseo/search/transformation.py
index fcd4ae70645..5cf7735325d 100644
--- a/litellm/llms/dataforseo/search/transformation.py
+++ b/litellm/llms/dataforseo/search/transformation.py
@@ -4,9 +4,11 @@ Calls DataForSEO SERP API to search the web.
DataForSEO API Reference: https://docs.dataforseo.com/v3/serp/google/organic/live/advanced/?bash
"""
+from collections.abc import Iterable, Mapping
from typing import Final, Literal
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.constants import DEFAULT_DATAFORSEO_LOCATION_CODE
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@@ -17,6 +19,9 @@ from litellm.llms.base_llm.search.transformation import (
)
from litellm.secret_managers.main import get_secret_str
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
+
class DataForSEOSearchConfig(BaseSearchConfig):
"""
@@ -201,18 +206,21 @@ class DataForSEOSearchConfig(BaseSearchConfig):
if task.get("status_code") == 20000 and "result" in task:
# Result is an array, take first element
if len(task["result"]) > 0:
- result: Final = task["result"][0]
+ result: Final = _JSON_OBJECT.validate_python(task["result"][0])
# Items contain the actual search results
- for item in result.get("items", []):
+ for raw_item in _JSON_VALUES.validate_python(result.get("items", [])):
+ item = _JSON_OBJECT.validate_python(raw_item)
# Only process organic search results
if item.get("type") == "organic":
- search_result = SearchResult(
- title=item.get("title", ""),
- url=item.get("url", ""),
- snippet=item.get("description", ""),
- date=None, # DataForSEO doesn't provide date in standard response
- last_updated=None,
+ search_result = SearchResult.model_validate(
+ {
+ "title": item.get("title", ""),
+ "url": item.get("url", ""),
+ "snippet": item.get("description", ""),
+ "date": None, # DataForSEO doesn't provide date in standard response
+ "last_updated": None,
+ }
)
results.append(search_result)
diff --git a/litellm/llms/deepagents/harness/transformation.py b/litellm/llms/deepagents/harness/transformation.py
index 1f8728230bc..cdac6944019 100644
--- a/litellm/llms/deepagents/harness/transformation.py
+++ b/litellm/llms/deepagents/harness/transformation.py
@@ -14,7 +14,7 @@ import json
from collections.abc import Iterator, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final
+from typing import TYPE_CHECKING, Any, Final, Protocol
from litellm.harness.options import DeepAgentsOptions
from litellm.harness.types import (
@@ -73,7 +73,7 @@ _APPROVAL_DECISIONS: Final = ("approve", "reject")
def chat_model_kwargs(
ctx: SessionContext,
-) -> dict[str, Any]: # mutable-ok: ChatLiteLLM constructor kwargs, splatted as **kwargs
+) -> dict[str, object]: # mutable-ok: ChatLiteLLM constructor kwargs, splatted as **kwargs
"""ChatLiteLLM constructor kwargs for gateway or SDK mode."""
if not ctx.model:
raise ValueError("Harness.DEEPAGENTS needs model=")
@@ -107,7 +107,7 @@ def blocked_tools(permissions: str, disable_tools: Sequence[str]) -> frozenset[s
def interrupt_config(
permissions: str, blocked: frozenset[str]
-) -> dict[str, Any] | None: # mutable-ok: deepagents create_deep_agent(interrupt_on=) takes a dict
+) -> dict[str, dict[str, list[str]]] | None: # mutable-ok: deepagents create_deep_agent(interrupt_on=) takes a dict
"""interrupt_on for permissions='ask': approve/reject every mutating built-in."""
if permissions != "ask":
return None
@@ -221,9 +221,14 @@ def _message_events(message: object, skip_tools: frozenset[str]) -> tuple[Event,
return ()
+class InterruptLike(Protocol):
+ @property
+ def id(self) -> str: ...
+
+
def interrupts_in(
update: object,
-) -> list[Any]: # mutable-ok: returns a list; existing callers/tests compare it to list literals
+) -> list[InterruptLike]: # mutable-ok: returns a list; existing callers/tests compare it to list literals
if not isinstance(update, Mapping):
return [] # mutable-ok: returns a list; existing callers/tests compare it to list literals
found = update.get("__interrupt__")
@@ -260,7 +265,7 @@ def approval_requests(
return list(kept) # mutable-ok: list return; callers/tests compare to lists
-def decision(allowed: bool, reason: str) -> dict[str, Any]: # mutable-ok: LangGraph resume payload (HITL decision dict)
+def decision(allowed: bool, reason: str) -> dict[str, str]: # mutable-ok: LangGraph resume payload (HITL decision dict)
if allowed:
return {"type": "approve"} # mutable-ok: LangGraph resume payload (HITL decision dict)
return { # mutable-ok: LangGraph HITL decision
@@ -273,7 +278,7 @@ def decision(allowed: bool, reason: str) -> dict[str, Any]: # mutable-ok: LangG
class TurnState:
"""Mutable state across the stream passes of one turn."""
- interrupts: tuple[Any, ...] = ()
+ interrupts: tuple[InterruptLike, ...] = ()
class DeepAgentsHarnessConfig(BaseHarnessConfig):
diff --git a/litellm/llms/duckduckgo/search/transformation.py b/litellm/llms/duckduckgo/search/transformation.py
index 93e70f65bcd..1816a0d24e4 100644
--- a/litellm/llms/duckduckgo/search/transformation.py
+++ b/litellm/llms/duckduckgo/search/transformation.py
@@ -4,10 +4,12 @@ Calls DuckDuckGo's Instant Answer API to search the web.
DuckDuckGo API Reference: https://duckduckgo.com/api
"""
+from collections.abc import Iterable, Mapping
from typing import Final, Literal, TypedDict
from urllib.parse import urlencode
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.search.transformation import (
@@ -17,6 +19,9 @@ from litellm.llms.base_llm.search.transformation import (
)
from litellm.secret_managers.main import get_secret_str
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
+
class _DuckDuckGoSearchRequestRequired(TypedDict):
"""Required fields for DuckDuckGo Search API request."""
@@ -159,7 +164,7 @@ class DuckDuckGoSearchConfig(BaseSearchConfig):
Returns:
SearchResponse with standardized format
"""
- response_json: Final = raw_response.json()
+ response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
# Extract max_results from the request URL params
query_params: Final = raw_response.request.url.params if raw_response.request else {}
@@ -178,17 +183,19 @@ class DuckDuckGoSearchConfig(BaseSearchConfig):
# Check if there's an Abstract with URL
if response_json.get("AbstractURL") and response_json.get("AbstractText"):
- abstract_result: Final = SearchResult(
- title=response_json.get("Heading", ""),
- url=response_json.get("AbstractURL", ""),
- snippet=response_json.get("AbstractText", ""),
- date=None,
- last_updated=None,
+ abstract_result: Final = SearchResult.model_validate(
+ {
+ "title": response_json.get("Heading", ""),
+ "url": response_json.get("AbstractURL", ""),
+ "snippet": response_json.get("AbstractText", ""),
+ "date": None,
+ "last_updated": None,
+ }
)
results.append(abstract_result)
# Process RelatedTopics
- related_topics: Final = response_json.get("RelatedTopics", [])
+ related_topics: Final = _JSON_VALUES.validate_python(response_json.get("RelatedTopics", []))
for topic in related_topics:
# Stop if we've reached max_results
if max_results is not None and len(results) >= max_results:
diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py
index cd904e493f2..c3c29e62f57 100644
--- a/litellm/llms/exa_ai/search/transformation.py
+++ b/litellm/llms/exa_ai/search/transformation.py
@@ -4,9 +4,11 @@ Calls Exa AI's /search endpoint to search the web.
Exa AI API Reference: https://docs.exa.ai/reference/search
"""
+from collections.abc import Iterable, Mapping
from typing import Final, TypedDict
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.search.transformation import (
@@ -16,6 +18,10 @@ from litellm.llms.base_llm.search.transformation import (
)
from litellm.secret_managers.main import get_secret_str
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
+_STRS: Final = TypeAdapter(Iterable[str], config=ConfigDict(hide_input_in_errors=True))
+
class _ExaAISearchRequestRequired(TypedDict):
"""Required fields for Exa AI Search API request."""
@@ -177,20 +183,23 @@ class ExaAISearchConfig(BaseSearchConfig):
Returns:
SearchResponse with standardized format
"""
- response_json: Final = raw_response.json()
+ response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
# Transform results to SearchResult objects
results: Final = []
- for result in response_json.get("results", []):
- search_result = SearchResult(
- title=result.get("title", ""),
- url=result.get("url", ""),
- snippet=result.get("text")
- or "\n\n".join(result.get("highlights") or [])
- or result.get("summary")
- or "",
- date=result.get("publishedDate"), # ISO 8601 datetime string
- last_updated=None, # Exa AI doesn't provide last_updated in response
+ for raw_result in _JSON_VALUES.validate_python(response_json.get("results", [])):
+ result = _JSON_OBJECT.validate_python(raw_result)
+ search_result = SearchResult.model_validate(
+ {
+ "title": result.get("title", ""),
+ "url": result.get("url", ""),
+ "snippet": result.get("text")
+ or "\n\n".join(_STRS.validate_python(result.get("highlights") or []))
+ or result.get("summary")
+ or "",
+ "date": result.get("publishedDate"), # ISO 8601 datetime string
+ "last_updated": None, # Exa AI doesn't provide last_updated in response
+ }
)
results.append(search_result)
diff --git a/litellm/llms/fal_ai/image_edit/transformation.py b/litellm/llms/fal_ai/image_edit/transformation.py
index d77769b130d..ea4e8cac5bf 100644
--- a/litellm/llms/fal_ai/image_edit/transformation.py
+++ b/litellm/llms/fal_ai/image_edit/transformation.py
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable
import httpx
from httpx._types import RequestFiles
+from pydantic import ConfigDict, TypeAdapter
from litellm.images.utils import ImageEditRequestUtils
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
@@ -23,6 +24,8 @@ from litellm.types.utils import FileTypes, ImageResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+_JSON_VALUE: Final = TypeAdapter(object)
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
DEFAULT_BASE_URL: Final[str] = "https://fal.run"
EDIT_SUFFIX: Final[str] = "/edit"
SUPPORTED_OPENAI_PARAMS: Final[tuple[str, ...]] = ("background", "mask", "n", "quality", "size")
@@ -163,7 +166,7 @@ class FalAIImageEditConfig(BaseImageEditConfig):
logging_obj: "LiteLLMLoggingObj",
) -> ImageResponse:
try:
- response_json: Final = raw_response.json()
+ response_json: Final = _JSON_VALUE.validate_python(raw_response.json())
except Exception as e:
raise self.get_error_class(
error_message=f"Error parsing Fal AI image edit response: {e}",
@@ -171,5 +174,7 @@ class FalAIImageEditConfig(BaseImageEditConfig):
headers=raw_response.headers,
)
model_response: Final = ImageResponse()
- model_response.data = list(fal_images_to_image_objects(response_json.get("images", ())))
+ model_response.data = list(
+ fal_images_to_image_objects(_JSON_OBJECT.validate_python(response_json).get("images", ()))
+ )
return model_response
diff --git a/litellm/llms/fal_ai/image_generation/bria_transformation.py b/litellm/llms/fal_ai/image_generation/bria_transformation.py
index 53da48e62ca..df8019747ef 100644
--- a/litellm/llms/fal_ai/image_generation/bria_transformation.py
+++ b/litellm/llms/fal_ai/image_generation/bria_transformation.py
@@ -1,6 +1,8 @@
+from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
from litellm.types.utils import ImageObject, ImageResponse
@@ -15,6 +17,9 @@ if TYPE_CHECKING:
else:
LiteLLMLoggingObj = Any
+_JSON_VALUE: Final = TypeAdapter(object)
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+
class FalAIBriaConfig(FalAIBaseConfig):
"""
@@ -206,7 +211,7 @@ class FalAIBriaConfig(FalAIBaseConfig):
}
"""
try:
- response_data: Final = raw_response.json()
+ response_data: Final = _JSON_VALUE.validate_python(raw_response.json())
except Exception as e:
raise self.get_error_class(
error_message=f"Error transforming image generation response: {e}",
@@ -218,7 +223,7 @@ class FalAIBriaConfig(FalAIBaseConfig):
model_response.data = []
# Handle Bria response format - uses "image" (singular) not "images"
- image_data: Final = response_data.get("image")
+ image_data: Final = _JSON_OBJECT.validate_python(response_data).get("image")
if image_data and isinstance(image_data, dict):
model_response.data.append(
ImageObject(
diff --git a/litellm/llms/fal_ai/image_generation/imagen4_transformation.py b/litellm/llms/fal_ai/image_generation/imagen4_transformation.py
index c5a5f0112ca..09e96926767 100644
--- a/litellm/llms/fal_ai/image_generation/imagen4_transformation.py
+++ b/litellm/llms/fal_ai/image_generation/imagen4_transformation.py
@@ -1,6 +1,8 @@
+from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, cast
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
@@ -16,6 +18,9 @@ if TYPE_CHECKING:
else:
LiteLLMLoggingObj = Any
+_JSON_VALUE: Final = TypeAdapter(object)
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+
class FalAIImagen4Config(FalAIBaseConfig):
"""
@@ -203,7 +208,7 @@ class FalAIImagen4Config(FalAIBaseConfig):
}
"""
try:
- response_data: Final = raw_response.json()
+ response_data: Final = _JSON_VALUE.validate_python(raw_response.json())
except Exception as e:
raise self.get_error_class(
error_message=f"Error transforming image generation response: {e}",
@@ -215,7 +220,8 @@ class FalAIImagen4Config(FalAIBaseConfig):
model_response.data = []
# Handle Imagen4 response format
- images: Final = response_data.get("images", [])
+ response_object: Final = _JSON_OBJECT.validate_python(response_data)
+ images: Final = response_object.get("images", [])
if isinstance(images, list):
for image_data in images:
if isinstance(image_data, dict):
@@ -239,7 +245,7 @@ class FalAIImagen4Config(FalAIBaseConfig):
hidden_params: Final = cast( # cast-ok: preserve mapping operations on dynamic response metadata
dict[str, object], getattr(model_response, HIDDEN_PARAMS_ATTR)
)
- if "seed" in response_data:
- hidden_params["seed"] = response_data["seed"]
+ if "seed" in response_object:
+ hidden_params["seed"] = response_object["seed"]
return model_response
diff --git a/litellm/llms/fal_ai/image_generation/transformation.py b/litellm/llms/fal_ai/image_generation/transformation.py
index 8f081c7228d..0bb5dada071 100644
--- a/litellm/llms/fal_ai/image_generation/transformation.py
+++ b/litellm/llms/fal_ai/image_generation/transformation.py
@@ -2,7 +2,7 @@ from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final
import httpx
-from pydantic import TypeAdapter
+from pydantic import ConfigDict, TypeAdapter
from typing_extensions import ReadOnly, TypedDict
from litellm.llms.base_llm.image_generation.transformation import (
@@ -31,6 +31,8 @@ class FalImageProviderSpecificFields(TypedDict, total=False):
_FAL_IMAGE_DATA: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
+_JSON_VALUE: Final = TypeAdapter(object)
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
def fal_images_to_image_objects(images: object) -> tuple[ImageObject, ...]:
@@ -124,7 +126,7 @@ class FalAIBaseConfig(BaseImageGenerationConfig):
Transform the image generation response to the litellm image response
"""
try:
- response_data: Final = raw_response.json()
+ response_data: Final = _JSON_VALUE.validate_python(raw_response.json())
except Exception as e:
raise self.get_error_class(
error_message=f"Error transforming image generation response: {e}",
@@ -134,7 +136,9 @@ class FalAIBaseConfig(BaseImageGenerationConfig):
if not model_response.data:
model_response.data = []
- model_response.data.extend(fal_images_to_image_objects(response_data.get("images", ())))
+ model_response.data.extend(
+ fal_images_to_image_objects(_JSON_OBJECT.validate_python(response_data).get("images", ()))
+ )
return model_response
diff --git a/litellm/llms/fastcrw/search/transformation.py b/litellm/llms/fastcrw/search/transformation.py
index 2ded3208f2e..0aef446c23b 100644
--- a/litellm/llms/fastcrw/search/transformation.py
+++ b/litellm/llms/fastcrw/search/transformation.py
@@ -8,9 +8,11 @@ or cloud). The search response uses the Firecrawl-compatible envelope
fastCRW API Reference: https://fastcrw.com/docs/rest-api
"""
+from collections.abc import Iterable, Mapping
from typing import Final, TypedDict
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.search.transformation import (
@@ -20,6 +22,9 @@ from litellm.llms.base_llm.search.transformation import (
)
from litellm.secret_managers.main import get_secret_str
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
+
class _FastCRWSearchRequestRequired(TypedDict):
"""Required fields for fastCRW Search API request."""
@@ -157,21 +162,24 @@ class FastCRWSearchConfig(BaseSearchConfig):
Returns:
SearchResponse with standardized format
"""
- response_json: Final = raw_response.json()
+ response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
results: Final = []
data: Final = response_json.get("data", [])
if isinstance(data, list):
- for result in data:
+ for raw_result in _JSON_VALUES.validate_python(data):
+ result = _JSON_OBJECT.validate_python(raw_result)
snippet = result.get("markdown") or result.get("description", "")
- search_result = SearchResult(
- title=result.get("title", ""),
- url=result.get("url", ""),
- snippet=snippet,
- date=None,
- last_updated=None,
+ search_result = SearchResult.model_validate(
+ {
+ "title": result.get("title", ""),
+ "url": result.get("url", ""),
+ "snippet": snippet,
+ "date": None,
+ "last_updated": None,
+ }
)
results.append(search_result)
diff --git a/litellm/llms/firecrawl/search/transformation.py b/litellm/llms/firecrawl/search/transformation.py
index 552388ce47c..736904c4e36 100644
--- a/litellm/llms/firecrawl/search/transformation.py
+++ b/litellm/llms/firecrawl/search/transformation.py
@@ -4,9 +4,11 @@ Calls Firecrawl's /search endpoint to search the web.
Firecrawl API Reference: https://docs.firecrawl.dev/api-reference/endpoint/search
"""
+from collections.abc import Iterable, Mapping
from typing import Final, TypedDict
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.search.transformation import (
@@ -16,6 +18,9 @@ from litellm.llms.base_llm.search.transformation import (
)
from litellm.secret_managers.main import get_secret_str
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
+
class _FirecrawlSearchRequestRequired(TypedDict):
"""Required fields for Firecrawl Search API request."""
@@ -170,7 +175,7 @@ class FirecrawlSearchConfig(BaseSearchConfig):
Returns:
SearchResponse with standardized format
"""
- response_json: Final = raw_response.json()
+ response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
# Transform results to SearchResult objects
results: Final = []
@@ -179,44 +184,54 @@ class FirecrawlSearchConfig(BaseSearchConfig):
if isinstance(data, list):
# Self-hosted Firecrawl (v1) format: data is a flat list of results
- for result in data:
+ for raw_result in _JSON_VALUES.validate_python(data):
+ result = _JSON_OBJECT.validate_python(raw_result)
snippet = result.get("markdown") or result.get("description", "")
- search_result = SearchResult(
- title=result.get("title", ""),
- url=result.get("url", ""),
- snippet=snippet,
- date=None,
- last_updated=None,
+ search_result = SearchResult.model_validate(
+ {
+ "title": result.get("title", ""),
+ "url": result.get("url", ""),
+ "snippet": snippet,
+ "date": None,
+ "last_updated": None,
+ }
)
results.append(search_result)
elif isinstance(data, dict):
# Firecrawl Cloud (v2) format: data is a dict with web/news keys
- web_results: Final = data.get("web", [])
+ sections: Final = _JSON_OBJECT.validate_python(data)
+ web_results: Final = _JSON_VALUES.validate_python(sections.get("web", []))
- for result in web_results:
+ for raw_result in web_results:
+ result = _JSON_OBJECT.validate_python(raw_result)
# Use markdown if available, otherwise fall back to description
snippet = result.get("markdown") or result.get("description", "")
- search_result = SearchResult(
- title=result.get("title", ""),
- url=result.get("url", ""),
- snippet=snippet,
- date=None,
- last_updated=None,
+ search_result = SearchResult.model_validate(
+ {
+ "title": result.get("title", ""),
+ "url": result.get("url", ""),
+ "snippet": snippet,
+ "date": None,
+ "last_updated": None,
+ }
)
results.append(search_result)
# Process news results if available (they have date field)
- news_results: Final = data.get("news", [])
- for result in news_results:
+ news_results: Final = _JSON_VALUES.validate_python(sections.get("news", []))
+ for raw_result in news_results:
+ result = _JSON_OBJECT.validate_python(raw_result)
snippet = result.get("markdown") or result.get("snippet", "")
- search_result = SearchResult(
- title=result.get("title", ""),
- url=result.get("url", ""),
- snippet=snippet,
- date=result.get("date"), # News results include date
- last_updated=None,
+ search_result = SearchResult.model_validate(
+ {
+ "title": result.get("title", ""),
+ "url": result.get("url", ""),
+ "snippet": snippet,
+ "date": result.get("date"), # News results include date
+ "last_updated": None,
+ }
)
results.append(search_result)
diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py
index 2c240544ee8..bd7c0d11b7d 100644
--- a/litellm/llms/fireworks_ai/chat/transformation.py
+++ b/litellm/llms/fireworks_ai/chat/transformation.py
@@ -1,8 +1,9 @@
import json
-from collections.abc import AsyncIterator, Iterator, Mapping
+from collections.abc import AsyncIterator, Iterable, Iterator, Mapping
from typing import TYPE_CHECKING, Final, Literal, cast
import httpx
+from pydantic import ConfigDict, TypeAdapter
import litellm
from litellm._logging import verbose_logger
@@ -51,6 +52,10 @@ from ..common_utils import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
+_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True))
+_STR: Final = TypeAdapter(str)
+
def _map_reasoning_effort(value: object) -> object:
effort: Final[object] = cast(Mapping[str, object], value).get("effort") if isinstance(value, Mapping) else value
@@ -757,7 +762,8 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig):
additional_headers: Final = get_response_headers(raw_response_headers)
- response: Final = ModelResponse(**completion_response)
+ completion_fields: Final = _JSON_DICT.validate_python(completion_response)
+ response: Final = ModelResponse(**completion_fields)
if response.model is not None:
response.model = "fireworks_ai/" + response.model
@@ -771,7 +777,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig):
response.hidden_params = {
"additional_headers": additional_headers,
- **_extract_fireworks_hidden_params(completion_response),
+ **_extract_fireworks_hidden_params(completion_fields),
}
return response
@@ -832,9 +838,9 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig):
f"Failed to fetch models from Fireworks AI. Status code: {response.status_code}, Response: {response.json()}"
)
- models: Final = response.json()["models"]
+ models: Final = _JSON_OBJECTS.validate_python(_JSON_DICT.validate_python(response.json())["models"])
- return ["fireworks_ai/" + model["name"] for model in models]
+ return ["fireworks_ai/" + _STR.validate_python(model["name"]) for model in models]
@staticmethod
def get_api_key(api_key: str | None = None) -> str | None:
diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py
index 4d13d0d2e6b..fe82a9cae81 100644
--- a/litellm/llms/gemini/image_generation/transformation.py
+++ b/litellm/llms/gemini/image_generation/transformation.py
@@ -1,4 +1,4 @@
-from collections.abc import Iterable, Mapping
+from collections.abc import Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final
import httpx
@@ -35,6 +35,10 @@ else:
_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True))
_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object]]] = TypeAdapter(
+ Mapping[str, object] | str | Sequence[object], config=ConfigDict(hide_input_in_errors=True)
+)
class GoogleImageGenConfig(BaseImageGenerationConfig):
@@ -199,19 +203,22 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
# Handle different response formats based on model
if is_gemini_image_model(model):
# Gemini Flash Image Preview models return in candidates format
- candidates: Final = response_data.get("candidates", [])
+ response_object: Final = _JSON_OBJECT.validate_python(response_data)
+ candidates: Final = _JSON_CONTAINER.validate_python(response_object.get("candidates", []))
for candidate in candidates:
- content = candidate.get("content", {})
- parts = content.get("parts", [])
- for part in parts:
+ content = _JSON_OBJECT.validate_python(_JSON_OBJECT.validate_python(candidate).get("content", {}))
+ parts = _JSON_CONTAINER.validate_python(content.get("parts", []))
+ for raw_part in parts:
+ part = raw_part if isinstance(raw_part, str) else _JSON_CONTAINER.validate_python(raw_part)
# Look for inlineData with image
if "inlineData" in part:
- inline_data = part["inlineData"]
+ part_object = _JSON_OBJECT.validate_python(part)
+ inline_data = _JSON_CONTAINER.validate_python(part_object["inlineData"])
if "data" in inline_data:
- thought_sig = part.get("thoughtSignature")
+ thought_sig = part_object.get("thoughtSignature")
model_response.data.append(
ImageObject(
- b64_json=inline_data["data"],
+ b64_json=_JSON_OBJECT.validate_python(inline_data)["data"],
url=None,
provider_specific_fields=(
{"thought_signature": thought_sig} if thought_sig else None
@@ -224,12 +231,14 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
model_response.usage = transform_gemini_image_usage(
_JSON_DICT.validate_python(response_data["usageMetadata"])
)
- web_search_requests: Final = get_gemini_image_web_search_requests(response_data)
+ web_search_requests: Final = get_gemini_image_web_search_requests(response_object)
if web_search_requests and model_response.usage is not None:
setattr(model_response.usage, "web_search_requests", web_search_requests)
else:
# Original Imagen format - predictions with generated images
- predictions: Final = _JSON_OBJECTS.validate_python(response_data.get("predictions", []))
+ predictions: Final = _JSON_OBJECTS.validate_python(
+ _JSON_OBJECT.validate_python(response_data).get("predictions", [])
+ )
for prediction in predictions:
# Google AI returns base64 encoded images in the prediction
model_response.data.append(
diff --git a/litellm/llms/google_pse/search/transformation.py b/litellm/llms/google_pse/search/transformation.py
index 0128db9a823..1ae465277c5 100644
--- a/litellm/llms/google_pse/search/transformation.py
+++ b/litellm/llms/google_pse/search/transformation.py
@@ -4,9 +4,11 @@ Calls Google Programmable Search Engine (PSE) API to search the web.
Google PSE API Reference: https://developers.google.com/custom-search/v1/reference/rest/v1/cse/list
"""
+from collections.abc import Iterable, Mapping
from typing import Final, Literal, TypedDict
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.search.transformation import (
@@ -16,6 +18,9 @@ from litellm.llms.base_llm.search.transformation import (
)
from litellm.secret_managers.main import get_secret_str
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
+
class _GooglePSESearchRequestRequired(TypedDict):
"""Required fields for Google PSE Search API request."""
@@ -239,17 +244,20 @@ class GooglePSESearchConfig(BaseSearchConfig):
Returns:
SearchResponse with standardized format
"""
- response_json: Final = raw_response.json()
+ response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
# Transform results to SearchResult objects
results: Final = []
- for item in response_json.get("items", []):
- search_result = SearchResult(
- title=item.get("title", ""),
- url=item.get("link", ""),
- snippet=item.get("snippet", ""),
- date=None, # Google PSE doesn't provide date in standard response
- last_updated=None, # Google PSE doesn't provide last_updated in response
+ for raw_item in _JSON_VALUES.validate_python(response_json.get("items", [])):
+ item = _JSON_OBJECT.validate_python(raw_item)
+ search_result = SearchResult.model_validate(
+ {
+ "title": item.get("title", ""),
+ "url": item.get("link", ""),
+ "snippet": item.get("snippet", ""),
+ "date": None, # Google PSE doesn't provide date in standard response
+ "last_updated": None, # Google PSE doesn't provide last_updated in response
+ }
)
results.append(search_result)
diff --git a/litellm/llms/infinity/embedding/transformation.py b/litellm/llms/infinity/embedding/transformation.py
index 26130e4ca6e..f289918cb3c 100644
--- a/litellm/llms/infinity/embedding/transformation.py
+++ b/litellm/llms/infinity/embedding/transformation.py
@@ -1,6 +1,8 @@
+from collections.abc import Mapping
from typing import Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
@@ -11,6 +13,9 @@ from litellm.types.utils import EmbeddingResponse, Usage
from ..common_utils import InfinityError
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_INT: Final = TypeAdapter(int)
+
class InfinityEmbeddingConfig(BaseEmbeddingConfig):
"""
@@ -124,9 +129,10 @@ class InfinityEmbeddingConfig(BaseEmbeddingConfig):
model_response.data = raw_response_json.get("data")
model_response.object = raw_response_json.get("object")
+ usage_payload: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {}))
usage: Final = Usage(
- prompt_tokens=raw_response_json.get("usage", {}).get("prompt_tokens", 0),
- total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0),
+ prompt_tokens=_INT.validate_python(usage_payload.get("prompt_tokens", 0) or 0),
+ total_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0),
)
model_response.usage = usage
return model_response
diff --git a/litellm/llms/langgraph/chat/transformation.py b/litellm/llms/langgraph/chat/transformation.py
index f1f4054099a..4f93f435957 100644
--- a/litellm/llms/langgraph/chat/transformation.py
+++ b/litellm/llms/langgraph/chat/transformation.py
@@ -10,7 +10,7 @@ Non-streaming endpoint: POST /runs/wait
import json
from collections.abc import Mapping
-from typing import TYPE_CHECKING, Any, Final, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Final, Optional, Union
import httpx
@@ -369,7 +369,7 @@ class LangGraphConfig(BaseConfig):
from litellm.utils import CustomStreamWrapper
if client is None or not isinstance(client, AsyncHTTPHandler):
- client = get_async_httpx_client(llm_provider=cast(Any, "langgraph"), params={})
+ client = get_async_httpx_client(llm_provider="langgraph", params={})
verbose_logger.debug("Making async streaming request to: %s", api_base)
diff --git a/litellm/llms/lemonade/chat/transformation.py b/litellm/llms/lemonade/chat/transformation.py
index c4079dfc73b..82712af9af3 100644
--- a/litellm/llms/lemonade/chat/transformation.py
+++ b/litellm/llms/lemonade/chat/transformation.py
@@ -2,10 +2,12 @@
Translate from OpenAI's `/v1/chat/completions` to Lemonade's `/v1/chat/completions`
"""
+from collections.abc import Iterable, Mapping
from typing import TYPE_CHECKING, Any, Final
from urllib.parse import quote
import httpx
+from pydantic import ConfigDict, TypeAdapter
import litellm
from litellm._logging import verbose_logger
@@ -21,6 +23,10 @@ from ...openai_like.chat.transformation import OpenAILikeChatConfig
if TYPE_CHECKING:
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
+_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True))
+_STR: Final = TypeAdapter(str)
+
class LemonadeChatConfig(OpenAILikeChatConfig):
_DEFAULT_API_KEY = "lemonade"
@@ -55,7 +61,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig):
response_format: dict | None = None,
tools: list | None = None,
) -> None:
- locals_: Final = locals().copy()
+ locals_: Final[Mapping[str, object]] = dict(locals())
for key, value in locals_.items():
if key != "self" and value is not None:
setattr(self.__class__, key, value)
@@ -104,8 +110,8 @@ class LemonadeChatConfig(OpenAILikeChatConfig):
f"Failed to fetch models from Lemonade. Status code: {response.status_code}, Response: {response.text}"
)
- model_list: Final = response.json().get("data", [])
- return ["lemonade/" + model["id"] for model in model_list]
+ model_list: Final = _JSON_OBJECTS.validate_python(_JSON_DICT.validate_python(response.json()).get("data", []))
+ return ["lemonade/" + _STR.validate_python(model["id"]) for model in model_list]
@staticmethod
def _get_positive_int(value: Any) -> int | None:
@@ -188,10 +194,11 @@ class LemonadeChatConfig(OpenAILikeChatConfig):
verbose_logger.debug("LemonadeError: Could not get model info.")
return self._get_default_model_info(model)
- max_input_tokens: Final = self._get_context_window(model_info)
- max_output_tokens: Final = self._get_positive_int(model_info.get("max_output_tokens"))
- max_tokens: Final = self._get_positive_int(model_info.get("max_tokens"))
- provider_specific_entry: Final = self._get_provider_specific_entry(model_info)
+ model_info_fields: Final = _JSON_DICT.validate_python(model_info)
+ max_input_tokens: Final = self._get_context_window(model_info_fields)
+ max_output_tokens: Final = self._get_positive_int(model_info_fields.get("max_output_tokens"))
+ max_tokens: Final = self._get_positive_int(model_info_fields.get("max_tokens"))
+ provider_specific_entry: Final = self._get_provider_specific_entry(model_info_fields)
model_info_response: Final = self._get_default_model_info(model)
model_info_response.update(
diff --git a/litellm/llms/linkup/search/transformation.py b/litellm/llms/linkup/search/transformation.py
index 8ecd5a24c08..f20a1f5285f 100644
--- a/litellm/llms/linkup/search/transformation.py
+++ b/litellm/llms/linkup/search/transformation.py
@@ -4,9 +4,11 @@ Calls Linkup's /search endpoint to search the web.
Linkup API Reference: https://docs.linkup.so/pages/documentation/api-reference/endpoint/post-search
"""
+from collections.abc import Iterable, Mapping
from typing import Final, Literal, TypedDict
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.search.transformation import (
@@ -16,6 +18,9 @@ from litellm.llms.base_llm.search.transformation import (
)
from litellm.secret_managers.main import get_secret_str
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
+
class _LinkupSearchRequestRequired(TypedDict):
"""Required fields for Linkup Search API request."""
@@ -165,35 +170,40 @@ class LinkupSearchConfig(BaseSearchConfig):
Returns:
SearchResponse with standardized format
"""
- response_json: Final = raw_response.json()
+ response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
# Transform results to SearchResult objects
results: Final = []
# Process results array
- raw_results: Final = response_json.get("results", [])
+ raw_results: Final = _JSON_VALUES.validate_python(response_json.get("results", []))
- for result in raw_results:
+ for raw_result in raw_results:
+ result = _JSON_OBJECT.validate_python(raw_result)
# Handle both text and image result types
result_type = result.get("type", "text")
if result_type == "text":
- search_result = SearchResult(
- title=result.get("name", ""),
- url=result.get("url", ""),
- snippet=result.get("content", ""),
- date=None,
- last_updated=None,
+ search_result = SearchResult.model_validate(
+ {
+ "title": result.get("name", ""),
+ "url": result.get("url", ""),
+ "snippet": result.get("content", ""),
+ "date": None,
+ "last_updated": None,
+ }
)
results.append(search_result)
elif result_type == "image":
# For image results, use the URL as both title and snippet if name not provided
- search_result = SearchResult(
- title=result.get("name", result.get("url", "")),
- url=result.get("url", ""),
- snippet=result.get("content", ""),
- date=None,
- last_updated=None,
+ search_result = SearchResult.model_validate(
+ {
+ "title": result.get("name", result.get("url", "")),
+ "url": result.get("url", ""),
+ "snippet": result.get("content", ""),
+ "date": None,
+ "last_updated": None,
+ }
)
results.append(search_result)
diff --git a/litellm/llms/mistral/audio_speech/transformation.py b/litellm/llms/mistral/audio_speech/transformation.py
index 04a7e3b9341..5a8960d0898 100644
--- a/litellm/llms/mistral/audio_speech/transformation.py
+++ b/litellm/llms/mistral/audio_speech/transformation.py
@@ -11,6 +11,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Final
import httpx
+from pydantic import ConfigDict, TypeAdapter, ValidationError
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.text_to_speech.transformation import (
@@ -23,6 +24,8 @@ if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.llms.openai import HttpxBinaryResponseContent
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+
class MistralTextToSpeechException(BaseLLMException):
pass
@@ -143,7 +146,7 @@ class MistralTextToSpeechConfig(BaseTextToSpeechConfig):
return request_data
def _requested_content_type(self, request: httpx.Request) -> str:
- request_body: Final = json.loads(request.content or b"{}")
+ request_body: Final = _JSON_OBJECT.validate_python(json.loads(request.content or b"{}"))
requested_format: Final = request_body.get("response_format")
if not isinstance(requested_format, str):
return "audio/mpeg"
@@ -158,7 +161,9 @@ class MistralTextToSpeechConfig(BaseTextToSpeechConfig):
from litellm.types.llms.openai import HttpxBinaryResponseContent
try:
- response_json: Final = raw_response.json()
+ response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
+ except ValidationError:
+ raise
except (json.JSONDecodeError, ValueError):
raise MistralTextToSpeechException(
status_code=raw_response.status_code,
diff --git a/litellm/llms/nlp_cloud/chat/transformation.py b/litellm/llms/nlp_cloud/chat/transformation.py
index 2ff48894619..0cca299de93 100644
--- a/litellm/llms/nlp_cloud/chat/transformation.py
+++ b/litellm/llms/nlp_cloud/chat/transformation.py
@@ -1,5 +1,6 @@
import json
import time
+from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final
import httpx
@@ -79,7 +80,7 @@ class NLPCloudConfig(BaseConfig):
num_beams: int | None = None,
num_return_sequences: int | None = None,
) -> None:
- locals_: Final = locals().copy()
+ locals_: Final[Mapping[str, object]] = dict(locals())
for key, value in locals_.items():
if key != "self" and value is not None:
setattr(self.__class__, key, value)
diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py
index 14ef6225666..6b038872b8d 100644
--- a/litellm/llms/oci/chat/transformation.py
+++ b/litellm/llms/oci/chat/transformation.py
@@ -759,7 +759,7 @@ class OCIStreamWrapper(CustomStreamWrapper):
self.sent_last_chunk = True
return self.model_response_creator(chunk={"choices": parsed.choices})
- def chunk_creator(self, chunk: Any) -> ModelResponseStream | None:
+ def chunk_creator(self, chunk: object) -> ModelResponseStream | None:
if not isinstance(chunk, str):
raise ValueError(f"Chunk is not a string: {chunk}")
if not chunk.startswith("data:"):
diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py
index 409e763dde1..af8d492316e 100644
--- a/litellm/llms/ollama/common_utils.py
+++ b/litellm/llms/ollama/common_utils.py
@@ -1,6 +1,6 @@
import base64
import io
-from typing import Any, Final
+from typing import Final
import httpx
@@ -174,7 +174,7 @@ class OllamaModelInfo(BaseLLMModelInfo):
model: str,
api_base: str | None = None,
api_key: str | None = None,
- ) -> dict[str, Any]:
+ ) -> dict[str, object]:
from litellm import module_level_client
model = self._strip_ollama_model_prefix(model)
@@ -223,7 +223,7 @@ class OllamaModelInfo(BaseLLMModelInfo):
model: str,
api_base: str | None = None,
api_key: str | None = None,
- ) -> dict[str, Any] | None:
+ ) -> dict[str, object] | None:
if self._is_static_ollama_model(model):
return None
return self.get_runtime_model_info(model=model, api_base=api_base, api_key=api_key)
diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py
index dcae8524c4f..09946251227 100644
--- a/litellm/llms/ollama/completion/transformation.py
+++ b/litellm/llms/ollama/completion/transformation.py
@@ -97,6 +97,7 @@ class _OllamaGenerateResponse(TypedDict):
_OLLAMA_GENERATE_RESPONSE: Final = TypeAdapter(_OllamaGenerateResponse)
+_JSON_VALUE: Final = TypeAdapter(object)
_JSON_CODE_FENCE: Final = re.compile(r"^```(?:json)?\s*(.*?)\s*(?:```)?$", re.DOTALL)
_JSON_OBJECT_START: Final = re.compile(r"(?:```(?:json)?\s*)?\{")
_JSON_OBJECT_PARTIAL_START: Final = re.compile(r"`{0,3}|```(?:j|js|jso|json)?\s*")
@@ -334,7 +335,9 @@ class OllamaConfig(BaseConfig):
model_response.choices[0].finish_reason = "stop"
else:
try:
- response_content: Final[object] = json.loads(_strip_json_code_fence(response_text))
+ response_content: Final = _JSON_VALUE.validate_python(
+ json.loads(_strip_json_code_fence(response_text))
+ )
# Check if this is a function call format with name/arguments structure
if (
diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py
index 0d7a553d633..02a05b060b8 100644
--- a/litellm/llms/openai/responses/count_tokens/handler.py
+++ b/litellm/llms/openai/responses/count_tokens/handler.py
@@ -6,7 +6,7 @@ Uses httpx for HTTP requests to OpenAI's /v1/responses/input_tokens endpoint.
import json
from collections.abc import Sequence
-from typing import Any, Final
+from typing import Final
import httpx
@@ -33,7 +33,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig):
timeout: float | httpx.Timeout | None = None,
tools: list[dict[str, object]] | None = None,
instructions: str | None = None,
- ) -> dict[str, Any]:
+ ) -> dict[str, object]:
"""
Handle a token counting request to OpenAI's Responses API.
diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py
index c05d943b7cf..ab3ddbd179d 100644
--- a/litellm/llms/openai/responses/count_tokens/token_counter.py
+++ b/litellm/llms/openai/responses/count_tokens/token_counter.py
@@ -81,12 +81,14 @@ class OpenAITokenCounter(BaseTokenCounter):
)
if result is not None:
- return TokenCountResponse(
- total_tokens=result.get("input_tokens", 0),
- request_model=request_model,
- model_used=model_to_use,
- tokenizer_type="openai_api",
- original_response=result,
+ return TokenCountResponse.model_validate(
+ {
+ "total_tokens": result.get("input_tokens", 0),
+ "request_model": request_model,
+ "model_used": model_to_use,
+ "tokenizer_type": "openai_api",
+ "original_response": result,
+ }
)
except OpenAIError as e:
verbose_logger.warning("OpenAI CountTokens API error: status=%s, message=%s", e.status_code, e.message)
diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py
index 9f596505f91..45d8ab0b4c4 100644
--- a/litellm/llms/openai/responses/count_tokens/transformation.py
+++ b/litellm/llms/openai/responses/count_tokens/transformation.py
@@ -118,7 +118,7 @@ class OpenAICountTokensConfig:
self,
model: str,
input: str | Sequence[object],
- tools: list[dict[str, Any]] | None = None,
+ tools: list[dict[str, object]] | None = None,
instructions: str | None = None,
) -> dict[str, object]:
"""
diff --git a/litellm/llms/openrouter/chat/transformation.py b/litellm/llms/openrouter/chat/transformation.py
index c63863bbca4..9aef15f15a3 100644
--- a/litellm/llms/openrouter/chat/transformation.py
+++ b/litellm/llms/openrouter/chat/transformation.py
@@ -6,11 +6,12 @@ Calls done in OpenAI/openai.py as OpenRouter is openai-compatible.
Docs: https://openrouter.ai/docs/parameters
"""
-from collections.abc import AsyncIterator, Iterator
+from collections.abc import AsyncIterator, Iterator, Mapping
from enum import Enum
-from typing import TYPE_CHECKING, Any, Final, cast
+from typing import TYPE_CHECKING, Final, cast
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
@@ -26,6 +27,8 @@ if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+
class CacheControlSupportedModels(str, Enum):
"""Models that support cache_control in content blocks."""
@@ -203,9 +206,9 @@ class OpenrouterConfig(OpenAIGPTConfig):
# Extract cost from OpenRouter response body
# OpenRouter returns cost information in the usage object when usage.include=true
try:
- response_json: Final = raw_response.json()
+ response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
if "usage" in response_json and response_json["usage"]:
- response_cost: Final = response_json["usage"].get("cost")
+ response_cost: Final = _JSON_OBJECT.validate_python(response_json["usage"]).get("cost")
if response_cost is not None:
# Store cost in hidden params for the cost calculator to use
if not hasattr(model_response, HIDDEN_PARAMS_ATTR):
@@ -218,7 +221,8 @@ class OpenrouterConfig(OpenAIGPTConfig):
additional_headers: Final = cast( # cast-ok: preserve mapping operations on response metadata
dict[str, object], hidden_params["additional_headers"]
)
- additional_headers["llm_provider-x-litellm-response-cost"] = float(response_cost)
+ if isinstance(response_cost, str | int | float):
+ additional_headers["llm_provider-x-litellm-response-cost"] = float(response_cost)
except Exception:
# If we can't extract cost, continue without it - don't fail the response
pass
@@ -237,7 +241,7 @@ class OpenrouterConfig(OpenAIGPTConfig):
streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse,
sync_stream: bool,
json_mode: bool | None = False,
- ) -> Any:
+ ) -> "OpenRouterChatCompletionStreamingHandler":
return OpenRouterChatCompletionStreamingHandler(
streaming_response=streaming_response,
sync_stream=sync_stream,
diff --git a/litellm/llms/opensandbox/sandbox/transformation.py b/litellm/llms/opensandbox/sandbox/transformation.py
index d2c558241f0..13c9d805638 100644
--- a/litellm/llms/opensandbox/sandbox/transformation.py
+++ b/litellm/llms/opensandbox/sandbox/transformation.py
@@ -1,9 +1,11 @@
import asyncio
import json
import time
+from collections.abc import Mapping
from typing import Final, cast
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.constants import (
OPEN_SANDBOX_API_BASE_ENV_VAR,
@@ -37,6 +39,9 @@ DEFAULT_READY_TIMEOUT: Final = OPEN_SANDBOX_READY_TIMEOUT
DEFAULT_POLL_INTERVAL: Final = OPEN_SANDBOX_POLL_INTERVAL
MAX_OUTPUT_BYTES: Final = SANDBOX_MAX_OUTPUT_BYTES
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_VALUE: Final = TypeAdapter(object)
+
class OpenSandboxSandboxConfig(BaseSandboxConfig):
def _http(self, client: AsyncHTTPHandler | None) -> AsyncHTTPHandler:
@@ -92,7 +97,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig):
headers=self._lifecycle_headers(key),
json=body,
)
- data: Final = response.json()
+ data: Final = _JSON_OBJECT.validate_python(response.json())
sandbox_id: Final = str(data["id"])
if self._sandbox_state(data) != "Running":
@@ -256,7 +261,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig):
url=f"{api_base}/sandboxes/{sandbox_id}",
headers=headers,
)
- data = response.json()
+ data = _JSON_VALUE.validate_python(response.json())
state = self._sandbox_state(data)
if state == "Running":
return
diff --git a/litellm/llms/ovhcloud/embedding/transformation.py b/litellm/llms/ovhcloud/embedding/transformation.py
index b78a81d05bf..cdf095e6ad1 100644
--- a/litellm/llms/ovhcloud/embedding/transformation.py
+++ b/litellm/llms/ovhcloud/embedding/transformation.py
@@ -3,9 +3,11 @@ This is OpenAI compatible - no transformation is applied
"""
+from collections.abc import Mapping
from typing import Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
@@ -16,6 +18,9 @@ from litellm.types.utils import EmbeddingResponse, Usage
from ..utils import OVHCloudException
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_INT: Final = TypeAdapter(int)
+
class OVHCloudEmbeddingConfig(BaseEmbeddingConfig):
def __init__(self) -> None:
@@ -107,9 +112,10 @@ class OVHCloudEmbeddingConfig(BaseEmbeddingConfig):
model_response.data = raw_response_json.get("data")
model_response.object = raw_response_json.get("object")
+ usage_payload: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {}))
usage: Final = Usage(
- prompt_tokens=raw_response_json.get("usage", {}).get("prompt_tokens", 0),
- total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0),
+ prompt_tokens=_INT.validate_python(usage_payload.get("prompt_tokens", 0) or 0),
+ total_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0),
)
model_response.usage = usage
diff --git a/litellm/llms/perplexity/chat/transformation.py b/litellm/llms/perplexity/chat/transformation.py
index 5b85e1e8ae6..017a74cd51e 100644
--- a/litellm/llms/perplexity/chat/transformation.py
+++ b/litellm/llms/perplexity/chat/transformation.py
@@ -7,6 +7,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
import litellm
from litellm._logging import verbose_logger
@@ -20,6 +21,7 @@ if TYPE_CHECKING:
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
_PERPLEXITY_ATTRIBUTION_HEADERS: Final[Mapping[str, str]] = MappingProxyType({"X-Pplx-Integration": "litellm"})
+_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True))
class PerplexityChatConfig(OpenAIGPTConfig):
@@ -110,7 +112,7 @@ class PerplexityChatConfig(OpenAIGPTConfig):
# Extract and enhance usage with Perplexity-specific fields
try:
- raw_response_json: Final = raw_response.json()
+ raw_response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
self._enhance_usage_with_perplexity_fields(model_response, raw_response_json)
self._add_citations_as_annotations(model_response, raw_response_json)
except Exception as e:
diff --git a/litellm/llms/perplexity/embedding/transformation.py b/litellm/llms/perplexity/embedding/transformation.py
index c93206db2bb..8e9e69e1099 100644
--- a/litellm/llms/perplexity/embedding/transformation.py
+++ b/litellm/llms/perplexity/embedding/transformation.py
@@ -13,9 +13,11 @@ This module decodes them into float arrays for OpenAI-compatible responses.
import base64
import struct
+from collections.abc import Mapping
from typing import Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
@@ -24,6 +26,9 @@ from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues
from litellm.types.utils import EmbeddingResponse, Usage
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_INT: Final = TypeAdapter(int)
+
class PerplexityEmbeddingError(BaseLLMException):
def __init__(
@@ -156,15 +161,17 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig):
raw_data: Final = raw_response_json.get("data", [])
decoded_data: Final[list[dict[str, object]]] = []
for item in raw_data:
- decoded_item = dict(item)
- decoded_item["embedding"] = self._decode_base64_embedding(item.get("embedding"))
+ decoded_item = dict(_JSON_OBJECT.validate_python(item))
+ decoded_item["embedding"] = self._decode_base64_embedding(decoded_item.get("embedding"))
decoded_data.append(decoded_item)
model_response.data = decoded_data
- usage_data: Final = raw_response_json.get("usage", {})
+ usage_data: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {}))
usage: Final = Usage(
- prompt_tokens=usage_data.get("prompt_tokens", 0) or usage_data.get("total_tokens", 0),
- total_tokens=usage_data.get("total_tokens", 0),
+ prompt_tokens=_INT.validate_python(
+ usage_data.get("prompt_tokens", 0) or usage_data.get("total_tokens", 0) or 0
+ ),
+ total_tokens=_INT.validate_python(usage_data.get("total_tokens", 0) or 0),
)
model_response.usage = usage
return model_response
diff --git a/litellm/llms/perplexity/search/transformation.py b/litellm/llms/perplexity/search/transformation.py
index 98299468839..5cfbdf207bb 100644
--- a/litellm/llms/perplexity/search/transformation.py
+++ b/litellm/llms/perplexity/search/transformation.py
@@ -2,9 +2,11 @@
Calls Perplexity's /search endpoint to search the web.
"""
+from collections.abc import Iterable, Mapping
from typing import Final, TypedDict
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.search.transformation import (
@@ -14,6 +16,9 @@ from litellm.llms.base_llm.search.transformation import (
)
from litellm.secret_managers.main import get_secret_str
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
+
class _PerplexitySearchRequestRequired(TypedDict):
"""Required fields for Perplexity Search API request."""
@@ -146,17 +151,20 @@ class PerplexitySearchConfig(BaseSearchConfig):
Returns:
SearchResponse with standardized format
"""
- response_json: Final = raw_response.json()
+ response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
# Transform results to SearchResult objects
results: Final = []
- for result in response_json.get("results", []):
- search_result = SearchResult(
- title=result.get("title", ""),
- url=result.get("url", ""),
- snippet=result.get("snippet", ""),
- date=result.get("date"),
- last_updated=result.get("last_updated"),
+ for raw_result in _JSON_VALUES.validate_python(response_json.get("results", [])):
+ result = _JSON_OBJECT.validate_python(raw_result)
+ search_result = SearchResult.model_validate(
+ {
+ "title": result.get("title", ""),
+ "url": result.get("url", ""),
+ "snippet": result.get("snippet", ""),
+ "date": result.get("date"),
+ "last_updated": result.get("last_updated"),
+ }
)
results.append(search_result)
diff --git a/litellm/llms/replicate/chat/handler.py b/litellm/llms/replicate/chat/handler.py
index e4233ca725d..ad267847e81 100644
--- a/litellm/llms/replicate/chat/handler.py
+++ b/litellm/llms/replicate/chat/handler.py
@@ -1,7 +1,7 @@
import asyncio
import json
import time
-from collections.abc import Callable
+from collections.abc import Callable, Iterable
from typing import Final
from pydantic import ConfigDict, TypeAdapter
@@ -23,6 +23,8 @@ from .transformation import ReplicateConfig
replicate_config: Final = ReplicateConfig()
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
+_POLLED_PREDICTION: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True))
+_TEXTS: Final = TypeAdapter(Iterable[str], config=ConfigDict(strict=True, hide_input_in_errors=True))
# Function to handle prediction response (streaming)
@@ -38,11 +40,11 @@ def handle_prediction_response_streaming(
print_verbose(f"replicate: polling endpoint: {prediction_url}")
response = http_client.get(prediction_url, headers=headers)
if response.status_code == 200:
- response_data = response.json()
+ response_data = _POLLED_PREDICTION.validate_python(response.json())
status = response_data["status"]
if "output" in response_data:
try:
- output_string = "".join(response_data["output"])
+ output_string = "".join(_TEXTS.validate_python(response_data["output"]))
except Exception:
raise ReplicateError(
status_code=422,
diff --git a/litellm/llms/replicate/chat/transformation.py b/litellm/llms/replicate/chat/transformation.py
index e8dd8816997..42e9eb2c3a6 100644
--- a/litellm/llms/replicate/chat/transformation.py
+++ b/litellm/llms/replicate/chat/transformation.py
@@ -1,4 +1,4 @@
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
from typing import TYPE_CHECKING, Any, Final
import httpx
@@ -29,6 +29,7 @@ else:
LoggingClass = Any
_TEXTS: Final = TypeAdapter(Iterable[str], config=ConfigDict(strict=True, hide_input_in_errors=True))
+_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
class ReplicateConfig(BaseConfig):
@@ -79,7 +80,7 @@ class ReplicateConfig(BaseConfig):
seed: int | None = None,
debug: bool | None = None,
) -> None:
- locals_: Final = locals().copy()
+ locals_: Final[Mapping[str, object]] = dict(locals())
for key, value in locals_.items():
if key != "self" and value is not None:
setattr(self.__class__, key, value)
@@ -250,7 +251,7 @@ class ReplicateConfig(BaseConfig):
original_response=raw_response.text,
additional_args={"complete_input_dict": request_data},
)
- raw_response_json: Final = raw_response.json()
+ raw_response_json: Final = _JSON_DICT.validate_python(raw_response.json())
if raw_response_json.get("status") != "succeeded":
raise ReplicateError(
status_code=422,
diff --git a/litellm/llms/sambanova/embedding/transformation.py b/litellm/llms/sambanova/embedding/transformation.py
index 5483f73cf71..cf39c0a092f 100644
--- a/litellm/llms/sambanova/embedding/transformation.py
+++ b/litellm/llms/sambanova/embedding/transformation.py
@@ -3,9 +3,11 @@ This is OpenAI compatible - no transformation is applied
"""
+from collections.abc import Mapping
from typing import Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
@@ -16,6 +18,9 @@ from litellm.types.utils import EmbeddingResponse, Usage
from ..common_utils import SambaNovaError
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_INT: Final = TypeAdapter(int)
+
class SambaNovaEmbeddingConfig(BaseEmbeddingConfig):
def __init__(self) -> None:
@@ -124,9 +129,10 @@ class SambaNovaEmbeddingConfig(BaseEmbeddingConfig):
model_response.data = raw_response_json.get("data")
model_response.object = raw_response_json.get("object")
+ usage_payload: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {}))
usage: Final = Usage(
- prompt_tokens=raw_response_json.get("usage", {}).get("prompt_tokens", 0),
- total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0),
+ prompt_tokens=_INT.validate_python(usage_payload.get("prompt_tokens", 0) or 0),
+ total_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0),
)
model_response.usage = usage
diff --git a/litellm/llms/sap/chat/transformation.py b/litellm/llms/sap/chat/transformation.py
index 342a4ff7025..34e8acd766a 100755
--- a/litellm/llms/sap/chat/transformation.py
+++ b/litellm/llms/sap/chat/transformation.py
@@ -7,6 +7,7 @@ from functools import cached_property
from typing import TYPE_CHECKING, Any, Final, Union
import httpx
+from pydantic import BaseModel
import litellm
from litellm.types.llms.openai import AllMessageValues
@@ -54,7 +55,7 @@ _SAP_MODEL_PARAMS_EXCLUDED_KEYS: Final[frozenset[str]] = frozenset(
def validate_dict(
- data: dict, model: type
+ data: dict, model: type[BaseModel]
) -> dict: # mutable-ok: pydantic validation boundary; both input and output are untyped wire dicts
return model(**data).model_dump(by_alias=True, exclude_unset=True)
diff --git a/litellm/llms/searchapi/search/transformation.py b/litellm/llms/searchapi/search/transformation.py
index ab4a68cef7f..088b807d02a 100644
--- a/litellm/llms/searchapi/search/transformation.py
+++ b/litellm/llms/searchapi/search/transformation.py
@@ -4,10 +4,12 @@ Calls SearchAPI.io's Google Search API endpoint.
SearchAPI.io API Reference: https://www.searchapi.io/docs/google
"""
+from collections.abc import Iterable, Mapping
from typing import Final, Literal, TypedDict, cast
from urllib.parse import urlencode
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.search.transformation import (
@@ -17,6 +19,9 @@ from litellm.llms.base_llm.search.transformation import (
)
from litellm.secret_managers.main import get_secret_str
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
+
class _SearchAPIRequestRequired(TypedDict):
"""Required fields for SearchAPI.io request."""
@@ -213,24 +218,27 @@ class SearchAPIConfig(BaseSearchConfig):
- organic_results[].snippet → SearchResult.snippet
- organic_results[].date → SearchResult.date
"""
- response_json: Final = raw_response.json()
+ response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
# Transform results to SearchResult objects
results: Final[list[SearchResult]] = []
# Process organic results
- for result in response_json.get("organic_results", []):
+ for raw_result in _JSON_VALUES.validate_python(response_json.get("organic_results", [])):
+ result = _JSON_OBJECT.validate_python(raw_result)
title = result.get("title", "")
url = result.get("link", "")
snippet = result.get("snippet", "")
date = result.get("date") # SearchAPI.io provides date in some results
- search_result = SearchResult(
- title=title,
- url=url,
- snippet=snippet,
- date=date,
- last_updated=None, # SearchAPI.io doesn't provide last_updated
+ search_result = SearchResult.model_validate(
+ {
+ "title": title,
+ "url": url,
+ "snippet": snippet,
+ "date": date,
+ "last_updated": None, # SearchAPI.io doesn't provide last_updated
+ }
)
results.append(search_result)
diff --git a/litellm/llms/searxng/search/transformation.py b/litellm/llms/searxng/search/transformation.py
index 4afb8b32e4b..265a3db72f5 100644
--- a/litellm/llms/searxng/search/transformation.py
+++ b/litellm/llms/searxng/search/transformation.py
@@ -4,9 +4,11 @@ Calls SearXNG's /search endpoint to search the web.
SearXNG API Reference: https://docs.searxng.org/dev/search_api.html
"""
+from collections.abc import Iterable, Mapping
from typing import Final, TypedDict
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.search.transformation import (
@@ -16,6 +18,9 @@ from litellm.llms.base_llm.search.transformation import (
)
from litellm.secret_managers.main import get_secret_str
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
+
class _SearXNGSearchRequestRequired(TypedDict):
"""Required fields for SearXNG Search API request."""
@@ -204,22 +209,25 @@ class SearXNGSearchConfig(BaseSearchConfig):
Returns:
SearchResponse with standardized format
"""
- response_json: Final = raw_response.json()
+ response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
# Transform results to SearchResult objects
# Note: SearXNG doesn't natively support limiting results via API params
# It returns ~20 results per page by default
results: Final = []
- for result in response_json.get("results", []):
+ for raw_result in _JSON_VALUES.validate_python(response_json.get("results", [])):
+ result = _JSON_OBJECT.validate_python(raw_result)
# Get date from either publishedDate or pubdate field
date = result.get("publishedDate") or result.get("pubdate")
- search_result = SearchResult(
- title=result.get("title", ""),
- url=result.get("url", ""),
- snippet=result.get("content", ""), # SearXNG uses "content" for snippet
- date=date,
- last_updated=None, # SearXNG doesn't provide last_updated in response
+ search_result = SearchResult.model_validate(
+ {
+ "title": result.get("title", ""),
+ "url": result.get("url", ""),
+ "snippet": result.get("content", ""), # SearXNG uses "content" for snippet
+ "date": date,
+ "last_updated": None, # SearXNG doesn't provide last_updated in response
+ }
)
results.append(search_result)
diff --git a/litellm/llms/serper/search/transformation.py b/litellm/llms/serper/search/transformation.py
index 54d5a1ff7fa..2539610ea59 100644
--- a/litellm/llms/serper/search/transformation.py
+++ b/litellm/llms/serper/search/transformation.py
@@ -4,9 +4,11 @@ Calls Serper's /search endpoint to search Google.
Serper API Reference: https://serper.dev
"""
+from collections.abc import Iterable, Mapping
from typing import Final, TypedDict
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.search.transformation import (
@@ -16,6 +18,9 @@ from litellm.llms.base_llm.search.transformation import (
)
from litellm.secret_managers.main import get_secret_str
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
+
class _SerperSearchRequestRequired(TypedDict):
"""Required fields for Serper Search API request."""
@@ -156,16 +161,19 @@ class SerperSearchConfig(BaseSearchConfig):
Returns:
SearchResponse with standardized format
"""
- response_json: Final = raw_response.json()
+ response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
results: Final = []
- for result in response_json.get("organic", []):
- search_result = SearchResult(
- title=result.get("title", ""),
- url=result.get("link", ""),
- snippet=result.get("snippet", ""),
- date=result.get("date"),
- last_updated=None,
+ for raw_result in _JSON_VALUES.validate_python(response_json.get("organic", [])):
+ result = _JSON_OBJECT.validate_python(raw_result)
+ search_result = SearchResult.model_validate(
+ {
+ "title": result.get("title", ""),
+ "url": result.get("link", ""),
+ "snippet": result.get("snippet", ""),
+ "date": result.get("date"),
+ "last_updated": None,
+ }
)
results.append(search_result)
diff --git a/litellm/llms/tavily/search/transformation.py b/litellm/llms/tavily/search/transformation.py
index 158346051f6..8643322e1a8 100644
--- a/litellm/llms/tavily/search/transformation.py
+++ b/litellm/llms/tavily/search/transformation.py
@@ -4,9 +4,11 @@ Calls Tavily's /search endpoint to search the web.
Tavily API Reference: https://docs.tavily.com/documentation/api-reference/endpoint/search
"""
+from collections.abc import Iterable, Mapping
from typing import Final, TypedDict
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.search.transformation import (
@@ -16,6 +18,9 @@ from litellm.llms.base_llm.search.transformation import (
)
from litellm.secret_managers.main import get_secret_str
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
+
class _TavilySearchRequestRequired(TypedDict):
"""Required fields for Tavily Search API request."""
@@ -174,17 +179,20 @@ class TavilySearchConfig(BaseSearchConfig):
Returns:
SearchResponse with standardized format
"""
- response_json: Final = raw_response.json()
+ response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
# Transform results to SearchResult objects
results: Final = []
- for result in response_json.get("results", []):
- search_result = SearchResult(
- title=result.get("title", ""),
- url=result.get("url", ""),
- snippet=result.get("content", ""), # Tavily uses "content" instead of "snippet"
- date=None, # Tavily doesn't provide date in response
- last_updated=None, # Tavily doesn't provide last_updated in response
+ for raw_result in _JSON_VALUES.validate_python(response_json.get("results", [])):
+ result = _JSON_OBJECT.validate_python(raw_result)
+ search_result = SearchResult.model_validate(
+ {
+ "title": result.get("title", ""),
+ "url": result.get("url", ""),
+ "snippet": result.get("content", ""), # Tavily uses "content" instead of "snippet"
+ "date": None, # Tavily doesn't provide date in response
+ "last_updated": None, # Tavily doesn't provide last_updated in response
+ }
)
results.append(search_result)
diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py
index 09ec2ba6ee4..7f3ef479ab5 100644
--- a/litellm/llms/vertex_ai/files/transformation.py
+++ b/litellm/llms/vertex_ai/files/transformation.py
@@ -672,7 +672,7 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_rows(
def _openai_batch_jsonl_entry_to_vertex_rows(
openai_entry: dict[str, Any],
- map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, object]],
+ map_openai_to_vertex_params: Callable[[dict[str, object]], dict[str, object]],
) -> tuple[Mapping[str, object], ...]:
"""
Transforms a single OpenAI JSONL batch entry into the Vertex rows it maps to.
diff --git a/litellm/llms/vertex_ai/fine_tuning/handler.py b/litellm/llms/vertex_ai/fine_tuning/handler.py
index c79b6ffce43..e1acea53192 100644
--- a/litellm/llms/vertex_ai/fine_tuning/handler.py
+++ b/litellm/llms/vertex_ai/fine_tuning/handler.py
@@ -5,6 +5,7 @@ from datetime import datetime
from typing import Final, Literal
import httpx
+from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_logger
@@ -23,6 +24,8 @@ from litellm.types.llms.vertex_ai import (
)
from litellm.types.utils import LiteLLMFineTuningJob
+_JSON_VALUE: Final = TypeAdapter(object)
+
class VertexFineTuningAPI(VertexLLM):
"""
@@ -341,4 +344,4 @@ class VertexFineTuningAPI(VertexLLM):
f"Error creating fine tuning job. Status code: {response.status_code}. Response: {response.text}"
)
- return response.json()
+ return _JSON_VALUE.validate_python(response.json())
diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py
index fddc075bfc6..f916d494aad 100644
--- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py
+++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py
@@ -1,13 +1,14 @@
import base64
import json
import os
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
from io import BufferedRandom, BufferedReader, BytesIO
from pathlib import Path
from typing import TYPE_CHECKING, Any, Final, cast
import httpx
from httpx._types import RequestFiles
+from pydantic import ConfigDict, TypeAdapter
import litellm
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
@@ -26,6 +27,11 @@ if TYPE_CHECKING:
else:
LiteLLMLoggingObj = Any
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object]]] = TypeAdapter(
+ Mapping[str, object] | str | Sequence[object], config=ConfigDict(hide_input_in_errors=True)
+)
+
class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
"""
@@ -208,15 +214,18 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
headers=raw_response.headers,
)
- predictions: Final = response_json.get("predictions", [])
+ predictions: Final = _JSON_CONTAINER.validate_python(response_json.get("predictions", []))
data_list: Final[list[ImageObject]] = []
- for prediction in predictions:
+ for raw_prediction in predictions:
+ prediction = (
+ raw_prediction if isinstance(raw_prediction, str) else _JSON_CONTAINER.validate_python(raw_prediction)
+ )
# Imagen returns images as bytesBase64Encoded
if "bytesBase64Encoded" in prediction:
data_list.append(
ImageObject(
- b64_json=prediction["bytesBase64Encoded"],
+ b64_json=_JSON_OBJECT.validate_python(prediction)["bytesBase64Encoded"],
url=None,
)
)
diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py
index 9f3126fd4ff..916e25a027e 100644
--- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py
+++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py
@@ -1,5 +1,5 @@
import os
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final
import httpx
@@ -35,6 +35,9 @@ else:
_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object]]] = TypeAdapter(
+ Mapping[str, object] | str | Sequence[object], config=ConfigDict(hide_input_in_errors=True)
+)
class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
@@ -308,25 +311,27 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
model_response.data = []
# Gemini image generation models return in candidates format
- candidates: Final = response_data.get("candidates", [])
+ response_object: Final = _JSON_OBJECT.validate_python(response_data)
+ candidates: Final = _JSON_CONTAINER.validate_python(response_object.get("candidates", []))
for candidate in candidates:
- content = candidate.get("content", {})
- parts = content.get("parts", [])
- for part in parts:
+ content = _JSON_OBJECT.validate_python(_JSON_OBJECT.validate_python(candidate).get("content", {}))
+ parts = _JSON_CONTAINER.validate_python(content.get("parts", []))
+ for raw_part in parts:
+ part = raw_part if isinstance(raw_part, str) else _JSON_CONTAINER.validate_python(raw_part)
# Look for inlineData with image
if "inlineData" in part:
- inline_data = part["inlineData"]
+ part_object = _JSON_OBJECT.validate_python(part)
+ inline_data = _JSON_CONTAINER.validate_python(part_object["inlineData"])
if "data" in inline_data:
- thought_sig = part.get("thoughtSignature")
+ thought_sig = part_object.get("thoughtSignature")
model_response.data.append(
ImageObject(
- b64_json=inline_data["data"],
+ b64_json=_JSON_OBJECT.validate_python(inline_data)["data"],
url=None,
provider_specific_fields=({"thought_signature": thought_sig} if thought_sig else None),
)
)
- response_object: Final = _JSON_OBJECT.validate_python(response_data)
if usage_metadata := response_object.get("usageMetadata", None):
model_response.usage = self._transform_image_usage(_JSON_DICT.validate_python(usage_metadata))
diff --git a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py
index ae6e08611ae..927c7b74e67 100644
--- a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py
+++ b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py
@@ -1,7 +1,9 @@
import os
+from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
import litellm
from litellm.llms.base_llm.image_generation.transformation import (
@@ -27,6 +29,12 @@ if TYPE_CHECKING:
else:
LiteLLMLoggingObj = Any
+_JSON_VALUE: Final = TypeAdapter(object)
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object]]] = TypeAdapter(
+ Mapping[str, object] | str | Sequence[object], config=ConfigDict(hide_input_in_errors=True)
+)
+
class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
"""
@@ -221,7 +229,7 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
Transform Imagen image generation response to litellm ImageResponse format
"""
try:
- response_data: Final = raw_response.json()
+ response_data: Final = _JSON_VALUE.validate_python(raw_response.json())
except Exception as e:
raise self.get_error_class(
error_message=f"Error transforming image generation response: {e}",
@@ -233,13 +241,18 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
model_response.data = []
# Imagen format - predictions with generated images
- predictions: Final = response_data.get("predictions", [])
- for prediction in predictions:
+ predictions: Final = _JSON_CONTAINER.validate_python(
+ _JSON_OBJECT.validate_python(response_data).get("predictions", [])
+ )
+ for raw_prediction in predictions:
+ prediction = (
+ raw_prediction if isinstance(raw_prediction, str) else _JSON_CONTAINER.validate_python(raw_prediction)
+ )
# Imagen returns images as bytesBase64Encoded
if "bytesBase64Encoded" in prediction:
model_response.data.append(
ImageObject(
- b64_json=prediction["bytesBase64Encoded"],
+ b64_json=_JSON_OBJECT.validate_python(prediction)["bytesBase64Encoded"],
url=None,
)
)
diff --git a/litellm/llms/vertex_ai/text_to_speech/text_to_speech_handler.py b/litellm/llms/vertex_ai/text_to_speech/text_to_speech_handler.py
index d053f564dee..655ae215f76 100644
--- a/litellm/llms/vertex_ai/text_to_speech/text_to_speech_handler.py
+++ b/litellm/llms/vertex_ai/text_to_speech/text_to_speech_handler.py
@@ -1,6 +1,8 @@
+from collections.abc import Mapping
from typing import Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
from typing_extensions import TypedDict
import litellm
@@ -12,6 +14,9 @@ from litellm.llms.openai.openai import HttpxBinaryResponseContent
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_STR: Final = TypeAdapter(str, config=ConfigDict(hide_input_in_errors=True))
+
class VertexInput(TypedDict, total=False):
text: str | None
@@ -150,12 +155,12 @@ class VertexTextToSpeechAPI(VertexLLM):
if response.status_code != 200:
raise Exception(f"Request failed with status code {response.status_code}, {response.text}")
############ Process the response ############
- _json_response: Final = response.json()
+ _json_response: Final = _JSON_OBJECT.validate_python(response.json())
response_content: Final = _json_response["audioContent"]
# Decode base64 to get binary content
- binary_data: Final = base64.b64decode(response_content)
+ binary_data: Final = base64.b64decode(_STR.validate_python(response_content))
# Create an httpx.Response object
response = httpx.Response(
@@ -187,12 +192,12 @@ class VertexTextToSpeechAPI(VertexLLM):
if response.status_code != 200:
raise Exception(f"Request did not return a 200 status code: {response.status_code}, {response.text}")
- _json_response: Final = response.json()
+ _json_response: Final = _JSON_OBJECT.validate_python(response.json())
response_content: Final = _json_response["audioContent"]
# Decode base64 to get binary content
- binary_data: Final = base64.b64decode(response_content)
+ binary_data: Final = base64.b64decode(_STR.validate_python(response_content))
# Create an httpx.Response object
response = httpx.Response(
diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py
index a78145fdcf4..2fae2a9de1c 100644
--- a/litellm/llms/vertex_ai/text_to_speech/transformation.py
+++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py
@@ -47,6 +47,7 @@ else:
_LyriaVoice: TypeAlias = str | dict | None
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_ARRAY: Final = TypeAdapter(tuple[object, ...], config=ConfigDict(hide_input_in_errors=True))
_STR: Final = TypeAdapter(str, config=ConfigDict(hide_input_in_errors=True))
@@ -650,27 +651,37 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
) -> "HttpxBinaryResponseContent":
from litellm.types.llms.openai import HttpxBinaryResponseContent
- response_json: Final = raw_response.json()
+ response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
base_model: Final = model.removeprefix("vertex_ai/")
model_info: Final = self._get_model_info(model=model)
- audio_data: str | None = None # rebind-ok: response parsing discovers audio data in provider-specific shapes
- mime_type: str | None = None # rebind-ok: response parsing discovers the MIME type beside the audio payload
+ audio_data: object = None # rebind-ok: response parsing discovers audio data in provider-specific shapes
+ mime_type: object = None # rebind-ok: response parsing discovers the MIME type beside the audio payload
if model_info["vertex_ai_audio_api"] == "lyria_predict":
- predictions: Final = response_json.get("predictions") or ()
+ predictions: Final = _JSON_ARRAY.validate_python(response_json.get("predictions") or ())
if predictions:
- audio_data = predictions[0].get("audioContent") or predictions[0].get("bytesBase64Encoded")
- mime_type = predictions[0].get("mimeType") # rebind-ok: predict response supplies its audio MIME type
+ prediction: Final = _JSON_OBJECT.validate_python(predictions[0])
+ audio_data = prediction.get("audioContent") or prediction.get("bytesBase64Encoded")
+ mime_type = prediction.get("mimeType") # rebind-ok: predict response supplies its audio MIME type
else:
- for step in response_json.get("steps") or response_json.get("outputs") or ():
- content_items = step.get("content") or () if step.get("type") == "model_output" else (step,)
+ steps: Final = _JSON_ARRAY.validate_python(response_json.get("steps") or response_json.get("outputs") or ())
+ for step in map(_JSON_OBJECT.validate_python, steps):
+ content_items = (
+ map(_JSON_OBJECT.validate_python, _JSON_ARRAY.validate_python(step.get("content") or ()))
+ if step.get("type") == "model_output"
+ else (step,)
+ )
for content in content_items:
if content.get("type") == "audio" and content.get("data"):
audio_data = content["data"]
mime_type = content.get("mime_type")
if audio_data is None:
raise ValueError(f"No generated audio found in Vertex AI {base_model} response")
- binary_data: Final = base64.b64decode(audio_data)
- media_type: Final = mime_type or speech_media_type_from_audio_bytes(binary_data) or DEFAULT_SPEECH_MEDIA_TYPE
+ binary_data: Final = base64.b64decode(_STR.validate_python(audio_data))
+ media_type: Final = (
+ (_STR.validate_python(mime_type) if mime_type else None)
+ or speech_media_type_from_audio_bytes(binary_data)
+ or DEFAULT_SPEECH_MEDIA_TYPE
+ )
return HttpxBinaryResponseContent(
httpx.Response(
status_code=raw_response.status_code,
diff --git a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py
index 9626c879dd0..f1afaa7be59 100644
--- a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py
+++ b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py
@@ -1,6 +1,7 @@
from typing import Final, Literal
import httpx
+from pydantic import ConfigDict, TypeAdapter
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
@@ -17,6 +18,9 @@ from litellm.types.utils import EmbeddingResponse
from .types import *
+_JSON_VALUE: Final = TypeAdapter(object)
+_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
+
class VertexEmbedding(VertexBase):
def __init__(self) -> None:
@@ -123,12 +127,12 @@ class VertexEmbedding(VertexBase):
except httpx.TimeoutException:
raise VertexAIError(status_code=408, message="Timeout error occurred.")
- _json_response: Final = response.json()
+ _json_response: Final = _JSON_VALUE.validate_python(response.json())
## LOGGING POST-CALL
logging_obj.post_call(input=input, api_key=None, original_response=_json_response)
model_response = litellm.vertexAITextEmbeddingConfig.transform_vertex_response_to_openai(
- response=_json_response, model=model, model_response=model_response
+ response=_JSON_DICT.validate_python(_json_response), model=model, model_response=model_response
)
return model_response
@@ -215,12 +219,12 @@ class VertexEmbedding(VertexBase):
except httpx.TimeoutException:
raise VertexAIError(status_code=408, message="Timeout error occurred.")
- _json_response: Final = response.json()
+ _json_response: Final = _JSON_VALUE.validate_python(response.json())
## LOGGING POST-CALL
logging_obj.post_call(input=input, api_key=None, original_response=_json_response)
model_response = litellm.vertexAITextEmbeddingConfig.transform_vertex_response_to_openai(
- response=_json_response, model=model, model_response=model_response
+ response=_JSON_DICT.validate_python(_json_response), model=model, model_response=model_response
)
return model_response
diff --git a/litellm/llms/voyage/embedding/transformation.py b/litellm/llms/voyage/embedding/transformation.py
index a9232ad27e2..53220d59f51 100644
--- a/litellm/llms/voyage/embedding/transformation.py
+++ b/litellm/llms/voyage/embedding/transformation.py
@@ -1,6 +1,8 @@
+from collections.abc import Mapping
from typing import Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
@@ -9,6 +11,9 @@ from litellm.llms.voyage.common_utils import get_default_base_url, get_voyage_ap
from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues
from litellm.types.utils import EmbeddingResponse, Usage
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_INT: Final = TypeAdapter(int)
+
class VoyageError(BaseLLMException):
def __init__(
@@ -123,9 +128,10 @@ class VoyageEmbeddingConfig(BaseEmbeddingConfig):
model_response.data = raw_response_json.get("data")
model_response.object = raw_response_json.get("object")
+ usage_payload: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {}))
usage: Final = Usage(
- prompt_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0),
- total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0),
+ prompt_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0),
+ total_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0),
)
model_response.usage = usage
return model_response
diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py
index f5e7fdb48f5..231efb596d8 100644
--- a/litellm/llms/voyage/embedding/transformation_contextual.py
+++ b/litellm/llms/voyage/embedding/transformation_contextual.py
@@ -7,6 +7,7 @@ from collections.abc import Mapping
from typing import Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
@@ -15,6 +16,9 @@ from litellm.llms.voyage.common_utils import get_default_base_url, get_voyage_ap
from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues
from litellm.types.utils import EmbeddingResponse, Usage
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_INT: Final = TypeAdapter(int)
+
class VoyageError(BaseLLMException):
def __init__(
@@ -149,9 +153,10 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
model_response.data = raw_response_json.get("data")
model_response.object = raw_response_json.get("object")
+ usage_payload: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {}))
usage: Final = Usage(
- prompt_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0),
- total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0),
+ prompt_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0),
+ total_tokens=_INT.validate_python(usage_payload.get("total_tokens", 0) or 0),
)
model_response.usage = usage
return model_response
diff --git a/litellm/llms/voyage/embedding/transformation_multimodal.py b/litellm/llms/voyage/embedding/transformation_multimodal.py
index 035765b6691..550c3e87eeb 100644
--- a/litellm/llms/voyage/embedding/transformation_multimodal.py
+++ b/litellm/llms/voyage/embedding/transformation_multimodal.py
@@ -6,9 +6,11 @@ containing content blocks, unlike standard Voyage embeddings which use
/v1/embeddings and a string/list `input` field.
"""
+from collections.abc import Mapping
from typing import Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
@@ -17,6 +19,9 @@ from litellm.llms.voyage.common_utils import get_default_base_url, get_voyage_ap
from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues
from litellm.types.utils import EmbeddingResponse, Usage
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_INT: Final = TypeAdapter(int)
+
class VoyageMultimodalEmbeddingError(BaseLLMException):
def __init__(
@@ -155,8 +160,8 @@ class VoyageMultimodalEmbeddingConfig(BaseEmbeddingConfig):
model_response.data = raw_response_json.get("data")
model_response.object = raw_response_json.get("object")
- usage_payload: Final = raw_response_json.get("usage", {})
- total_tokens: Final = usage_payload.get("total_tokens", 0)
+ usage_payload: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {}))
+ total_tokens: Final = _INT.validate_python(usage_payload.get("total_tokens", 0) or 0)
model_response.usage = Usage(
prompt_tokens=total_tokens,
total_tokens=total_tokens,
diff --git a/litellm/llms/watsonx/completion/transformation.py b/litellm/llms/watsonx/completion/transformation.py
index e01a0ff85e5..5c0416bd6ca 100644
--- a/litellm/llms/watsonx/completion/transformation.py
+++ b/litellm/llms/watsonx/completion/transformation.py
@@ -1,5 +1,5 @@
import time
-from collections.abc import AsyncIterator, Iterator
+from collections.abc import AsyncIterator, Iterator, Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final
@@ -102,7 +102,7 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig):
stream: bool | None = None,
**kwargs,
) -> None:
- locals_: Final = locals().copy()
+ locals_: Final[Mapping[str, object]] = dict(locals())
for key, value in locals_.items():
if key != "self" and value is not None:
setattr(self.__class__, key, value)
diff --git a/litellm/llms/xai/common_utils.py b/litellm/llms/xai/common_utils.py
index cf76a851a86..4843caeac59 100644
--- a/litellm/llms/xai/common_utils.py
+++ b/litellm/llms/xai/common_utils.py
@@ -1,6 +1,8 @@
+from collections.abc import Iterable, Mapping
from typing import Final
import httpx
+from pydantic import ConfigDict, TypeAdapter
import litellm
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo
@@ -9,6 +11,9 @@ from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ProviderSpecificModelInfo
USD_TICKS_PER_DOLLAR: Final = 10_000_000_000
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
+_MODEL_ID: Final = TypeAdapter(str, config=ConfigDict(strict=True, hide_input_in_errors=True))
def xai_reported_cost_in_usd(cost_in_usd_ticks: object) -> float | None:
@@ -97,11 +102,11 @@ class XAIModelInfo(BaseLLMModelInfo):
f"Failed to fetch models from XAI. Status code: {response.status_code}, Response: {response.text}"
)
- models: Final = response.json()["data"]
+ models: Final = _JSON_VALUES.validate_python(_JSON_OBJECT.validate_python(response.json())["data"])
litellm_model_names: Final = []
for model in models:
- stripped_model_name = model["id"]
+ stripped_model_name = _MODEL_ID.validate_python(_JSON_OBJECT.validate_python(model)["id"])
litellm_model_name = "xai/" + stripped_model_name
litellm_model_names.append(litellm_model_name)
return litellm_model_names
diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py
index 74f85d5fa04..d1a8084f4b2 100644
--- a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py
+++ b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py
@@ -142,7 +142,7 @@ class MCPGuardrailTranslationHandler(BaseTranslation):
data: dict[str, Any],
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
- ) -> dict[str, Any]:
+ ) -> dict[str, object]:
mcp_tool_name: Final = data.get("mcp_tool_name") or data.get("name")
mcp_arguments: Final[object] = data.get("mcp_arguments") or data.get("arguments")
mcp_tool_description: Final = data.get("mcp_tool_description") or data.get("description")
@@ -228,9 +228,9 @@ class MCPGuardrailTranslationHandler(BaseTranslation):
response: "CallToolResult",
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
- user_api_key_dict: Any | None = None,
+ user_api_key_dict: object | None = None,
request_data: dict | None = None,
- ) -> Any:
+ ) -> object:
"""Scan the text content of an MCP tool result and write masked text back.
The content list is rewritten in place (only the entries the guardrail
diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
index 6beea5bfd71..a7375b4f17c 100644
--- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
@@ -6023,7 +6023,7 @@ class MCPServerManager:
server=server,
)
- hook_result: Final[dict[str, Any]] = {}
+ hook_result: Final[dict[str, object]] = {}
if proxy_logging_obj is None:
return hook_result
diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py
index 8ef573aa489..d0133739787 100644
--- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py
+++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py
@@ -157,7 +157,7 @@ def _sanitize_path_parameter_value(param_value: object, param_name: str) -> str:
return quote(value_str, safe="")
-def load_openapi_spec(filepath: str) -> dict[str, Any]:
+def load_openapi_spec(filepath: str) -> dict[str, object]:
"""
Sync wrapper. For URL specs, use the shared/custom MCP httpx client.
"""
diff --git a/litellm/proxy/_experimental/mcp_server/tool_registry.py b/litellm/proxy/_experimental/mcp_server/tool_registry.py
index dd25f09cc57..1f1d952f96c 100644
--- a/litellm/proxy/_experimental/mcp_server/tool_registry.py
+++ b/litellm/proxy/_experimental/mcp_server/tool_registry.py
@@ -3,7 +3,7 @@ import json
from collections.abc import Callable, Generator, Mapping
from contextlib import contextmanager
from contextvars import ContextVar
-from typing import TYPE_CHECKING, Any, Final
+from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_logger
from litellm.proxy.types_utils.utils import get_instance_fn
@@ -56,7 +56,7 @@ class MCPToolRegistry:
self,
name: str,
description: str,
- input_schema: dict[str, Any],
+ input_schema: dict[str, object],
handler: Callable,
*,
server_id: str | None = None,
@@ -117,7 +117,7 @@ class MCPToolRegistry:
def load_tools_from_config(
self,
- mcp_tools_config: dict[str, Any] | None = None,
+ mcp_tools_config: dict[str, object] | None = None,
config_file_path: str | None = None,
) -> None:
"""
diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py
index 8c9fdfcf67d..304f282d1f8 100644
--- a/litellm/proxy/agent_endpoints/a2a_endpoints.py
+++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py
@@ -303,7 +303,7 @@ async def _a2a_sse_event_source(
error_body: Final = await resp.aread()
error_event: Mapping[str, object] | None = None
try:
- parsed: Final = json.loads(error_body)
+ parsed: Final = _DECODED_JSON.validate_python(json.loads(error_body))
if isinstance(parsed, dict) and "error" in parsed:
error_event = normalize_a2a_jsonrpc_response(parsed, request_id=request_id)
except Exception:
diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py
index f2c5d066714..d4727cc8d8d 100644
--- a/litellm/proxy/anthropic_endpoints/endpoints.py
+++ b/litellm/proxy/anthropic_endpoints/endpoints.py
@@ -86,7 +86,7 @@ def _anthropic_error_json_response(exc: ProxyException, request: Request) -> JSO
return JSONResponse(status_code=status_code, content=content, headers=exc.headers)
-def _strip_total_tokens_from_anthropic_response(response: Any) -> None:
+def _strip_total_tokens_from_anthropic_response(response: object) -> None:
"""Remove the OpenAI-flavored `usage.total_tokens` field that LiteLLM
injects into Anthropic /v1/messages responses.
diff --git a/litellm/proxy/auth/ip_address_utils.py b/litellm/proxy/auth/ip_address_utils.py
index 5d32cc7eb4c..f1ca2c70f3b 100644
--- a/litellm/proxy/auth/ip_address_utils.py
+++ b/litellm/proxy/auth/ip_address_utils.py
@@ -199,7 +199,7 @@ class IPAddressUtils:
@staticmethod
def is_request_https(
request: Request,
- general_settings: Mapping[str, Any] | None = None,
+ general_settings: Mapping[str, object] | None = None,
) -> bool:
"""
Whether this request's PUBLIC-facing origin is HTTPS, for deciding
diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py
index 3101f19ced2..b70cdb08b96 100644
--- a/litellm/proxy/auth/model_checks.py
+++ b/litellm/proxy/auth/model_checks.py
@@ -2,7 +2,7 @@
## Common checks for /v1/models and `/model/info`
import copy
from collections.abc import Sequence
-from typing import Any, Final
+from typing import Final
import litellm
from litellm._logging import verbose_proxy_logger
@@ -336,8 +336,8 @@ def get_known_models_from_wildcard(wildcard_model: str, litellm_params: LiteLLM_
def expand_wildcard_deployments_for_model_info(
- deployments: list[dict[str, Any]],
-) -> list[dict[str, Any]]:
+ deployments: list[dict[str, object]],
+) -> list[dict[str, object]]:
"""Expand wildcard deployments into one row per known provider model.
PR #30025 changed /model/info to read from llm_router.model_list (correct,
diff --git a/litellm/proxy/auth/trusted_proxy_utils.py b/litellm/proxy/auth/trusted_proxy_utils.py
index 55cf3e98444..04c09abd5dd 100644
--- a/litellm/proxy/auth/trusted_proxy_utils.py
+++ b/litellm/proxy/auth/trusted_proxy_utils.py
@@ -1,4 +1,4 @@
-from typing import Any, Final
+from typing import Final
from fastapi import Request
@@ -22,7 +22,7 @@ def _get_proxy_general_settings() -> dict[str, object]:
def get_trusted_proxy_cidrs(
- general_settings: dict[str, Any] | None = None,
+ general_settings: dict[str, object] | None = None,
) -> list[str]:
"""Operator-configured trusted reverse-proxy CIDRs, normalized to strings.
@@ -48,7 +48,7 @@ def _get_direct_client_ip(request: Request) -> str | None:
def require_trusted_proxy_request(
*,
request: Request,
- general_settings: dict[str, Any] | None = None,
+ general_settings: dict[str, object] | None = None,
feature_name: str,
setting_name: str = TRUSTED_PROXY_RANGES_KEY,
) -> None:
diff --git a/litellm/proxy/caching_routes.py b/litellm/proxy/caching_routes.py
index c621ec6a33a..597080be3b3 100644
--- a/litellm/proxy/caching_routes.py
+++ b/litellm/proxy/caching_routes.py
@@ -40,7 +40,9 @@ def _extract_cache_params() -> dict[str, object]:
return {}
try:
cache_params: Final = vars(litellm.cache.cache)
- cleaned_params: Final = HealthCheckCacheParams(**cache_params).model_dump() if cache_params else {}
+ cleaned_params: Final = (
+ HealthCheckCacheParams.model_validate(dict(**cache_params)).model_dump() if cache_params else {}
+ )
return masker.mask_dict(cleaned_params)
except (AttributeError, TypeError) as e:
verbose_proxy_logger.debug("Error extracting cache params: %s", e)
diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py
index 684449acd0c..9e2e6966ed7 100644
--- a/litellm/proxy/client/cli/commands/auth.py
+++ b/litellm/proxy/client/cli/commands/auth.py
@@ -3,7 +3,7 @@ import sys
import time
import webbrowser
from collections.abc import Callable, Mapping, Sequence
-from typing import Any, Final, TypeVar
+from typing import Any, Final, Protocol, TypeVar
from urllib.parse import urlencode
import click
@@ -96,6 +96,24 @@ class CliContextObj(TypedDict):
api_key_from_token_file: ReadOnly[NotRequired[bool]]
+class _CliContextObjView(Protocol):
+ @property
+ def obj(self) -> CliContextObj: ...
+
+
+class _OptionalCliContextObjView(Protocol):
+ @property
+ def obj(self) -> CliContextObj | None: ...
+
+
+def _cli_context_obj(view: _CliContextObjView) -> CliContextObj:
+ return view.obj
+
+
+def _optional_cli_context_obj(view: _OptionalCliContextObjView) -> CliContextObj | None:
+ return view.obj
+
+
class CliPollData(TypedDict, total=False):
status: str
key: str
@@ -197,7 +215,7 @@ def keychain_unreadable_notice(vault: SecretVault) -> str:
def context_secret_vault(ctx: click.Context) -> SecretVault:
"""Where this invocation reads and writes secret material; injectable through ctx.obj for tests"""
- ctx_obj: Final[CliContextObj | None] = ctx.obj
+ ctx_obj: Final = _optional_cli_context_obj(ctx)
if ctx_obj is None:
return SYSTEM_KEYRING
return ctx_obj.get("secret_vault") or SYSTEM_KEYRING
@@ -880,7 +898,7 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None:
"""Login to LiteLLM proxy using SSO authentication"""
from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER
- ctx_obj: Final[CliContextObj] = ctx.obj
+ ctx_obj: Final = _cli_context_obj(ctx)
base_url: Final = ctx_obj["base_url"]
if config_claude:
settings_path: Final = claude_settings_path(os.environ)
@@ -1030,7 +1048,7 @@ def print_token(ctx: click.Context):
# explicitly pointed us at a server, trust whichever one `lite login`
# actually issued this token for -- that's the whole point of not
# needing a wrapper command.
- ctx_obj: Final[CliContextObj] = ctx.obj
+ ctx_obj: Final = _cli_context_obj(ctx)
issued_for_this_server: Final = token_data.get("base_url") == ctx_obj.get("base_url", "").rstrip("/")
if ctx_obj.get("base_url_explicit") and not issued_for_this_server:
click.echo("Not authenticated for this server. Run 'lite login'.", err=True)
diff --git a/litellm/proxy/client/cli/commands/http.py b/litellm/proxy/client/cli/commands/http.py
index 722c3afac23..529564c685c 100644
--- a/litellm/proxy/client/cli/commands/http.py
+++ b/litellm/proxy/client/cli/commands/http.py
@@ -1,13 +1,32 @@
import json as json_lib
-from typing import Final
+from typing import Final, Protocol
import click
import requests
import rich
+from typing_extensions import ReadOnly, TypedDict
from ...http_client import HTTPClient
+class _CliContext(TypedDict):
+ base_url: ReadOnly[str]
+ api_key: ReadOnly[str | None]
+
+
+class _CliContextView(Protocol):
+ @property
+ def obj(self) -> _CliContext: ...
+
+
+class _JsonBodyView(TypedDict):
+ body: ReadOnly[object]
+
+
+def _cli_context(view: _CliContextView) -> _CliContext:
+ return view.obj
+
+
@click.group()
def http():
"""Make HTTP requests to the LiteLLM proxy server"""
@@ -79,21 +98,23 @@ def request(
# If not JSON, use as raw data
request_data = data
- client: Final = HTTPClient(ctx.obj["base_url"], ctx.obj["api_key"])
+ client: Final = HTTPClient(_cli_context(ctx)["base_url"], _cli_context(ctx)["api_key"])
try:
- response: Final = client.request(
- method=method,
- uri=uri,
- data=request_data,
- json=json_data,
- headers=headers,
- )
- rich.print_json(data=response)
+ response: Final[_JsonBodyView] = {
+ "body": client.request(
+ method=method,
+ uri=uri,
+ data=request_data,
+ json=json_data,
+ headers=headers,
+ )
+ }
+ rich.print_json(data=response["body"])
except requests.exceptions.HTTPError as e:
click.echo(f"Error: HTTP {e.response.status_code}", err=True)
try:
- error_body: Final = e.response.json()
- rich.print_json(data=error_body)
+ error_body: Final[_JsonBodyView] = {"body": e.response.json()}
+ rich.print_json(data=error_body["body"])
except json_lib.JSONDecodeError:
click.echo(e.response.text, err=True)
raise click.Abort()
diff --git a/litellm/proxy/client/cli/commands/keys.py b/litellm/proxy/client/cli/commands/keys.py
index 0ab3d2480d9..558c9139f5b 100644
--- a/litellm/proxy/client/cli/commands/keys.py
+++ b/litellm/proxy/client/cli/commands/keys.py
@@ -2,7 +2,7 @@ import builtins
import json
from collections.abc import Mapping, Sequence
from datetime import datetime
-from typing import Any, Final, Literal
+from typing import Any, Final, Literal, Protocol
import click
import requests
@@ -20,8 +20,9 @@ class _CliContext(TypedDict):
api_key: ReadOnly[str | None]
-class _CliContextView(TypedDict):
- obj: ReadOnly[_CliContext]
+class _CliContextView(Protocol):
+ @property
+ def obj(self) -> _CliContext: ...
class _KeyRowsView(TypedDict):
@@ -32,6 +33,10 @@ class _JsonBodyView(TypedDict):
body: ReadOnly[object]
+def _cli_context(view: _CliContextView) -> _CliContext:
+ return view.obj
+
+
@click.group()
def keys():
"""Manage API keys for the LiteLLM proxy server"""
@@ -74,8 +79,8 @@ def list(
return_full_object: bool,
):
"""List all API keys"""
- context: Final[_CliContextView] = {"obj": ctx.obj}
- client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
+ context: Final = _cli_context(ctx)
+ client: Final = KeysManagementClient(context["base_url"], context["api_key"])
response: Final = client.list(
page=page,
size=size,
@@ -140,8 +145,8 @@ def generate(
config: str | None,
):
"""Generate a new API key"""
- context: Final[_CliContextView] = {"obj": ctx.obj}
- client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
+ context: Final = _cli_context(ctx)
+ client: Final = KeysManagementClient(context["base_url"], context["api_key"])
try:
models_list: Final = [m.strip() for m in models.split(",")] if models else None
aliases_dict: Final = json.loads(aliases) if aliases else None
@@ -177,8 +182,8 @@ def generate(
@click.pass_context
def delete(ctx: click.Context, keys: str | None, key_aliases: str | None):
"""Delete API keys by key or alias"""
- context: Final[_CliContextView] = {"obj": ctx.obj}
- client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
+ context: Final = _cli_context(ctx)
+ client: Final = KeysManagementClient(context["base_url"], context["api_key"])
keys_list: Final = [k.strip() for k in keys.split(",")] if keys else None
aliases_list: Final = [a.strip() for a in key_aliases.split(",")] if key_aliases else None
try:
@@ -377,8 +382,8 @@ def import_keys(
# Create clients for both source and destination
source_client: Final = KeysManagementClient(source_base_url, source_api_key)
- context: Final[_CliContextView] = {"obj": ctx.obj}
- dest_client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
+ context: Final = _cli_context(ctx)
+ dest_client: Final = KeysManagementClient(context["base_url"], context["api_key"])
try:
# Get all keys from source instance with pagination
diff --git a/litellm/proxy/client/cli/commands/models.py b/litellm/proxy/client/cli/commands/models.py
index f2b38c6eab4..ac7d5ee45ea 100644
--- a/litellm/proxy/client/cli/commands/models.py
+++ b/litellm/proxy/client/cli/commands/models.py
@@ -378,7 +378,7 @@ def _print_summary_table(provider_counts):
rich.print(summary_table)
-def get_model_list_from_yaml_file(yaml_file: str) -> list[dict[str, Any]]:
+def get_model_list_from_yaml_file(yaml_file: str) -> list[dict[str, object]]:
"""Load and validate the model list from a YAML file."""
with open(yaml_file, "r") as f:
data: Final = yaml.safe_load(f)
diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py
index 682dd61ac42..cbbf4247624 100644
--- a/litellm/proxy/client/cli/main.py
+++ b/litellm/proxy/client/cli/main.py
@@ -1,7 +1,7 @@
# stdlib imports
# third party imports
-from typing import Final
+from typing import Final, Protocol
import click
@@ -37,6 +37,15 @@ from .commands.users import users
from .interface import interactive_shell
+class _CliContextObjView(Protocol):
+ @property
+ def obj(self) -> CliContextObj: ...
+
+
+def _cli_context_obj(view: _CliContextObjView) -> CliContextObj:
+ return view.obj
+
+
def print_version(base_url: str, api_key: str | None):
"""Print CLI and server version info."""
click.echo(f"LiteLLM Proxy CLI Version: {litellm_version}")
@@ -135,7 +144,7 @@ def cli(ctx: click.Context, show_version: bool, base_url: str | None, api_key: s
@click.pass_context
def version(ctx: click.Context):
"""Show the LiteLLM Proxy CLI and server version."""
- ctx_obj: Final[CliContextObj] = ctx.obj
+ ctx_obj: Final = _cli_context_obj(ctx)
print_version(ctx_obj.get("base_url"), ctx_obj.get("api_key"))
diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py
index 6b87960fcc0..00729c12cc1 100644
--- a/litellm/proxy/common_request_processing.py
+++ b/litellm/proxy/common_request_processing.py
@@ -27,7 +27,7 @@ import httpx
import orjson
from fastapi import HTTPException, Request, status
from fastapi.responses import JSONResponse, Response, StreamingResponse
-from pydantic import BaseModel, TypeAdapter, ValidationError
+from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from starlette.types import Receive, Scope, Send
from typing_extensions import Never
@@ -556,6 +556,7 @@ def _stream_usage_tracking_updates(
}
+@runtime_checkable
class _UpstreamHttpResponse(Protocol):
@property
def status_code(self) -> int: ...
@@ -570,6 +571,11 @@ def _as_upstream_response(response: _UpstreamHttpResponse) -> _UpstreamHttpRespo
return response
+_UPSTREAM_RESPONSE: Final = TypeAdapter(
+ _UpstreamHttpResponse, config=ConfigDict(arbitrary_types_allowed=True, hide_input_in_errors=True)
+)
+
+
class _ReadsHeaderValues(Protocol):
def get(self, key: str, default: str = "") -> str: ...
@@ -2562,7 +2568,7 @@ class ProxyBaseLLMRequestProcessing:
@staticmethod
def _response_cost_from_logging_obj(
*,
- response: Any,
+ response: object,
logging_obj: LiteLLMLoggingObj,
) -> float | str:
"""
@@ -2618,7 +2624,7 @@ class ProxyBaseLLMRequestProcessing:
is_streaming_request: bool | None = False,
contents: list[object] | None = None,
skip_pre_call_logic: bool = False,
- ) -> Any:
+ ) -> object:
"""Run the request, sending SSE keepalives while the upstream is still silent.
Everything below this point, the upstream call included, happens before the
@@ -3261,7 +3267,7 @@ class ProxyBaseLLMRequestProcessing:
if isinstance(result, Response):
return result
- upstream: Final = _as_upstream_response(result)
+ upstream: Final = _UPSTREAM_RESPONSE.validate_python(result)
content: Final = await upstream.aread()
return Response(
content=content,
diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py
index 80d80224617..81ac9eba665 100644
--- a/litellm/proxy/common_utils/http_parsing_utils.py
+++ b/litellm/proxy/common_utils/http_parsing_utils.py
@@ -545,7 +545,7 @@ async def convert_upload_files_to_file_data(
return data
-async def get_request_body(request: Request) -> dict[str, Any]:
+async def get_request_body(request: Request) -> dict[str, object]:
"""
Read the request body and parse it as JSON.
"""
diff --git a/litellm/proxy/common_utils/openapi_schema_compat.py b/litellm/proxy/common_utils/openapi_schema_compat.py
index 903e6c1cb24..1826daa2762 100644
--- a/litellm/proxy/common_utils/openapi_schema_compat.py
+++ b/litellm/proxy/common_utils/openapi_schema_compat.py
@@ -5,7 +5,7 @@ FastAPI 0.120+ has stricter schema generation that fails on certain types like o
This module provides a compatibility layer to handle these cases gracefully.
"""
-from typing import Any, Final
+from typing import Final
from litellm._logging import verbose_proxy_logger
@@ -16,7 +16,7 @@ def get_openapi_schema_with_compat(
version: str,
description: str,
routes: list,
-) -> dict[str, Any]:
+) -> dict[str, object]:
"""
Generate OpenAPI schema with compatibility handling for FastAPI 0.120+.
diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py
index 2a2ef5217b8..64c935ce0d0 100644
--- a/litellm/proxy/guardrails/_content_utils.py
+++ b/litellm/proxy/guardrails/_content_utils.py
@@ -146,7 +146,7 @@ def iter_message_text(data: Mapping[str, object]) -> Iterator[str]:
yield from _iter_text_parts_in_content(message.get("content"))
-def walk_user_text(data: dict[str, Any], visit: Callable[[str], str]) -> int:
+def walk_user_text(data: dict[str, object], visit: Callable[[str], str]) -> int:
"""Rewrite every text fragment in place via ``visit``.
Mutates ``data["messages"]`` and ``data["input"]``. Returns the number
@@ -286,7 +286,7 @@ def has_non_string_content(data: Mapping[str, object]) -> bool:
return False
-def build_inspection_messages(data: dict[str, Any]) -> list[dict[str, str]]:
+def build_inspection_messages(data: dict[str, object]) -> list[dict[str, str]]:
"""Synthesize a chat-style messages list for posting to a guardrail API.
Each returned message has a plain-string ``content`` — multimodal text
diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py
index 3a8bd54c587..50a14375107 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py
@@ -8,7 +8,7 @@
import os
from collections.abc import AsyncGenerator, AsyncIterable
from datetime import datetime
-from typing import Any, Final
+from typing import Final
import httpx
@@ -86,7 +86,7 @@ class DynamoAIGuardrails(CustomGuardrail):
async def _call_dynamoai_guardrails(
self,
- messages: list[dict[str, Any]],
+ messages: list[dict[str, object]],
event_type: GuardrailEventHooks,
text_type: str = "input",
request_data: dict | None = None,
@@ -405,7 +405,7 @@ class DynamoAIGuardrails(CustomGuardrail):
# to avoid sending empty content to DynamoAI (e.g., during tool calls)
if isinstance(response, litellm.ModelResponse):
has_text_content = False
- dynamoai_messages: Final[list[dict[str, str]]] = []
+ dynamoai_messages: Final[list[dict[str, object]]] = []
for choice in response.choices:
if isinstance(choice, litellm.Choices):
diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py
index c745a6c91b3..25df7a20e14 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py
@@ -8,7 +8,7 @@
import fnmatch
import os
from collections.abc import Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Final, Literal, Optional
+from typing import TYPE_CHECKING, Final, Literal, Optional
import httpx
@@ -194,7 +194,7 @@ class GenericGuardrailAPI(CustomGuardrail):
def __init__(
self,
- headers: dict[str, Any] | None = None,
+ headers: dict[str, str] | None = None,
api_base: str | None = None,
api_key: str | None = None,
additional_provider_specific_params: Mapping[str, object] | None = None,
diff --git a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py
index 5da64de329b..380562aacac 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py
@@ -6,7 +6,7 @@
# +-------------------------------------------------------------+
import os
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Mapping
from datetime import datetime
from typing import Any, Final
@@ -40,7 +40,7 @@ class IBMGuardrailDetector(CustomGuardrail):
base_url: str | None = None,
detector_id: str | None = None,
is_detector_server: bool = True,
- detector_params: dict[str, Any] | None = None,
+ detector_params: Mapping[str, object] | None = None,
extra_headers: dict[str, str] | None = None,
score_threshold: float | None = None,
block_on_detection: bool = True,
diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py
index 6ea4317a174..c970ac263e7 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py
@@ -371,8 +371,10 @@ class LLMShieldProxyGuardrail(CustomGuardrail):
return response
pending: Final[SlotSink] = []
+ choice: object
+ tool_call: object
for choice in choices:
- message = getattr(choice, "message", None)
+ message: object = getattr(choice, "message", None)
if message is None:
text = read_field(choice, "text")
if isinstance(text, str) and text:
@@ -382,11 +384,11 @@ class LLMShieldProxyGuardrail(CustomGuardrail):
if isinstance(content, str) and content:
pending.append((content, functools.partial(setattr, message, "content")))
for tool_call in getattr(message, "tool_calls", None) or ():
- function = getattr(tool_call, "function", None)
+ function: object = getattr(tool_call, "function", None)
arguments = getattr(function, "arguments", None) if function is not None else None
if isinstance(arguments, str) and arguments:
pending.append((arguments, functools.partial(setattr, function, "arguments")))
- legacy = getattr(message, "function_call", None)
+ legacy: object = getattr(message, "function_call", None)
legacy_arguments = getattr(legacy, "arguments", None) if legacy is not None else None
if isinstance(legacy_arguments, str) and legacy_arguments:
pending.append((legacy_arguments, functools.partial(setattr, legacy, "arguments")))
@@ -448,6 +450,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail):
function_call_output holds `output` -- so the two directions stay symmetric.
"""
slots: Final[SlotSink] = []
+ item: object
for item in getattr(response, "output", None) or ():
collect_response_item(item, slots)
return tuple(slots)
@@ -501,6 +504,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail):
continue
restored_chunk = detached(chunk)
last_chunk = restored_chunk
+ choice: object
for choice in getattr(restored_chunk, "choices", None) or ():
await self._restore_choice(choice, carries, session_id)
yield restored_chunk
@@ -532,6 +536,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail):
await self._restore_content_window(delta, (index, None), carries, session_id, is_final)
+ tool_call: object
for tool_call in getattr(delta, "tool_calls", None) or ():
await self._restore_tool_call_window(tool_call, index, carries, session_id)
diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py
index a60059fcc4d..0c1be621923 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py
@@ -1375,6 +1375,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
# 2. Process tool calls
tool_calls = getattr(message, "tool_calls", None)
if tool_calls:
+ tool_call: object
for tool_call in tool_calls:
function = getattr(tool_call, "function", None)
if function and hasattr(function, "arguments"):
diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py
index b2139779925..1a571047977 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py
@@ -7,7 +7,7 @@ then builds a SemanticRouter for prompt matching.
import os
from collections.abc import Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Final
+from typing import TYPE_CHECKING, Final
import yaml
@@ -28,7 +28,7 @@ class SemanticGuardRouteLoader:
"""Loads route definitions from YAML templates and custom configs, builds SemanticRouter."""
@staticmethod
- def load_builtin_template(template_name: str) -> dict[str, Any]:
+ def load_builtin_template(template_name: str) -> dict[str, object]:
"""Load a built-in route template YAML by name."""
file_path: Final = os.path.join(ROUTE_TEMPLATES_DIR, f"{template_name}.yaml")
if not os.path.exists(file_path):
@@ -50,7 +50,7 @@ class SemanticGuardRouteLoader:
return sorted(templates)
@staticmethod
- def load_custom_routes_file(file_path: str) -> list[dict[str, Any]]:
+ def load_custom_routes_file(file_path: str) -> list[dict[str, object]]:
"""Load custom routes from a YAML file."""
if not os.path.exists(file_path):
raise ValueError(f"SemanticGuard: custom routes file not found: {file_path}")
diff --git a/litellm/proxy/health_check_utils/shared_health_check_manager.py b/litellm/proxy/health_check_utils/shared_health_check_manager.py
index 764d570fbea..62c973196b9 100644
--- a/litellm/proxy/health_check_utils/shared_health_check_manager.py
+++ b/litellm/proxy/health_check_utils/shared_health_check_manager.py
@@ -192,12 +192,12 @@ class SharedHealthCheckManager:
@with_service_target(HEALTH_CHECKS_TARGET)
async def perform_shared_health_check(
self,
- model_list: list[dict[str, Any]],
+ model_list: list[Mapping[str, object]],
details: bool = True,
max_concurrency: int | None = None,
health_check_skip_disabled_background_models: bool = False,
router: "Router | None" = None,
- ) -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, Any]]:
+ ) -> tuple[list[Mapping[str, object]], list[Mapping[str, object]], dict[str, object]]:
"""
Perform health check with shared state coordination.
diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py
index ba319be01cb..5ea4e85a730 100644
--- a/litellm/proxy/litellm_pre_call_utils.py
+++ b/litellm/proxy/litellm_pre_call_utils.py
@@ -1550,7 +1550,7 @@ class LiteLLMProxyRequestSetup:
headers: dict,
request_data: Mapping[str, object],
user_api_key_dict: UserAPIKeyAuth,
- general_settings: dict[str, Any] | None = None,
+ general_settings: dict[str, object] | None = None,
) -> LitellmDataForBackendLLMCall:
"""
- Adds user from headers
diff --git a/litellm/proxy/mcp_tools.py b/litellm/proxy/mcp_tools.py
index eaf9e4b2665..239bb4b2fcf 100644
--- a/litellm/proxy/mcp_tools.py
+++ b/litellm/proxy/mcp_tools.py
@@ -1,7 +1,4 @@
-from typing import Any
-
-
-def get_current_time(params: dict[str, Any] | None = None) -> str:
+def get_current_time(params: dict[str, object] | None = None) -> str:
"""
Get the current time (hardcoded sample implementation)
@@ -18,7 +15,7 @@ def get_current_time(params: dict[str, Any] | None = None) -> str:
return "10:30:45 AM"
-def get_current_date(params: dict[str, Any] | None = None) -> str:
+def get_current_date(params: dict[str, object] | None = None) -> str:
"""
Get the current date (hardcoded sample implementation)
diff --git a/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py b/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py
index 3d408accba3..457af931f56 100644
--- a/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py
+++ b/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py
@@ -96,7 +96,7 @@ class FileContentStreamingHandler:
*,
custom_llm_provider: str,
file_id: str,
- data: dict[str, Any],
+ data: dict[str, object],
proxy_logging_obj: "ProxyLogging",
user_api_key_dict: "UserAPIKeyAuth",
version: str,
diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
index d787c16a534..23ef89b3c5c 100644
--- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
+++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
@@ -35,6 +35,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
PassthroughStandardLoggingPayload,
)
from litellm.types.utils import (
+ CompletionTokensDetailsWrapper,
LiteLLMBatch,
ModelResponse,
TextCompletionResponse,
@@ -306,13 +307,13 @@ class AnthropicPassthroughLoggingHandler:
if not line.startswith("data:"):
continue
try:
- data = json.loads(line[len("data:") :].strip())
+ data = _JSON_OBJECT.validate_python(json.loads(line[len("data:") :].strip()))
except (json.JSONDecodeError, ValueError):
continue
if not isinstance(data, dict):
continue
if data.get("type") == "message_start":
- value = (data.get("message") or {}).get(field)
+ value = _JSON_OBJECT.validate_python(data.get("message") or {}).get(field)
if isinstance(value, str) and value:
return value
return None
@@ -398,7 +399,7 @@ class AnthropicPassthroughLoggingHandler:
# Anthropic costing reads completion_tokens_details.text_tokens, so the
# stale message_start placeholder there must be corrected too or spend
# stays undercounted even after completion_tokens is fixed.
- details: Final = getattr(usage, "completion_tokens_details", None)
+ details: Final[CompletionTokensDetailsWrapper | None] = getattr(usage, "completion_tokens_details", None)
if details is not None and getattr(details, "text_tokens", None) is not None:
details.text_tokens = recovered_output_tokens
return usage
diff --git a/litellm/proxy/plugin_routes.py b/litellm/proxy/plugin_routes.py
index eb6fe7dd177..1dfe5cbc6bb 100644
--- a/litellm/proxy/plugin_routes.py
+++ b/litellm/proxy/plugin_routes.py
@@ -296,8 +296,8 @@ async def plugin_proxy(
# Forward caller identity so the plugin can enforce its own access control.
# The plugin MUST NOT trust these as credentials — they are informational.
# The plugin_key above is the only authentication mechanism.
- user_id: Final = getattr(user_api_key_dict, "user_id", None)
- user_role: Final = getattr(user_api_key_dict, "user_role", None)
+ user_id: Final[object] = getattr(user_api_key_dict, "user_id", None)
+ user_role: Final[object] = getattr(user_api_key_dict, "user_role", None)
if user_id:
forward_headers["x-litellm-user-id"] = str(user_id)
if user_role:
diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py
index d81471b3c1a..c751091ff9b 100644
--- a/litellm/proxy/policy_engine/attachment_registry.py
+++ b/litellm/proxy/policy_engine/attachment_registry.py
@@ -80,7 +80,7 @@ class AttachmentRegistry:
self._config_attachments: tuple[PolicyAttachment, ...] = ()
self._initialized: bool = False
- def load_attachments(self, attachments_config: list[dict[str, Any]]) -> None:
+ def load_attachments(self, attachments_config: list[dict[str, object]]) -> None:
"""
Load attachments from a configuration list.
diff --git a/litellm/proxy/policy_engine/init_policies.py b/litellm/proxy/policy_engine/init_policies.py
index 061a852b701..b6e84109222 100644
--- a/litellm/proxy/policy_engine/init_policies.py
+++ b/litellm/proxy/policy_engine/init_policies.py
@@ -98,7 +98,7 @@ def _print_policies_on_startup(
async def init_policies(
policies_config: dict[str, Any],
- policy_attachments_config: list[dict[str, Any]] | None = None,
+ policy_attachments_config: list[dict[str, object]] | None = None,
prisma_client: Optional["PrismaClient"] = None,
validate_db: bool = True,
fail_on_error: bool = True,
@@ -184,8 +184,8 @@ async def init_policies(
def init_policies_sync(
- policies_config: dict[str, Any],
- policy_attachments_config: list[dict[str, Any]] | None = None,
+ policies_config: dict[str, object],
+ policy_attachments_config: list[dict[str, object]] | None = None,
fail_on_error: bool = True,
) -> None:
"""
@@ -218,7 +218,7 @@ def init_policies_sync(
)
-def get_policies_summary() -> dict[str, Any]:
+def get_policies_summary() -> dict[str, object]:
"""
Get a summary of loaded policies for debugging/display.
diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py
index cfc1f481988..c5f7e14e61d 100644
--- a/litellm/proxy/policy_engine/policy_registry.py
+++ b/litellm/proxy/policy_engine/policy_registry.py
@@ -22,6 +22,8 @@ from typing import (
cast, # noqa: TID251 # prisma types the condition/pipeline Json columns as str, but reads return decoded values
)
+from typing_extensions import ReadOnly
+
from litellm._logging import verbose_proxy_logger
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import PolicyRepository
@@ -160,6 +162,15 @@ def _row_to_policy_db_response(row: _PolicyRow) -> PolicyDBResponse:
)
+class _AttributeView(TypedDict):
+ value: ReadOnly[object]
+
+
+def _attribute_of(source: object, name: str) -> object:
+ attribute: Final[_AttributeView] = {"value": getattr(source, name)}
+ return attribute["value"]
+
+
class PolicyRegistry:
"""
In-memory registry for storing and managing policies.
@@ -1020,8 +1031,8 @@ class PolicyRegistry:
)
field_diffs: Final[dict[str, dict[str, object]]] = {}
for field in compare_fields:
- val_a = getattr(resp_a, field)
- val_b = getattr(resp_b, field)
+ val_a = _attribute_of(resp_a, field)
+ val_b = _attribute_of(resp_b, field)
if val_a != val_b:
field_diffs[field] = {"version_a": val_a, "version_b": val_b}
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 3a831062b0a..bdaa4af4e26 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -15117,7 +15117,7 @@ async def _apply_search_filter_to_models(
size: int = 50,
sort_by: str | None = None,
model_name: str | None = None,
-) -> tuple[list[dict[str, Any]], int | None]:
+) -> tuple[list[dict[str, object]], int | None]:
"""
Apply search filter to models, querying database for additional matching models.
@@ -15267,10 +15267,10 @@ def _normalize_datetime_for_sorting(dt: object) -> datetime | None:
def _sort_models(
- all_models: list[dict[str, Any]],
+ all_models: list[dict[str, object]],
sort_by: str | None,
sort_order: str = "asc",
-) -> list[dict[str, Any]]:
+) -> list[dict[str, object]]:
"""
Sort models by the specified field and order.
@@ -15560,7 +15560,7 @@ async def _filter_models_by_team_id(
prisma_client: PrismaClient,
llm_router: Router,
user_api_key_dict: UserAPIKeyAuth | None = None,
-) -> list[dict[str, Any]]:
+) -> list[dict[str, object]]:
"""
Filter models by team ID. Returns models where:
- team_id matches the model's BYOK team_id, OR
diff --git a/litellm/proxy/read_model_list.py b/litellm/proxy/read_model_list.py
index a1830e7f2bc..92bfa8657cb 100644
--- a/litellm/proxy/read_model_list.py
+++ b/litellm/proxy/read_model_list.py
@@ -17,10 +17,10 @@ secret manager, which is what makes that resolution work. Returns the resolved
from __future__ import annotations
import asyncio
-from typing import Any, Final
+from typing import Final
-def read_model_list(config_path: str) -> list[dict[str, Any]]:
+def read_model_list(config_path: str) -> list[dict[str, object]]:
"""Load ``config_path`` via the proxy's own reader and return its
resolved ``model_list``."""
from litellm.proxy.proxy_server import ProxyConfig
diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py
index b13042dfb6c..5f6a7432afb 100644
--- a/litellm/proxy/response_polling/background_streaming.py
+++ b/litellm/proxy/response_polling/background_streaming.py
@@ -10,11 +10,10 @@ https://platform.openai.com/docs/api-reference/responses-streaming
import asyncio
import json
-from collections.abc import Callable, Mapping, Sequence
-from typing import TYPE_CHECKING, Final, TypeAlias
+from collections.abc import AsyncIterable, Callable, Mapping, Sequence
+from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, runtime_checkable
from fastapi import Request, Response
-from fastapi.responses import StreamingResponse
from starlette.types import Message
from typing_extensions import ReadOnly, TypedDict
@@ -75,6 +74,12 @@ class _StreamEventParser:
parse: Callable[[str], _StreamEvent] = staticmethod(json.loads)
+@runtime_checkable
+class _StreamsBody(Protocol):
+ @property
+ def body_iterator(self) -> AsyncIterable[str | bytes | memoryview]: ...
+
+
def _sse_frame_data(frame: str) -> str | None:
return next((line[6:].strip() for line in frame.splitlines() if line.startswith("data: ")), None)
@@ -141,7 +146,7 @@ async def background_streaming_task(
# Make streaming request.
# Pre-call checks (rate limits, guardrails, budget) were already run
# before polling ID creation, so skip them here to avoid double-counting.
- response: Final[StreamingResponse] = await processor.base_process_llm_request(
+ response: Final = await processor.base_process_llm_request(
request=detach_request_from_client(request),
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
@@ -216,13 +221,13 @@ async def background_streaming_task(
last_update_time = current_time
# Handle StreamingResponse
- if not hasattr(response, "body_iterator"):
+ if not isinstance(response, _StreamsBody):
verbose_proxy_logger.warning(
"background_streaming_task: response for %s has no body_iterator; this may indicate a misconfiguration or provider error",
polling_id,
)
- if hasattr(response, "body_iterator"):
+ if isinstance(response, _StreamsBody):
async for chunk in response.body_iterator:
# Parse chunk
if isinstance(chunk, bytes):
diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py
index 4ae227162d2..3f82168d958 100644
--- a/litellm/proxy/response_polling/polling_handler.py
+++ b/litellm/proxy/response_polling/polling_handler.py
@@ -3,8 +3,9 @@ Response Polling Handler for Background Responses with Cache
"""
import json
+from collections.abc import Mapping
from datetime import datetime, timezone
-from typing import Any, Final
+from typing import Final
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
@@ -44,7 +45,7 @@ class ResponsePollingHandler:
async def create_initial_state(
self,
polling_id: str,
- request_data: dict[str, Any],
+ request_data: Mapping[str, object],
) -> ResponsesAPIResponse:
"""
Create initial state in Redis for a polling request
@@ -62,14 +63,16 @@ class ResponsePollingHandler:
created_timestamp: Final = int(datetime.now(timezone.utc).timestamp())
# Create OpenAI-compliant response object
- response: Final = ResponsesAPIResponse(
- id=polling_id,
- object="response",
- status="queued", # OpenAI native status
- created_at=created_timestamp,
- output=[],
- metadata=request_data.get("metadata", {}),
- usage=None,
+ response: Final = ResponsesAPIResponse.model_validate(
+ {
+ "id": polling_id,
+ "object": "response",
+ "status": "queued", # OpenAI native status
+ "created_at": created_timestamp,
+ "output": [],
+ "metadata": request_data.get("metadata", {}),
+ "usage": None,
+ }
)
cache_key: Final = self.get_cache_key(polling_id)
@@ -218,7 +221,7 @@ class ResponsePollingHandler:
)
@with_service_target(_RESPONSE_POLLING_TARGET)
- async def get_state(self, polling_id: str) -> dict[str, Any] | None:
+ async def get_state(self, polling_id: str) -> Mapping[str, object] | None:
"""Get current polling state from Redis"""
if not self.redis_cache:
return None
diff --git a/litellm/proxy/types_utils/utils.py b/litellm/proxy/types_utils/utils.py
index c3fd237c0d8..62b87dbbf36 100644
--- a/litellm/proxy/types_utils/utils.py
+++ b/litellm/proxy/types_utils/utils.py
@@ -5,6 +5,12 @@ import os
from collections.abc import Callable
from typing import Any, Final, Literal, get_type_hints
+from typing_extensions import ReadOnly, TypedDict
+
+
+class _AttributeView(TypedDict):
+ value: ReadOnly[object]
+
def get_instance_fn(value: str, config_file_path: str | None = None) -> Any:
module_name = value
@@ -52,7 +58,8 @@ def get_instance_fn(value: str, config_file_path: str | None = None) -> Any:
module = importlib.import_module(module_name)
# Get the instance from the module
- instance: Final[object] = getattr(module, instance_name)
+ attribute: Final[_AttributeView] = {"value": getattr(module, instance_name)}
+ instance: Final[object] = attribute["value"]
return instance
except ImportError as e:
@@ -167,7 +174,8 @@ def _load_instance_from_remote_storage(remote_url: str, config_file_path: str |
spec.loader.exec_module(module)
# Get the instance
- instance: Final[object] = getattr(module, instance_name)
+ attribute: Final[_AttributeView] = {"value": getattr(module, instance_name)}
+ instance: Final[object] = attribute["value"]
# Clean up the temporary file
try:
diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py
index 8363aaee99a..b97c3a949d5 100644
--- a/litellm/proxy/vector_store_endpoints/utils.py
+++ b/litellm/proxy/vector_store_endpoints/utils.py
@@ -5,6 +5,7 @@ from types import MappingProxyType
from typing import Final, Literal
from fastapi import HTTPException, Request
+from pydantic import ConfigDict, TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@@ -21,6 +22,8 @@ from litellm.types.utils import LlmProviders
from litellm.types.vector_stores import LiteLLM_ManagedVectorStore
from litellm.utils import ProviderConfigManager
+_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
+
def _normalize_litellm_params(
vector_store: LiteLLM_ManagedVectorStore,
@@ -29,8 +32,7 @@ def _normalize_litellm_params(
if isinstance(litellm_params, str):
normalized: Final = LiteLLM_ManagedVectorStore(**dict(vector_store))
try:
- parsed: Final = json.loads(litellm_params)
- normalized["litellm_params"] = parsed if isinstance(parsed, dict) else {}
+ normalized["litellm_params"] = _JSON_OBJECT.validate_python(json.loads(litellm_params))
except (TypeError, ValueError):
normalized["litellm_params"] = {}
return normalized
diff --git a/litellm/proxy/video_endpoints/utils.py b/litellm/proxy/video_endpoints/utils.py
index a38226cc253..b2ec2eda761 100644
--- a/litellm/proxy/video_endpoints/utils.py
+++ b/litellm/proxy/video_endpoints/utils.py
@@ -47,7 +47,7 @@ def get_custom_provider_from_data(data: dict[str, Any]) -> str | None:
return None
-def encode_character_id_in_response(response: Any, custom_llm_provider: str, model_id: str | None) -> Any:
+def encode_character_id_in_response(response: Any, custom_llm_provider: str, model_id: str | None) -> object:
if isinstance(response, dict) and response.get("id"):
response["id"] = encode_character_id_with_provider(
character_id=response["id"],
diff --git a/litellm/rag/rag_query.py b/litellm/rag/rag_query.py
index 9325547c17d..9f608dbc99a 100644
--- a/litellm/rag/rag_query.py
+++ b/litellm/rag/rag_query.py
@@ -70,7 +70,7 @@ class RAGQuery:
return None
@staticmethod
- def build_context_message(context_chunks: list[Any]) -> ChatCompletionUserMessage:
+ def build_context_message(context_chunks: Sequence[_SearchResultView | str | None]) -> ChatCompletionUserMessage:
"""
Process search results and build a context message.
"""
diff --git a/litellm/router.py b/litellm/router.py
index cdede46ef5c..97d4b481369 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -3384,7 +3384,7 @@ class Router:
async def _aresponses_streaming_iterator(
self,
response: "BaseResponsesAPIStreamingIterator",
- initial_kwargs: dict[str, Any],
+ initial_kwargs: dict[str, object],
) -> "BaseResponsesAPIStreamingIterator":
"""
Wrap a Responses-API streaming iterator so MidStreamFallbackError
@@ -5596,7 +5596,7 @@ class Router:
async def _aanthropic_messages_streaming_iterator(
self,
response: AsyncIterator[bytes],
- initial_kwargs: dict[str, Any], # mutable-ok: mutated in-place before re-entering the fallback chain
+ initial_kwargs: dict[str, object], # mutable-ok: mutated in-place before re-entering the fallback chain
) -> AsyncIterator[bytes]:
"""
Wrap an anthropic_messages (/v1/messages) streaming response so a
diff --git a/litellm/secret_managers/google_secret_manager.py b/litellm/secret_managers/google_secret_manager.py
index 6d6543cee4b..15e8783b96b 100644
--- a/litellm/secret_managers/google_secret_manager.py
+++ b/litellm/secret_managers/google_secret_manager.py
@@ -1,7 +1,10 @@
import base64
import os
+from collections.abc import Mapping
from typing import Final
+from pydantic import ConfigDict, TypeAdapter
+
import litellm
from litellm._logging import verbose_logger
from litellm.caching.caching import InMemoryCache
@@ -11,6 +14,9 @@ from litellm.llms.custom_httpx.http_handler import get_httpx_client
from litellm.types.proxy.common_proxy_errors import CommonProxyErrors
from litellm.types.secret_managers.main import KeyManagementSystem
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
+_BASE64_TEXT: Final = TypeAdapter(str, config=ConfigDict(strict=True, hide_input_in_errors=True))
+
class GoogleSecretManager(GCSBucketBase):
def __init__(
@@ -86,12 +92,14 @@ class GoogleSecretManager(GCSBucketBase):
)
# Parse the JSON response and return the secret value
- secret_data: Final = response.json()
- _base64_encoded_value: Final = secret_data.get("payload", {}).get("data")
+ secret_data: Final = _JSON_OBJECT.validate_python(response.json())
+ _base64_encoded_value: Final = _JSON_OBJECT.validate_python(secret_data.get("payload", {})).get("data")
# decode the base64 encoded value
if _base64_encoded_value is not None:
- _decoded_value: Final = base64.b64decode(_base64_encoded_value).decode("utf-8")
+ _decoded_value: Final = base64.b64decode(_BASE64_TEXT.validate_python(_base64_encoded_value)).decode(
+ "utf-8"
+ )
self.cache.set_cache(secret_name, _decoded_value) # Cache the retrieved secret
return _decoded_value
diff --git a/litellm/utils.py b/litellm/utils.py
index 3f10d2a3ebd..4733051f4f6 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -5756,7 +5756,7 @@ def get_model_cost_key(potential_key: str) -> str | None:
_get_model_cost_key = get_model_cost_key
-def _get_model_info_from_model_cost(key: str) -> dict[str, Any]:
+def _get_model_info_from_model_cost(key: str) -> dict[str, object]:
return litellm.model_cost[key]
@@ -8525,7 +8525,7 @@ _get_bundled_model_cost_map = get_bundled_model_cost_map
def _get_model_cost_entry_for_provider_config(
model: str,
provider: LlmProviders,
-) -> dict[str, Any]:
+) -> dict[str, object]:
candidate_keys: Final = (model, f"{provider.value}/{model}")
for model_key in candidate_keys:
model_info = litellm.model_cost.get(model_key)
diff --git a/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py b/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py
index 1c87fb7564d..09f09170241 100644
--- a/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py
+++ b/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py
@@ -618,3 +618,53 @@ class TestHandlerIntegration:
api_base=None,
)
assert result is not None
+
+
+class TestStreamingContentType:
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ "response_headers, expected_events",
+ [
+ ([("content-type", "application/json")], [{"jsonrpc": "2.0", "id": "req-001", "result": {}}]),
+ (
+ [("Content-Type", "Application/JSON; charset=utf-8")],
+ [{"jsonrpc": "2.0", "id": "req-001", "result": {}}],
+ ),
+ (
+ [("content-type", "text/plain"), ("CONTENT-TYPE", "application/json")],
+ [{"jsonrpc": "2.0", "id": "req-001", "result": {}}],
+ ),
+ (
+ [("content-type", "application/json"), ("content-type", "text/event-stream")],
+ [{"jsonrpc": "2.0", "id": "req-001", "result": {}}],
+ ),
+ ([], []),
+ ([("content-type", "")], []),
+ ([("content-type", "text/event-stream")], []),
+ ([("x-content-type", "application/json")], []),
+ ([("content-type", "application"), ("content-type", "/json")], []),
+ ],
+ )
+ async def test_streaming_reads_a_single_json_body_only_when_the_content_type_says_json(
+ self, httpx_transport, response_headers, expected_events
+ ):
+ from litellm.a2a_protocol.providers.bedrock_agentcore.config import (
+ BedrockAgentCoreA2AConfig,
+ )
+
+ with respx.mock(assert_all_called=True) as router:
+ router.post(url__regex=r".*/invocations.*").mock(
+ return_value=httpx.Response(
+ 200, headers=response_headers, content=b'{"jsonrpc": "2.0", "id": "req-001", "result": {}}'
+ )
+ )
+ events = [
+ event
+ async for event in BedrockAgentCoreA2AConfig().handle_streaming(
+ request_id="req-001",
+ params=SAMPLE_PARAMS,
+ litellm_params=SAMPLE_LITELLM_PARAMS,
+ )
+ ]
+
+ assert events == expected_events
diff --git a/tests/unit/integrations/focus/test_focus_logger.py b/tests/unit/integrations/focus/test_focus_logger.py
new file mode 100644
index 00000000000..5ff5c40df47
--- /dev/null
+++ b/tests/unit/integrations/focus/test_focus_logger.py
@@ -0,0 +1,92 @@
+import pytest
+
+from litellm.integrations.focus.focus_logger import FocusLogger
+from litellm.proxy import proxy_server
+
+FOCUS_LOCK = "cronjob_lock:focus_export_usage_data"
+VANTAGE_LOCK = "cronjob_lock:vantage_export_usage_data"
+
+
+class _InMemoryRedis:
+ def __init__(self, locks: dict[str, str]) -> None:
+ self.locks = locks
+ self.granted: list[tuple[str, str]] = []
+
+ async def async_set_cache(self, key: str, value: str, nx: bool, ttl: int) -> bool:
+ if nx and key in self.locks:
+ return False
+ self.locks[key] = value
+ self.granted.append((key, value))
+ return True
+
+ async def async_get_cache(self, key: str) -> str | None:
+ return self.locks.get(key)
+
+ async def async_delete_cache(self, key: str) -> int:
+ return 1 if self.locks.pop(key, None) is not None else 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("proxy_logging_registered", [True, False])
+async def test_scheduled_focus_export_without_a_redis_lock_runs_until_the_database_boundary(
+ monkeypatch: pytest.MonkeyPatch, proxy_logging_registered: bool
+) -> None:
+ monkeypatch.setattr(proxy_server.proxy_logging_obj.db_spend_update_writer.pod_lock_manager, "redis_cache", None)
+ if not proxy_logging_registered:
+ monkeypatch.setattr(proxy_server, "proxy_logging_obj", None)
+ monkeypatch.setattr(proxy_server, "prisma_client", None)
+ logger = FocusLogger(provider="s3", destination_config={"bucket_name": "focus-exports"})
+
+ with pytest.raises(RuntimeError, match="Database not connected"):
+ await logger.initialize_focus_export_job()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ ("locks_before", "locks_after", "granted"),
+ [
+ pytest.param({}, {}, [(FOCUS_LOCK, "this-pod")], id="lock-free"),
+ pytest.param({FOCUS_LOCK: "this-pod"}, {}, [], id="lock-already-held-by-this-pod"),
+ pytest.param(
+ {VANTAGE_LOCK: "another-pod"},
+ {VANTAGE_LOCK: "another-pod"},
+ [(FOCUS_LOCK, "this-pod")],
+ id="vantage-lock-held-by-another-pod",
+ ),
+ ],
+)
+async def test_scheduled_focus_export_runs_under_the_focus_pod_lock_and_releases_it_when_the_export_fails(
+ monkeypatch: pytest.MonkeyPatch,
+ locks_before: dict[str, str],
+ locks_after: dict[str, str],
+ granted: list[tuple[str, str]],
+) -> None:
+ lock_manager = proxy_server.proxy_logging_obj.db_spend_update_writer.pod_lock_manager
+ redis = _InMemoryRedis(dict(locks_before))
+ monkeypatch.setattr(lock_manager, "redis_cache", redis)
+ monkeypatch.setattr(lock_manager, "pod_id", "this-pod")
+ monkeypatch.setattr(proxy_server, "prisma_client", None)
+ logger = FocusLogger(provider="s3", destination_config={"bucket_name": "focus-exports"})
+
+ with pytest.raises(RuntimeError, match="Database not connected"):
+ await logger.initialize_focus_export_job()
+
+ assert redis.locks == locks_after
+ assert redis.granted == granted
+
+
+@pytest.mark.asyncio
+async def test_scheduled_focus_export_is_skipped_while_another_pod_holds_the_focus_lock(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ lock_manager = proxy_server.proxy_logging_obj.db_spend_update_writer.pod_lock_manager
+ redis = _InMemoryRedis({FOCUS_LOCK: "another-pod"})
+ monkeypatch.setattr(lock_manager, "redis_cache", redis)
+ monkeypatch.setattr(lock_manager, "pod_id", "this-pod")
+ monkeypatch.setattr(proxy_server, "prisma_client", None)
+ logger = FocusLogger(provider="s3", destination_config={"bucket_name": "focus-exports"})
+
+ await logger.initialize_focus_export_job()
+
+ assert redis.locks == {FOCUS_LOCK: "another-pod"}
+ assert redis.granted == []
diff --git a/tests/unit/integrations/test_argilla.py b/tests/unit/integrations/test_argilla.py
new file mode 100644
index 00000000000..b8948acf011
--- /dev/null
+++ b/tests/unit/integrations/test_argilla.py
@@ -0,0 +1,110 @@
+import asyncio
+import json
+import re
+
+import httpx
+import pytest
+import respx
+
+import litellm
+from litellm.integrations.argilla import ArgillaLogger
+
+ARGILLA_BASE_URL = "https://argilla.example.test"
+USER_MESSAGE = {"role": "user", "content": "where is my order?"}
+
+
+async def _cancel_background_tasks() -> None:
+ background_tasks = [task for task in asyncio.all_tasks() if task is not asyncio.current_task()]
+ for task in background_tasks:
+ task.cancel()
+ await asyncio.gather(*background_tasks, return_exceptions=True)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ ("transformation_object", "standard_logging_object", "expected_fields"),
+ [
+ pytest.param(
+ {"prompt": "messages", "answer": "response"},
+ {"messages": [USER_MESSAGE], "response": "it ships tomorrow"},
+ {"prompt": [USER_MESSAGE], "answer": "it ships tomorrow"},
+ id="message-list-and-text-response",
+ ),
+ pytest.param(
+ {"answer": "response", "prompt": "messages"},
+ {"messages": USER_MESSAGE, "response": {"choices": [{"message": {"content": "it ships tomorrow"}}]}},
+ {"answer": "it ships tomorrow", "prompt": [USER_MESSAGE]},
+ id="single-message-and-chat-completion-response",
+ ),
+ pytest.param(
+ {"prompt": "messages", "context": "messages", "answer": "response", "label": "response"},
+ {"messages": [USER_MESSAGE, USER_MESSAGE], "response": {"choices": [{"message": {}}]}},
+ {
+ "prompt": [USER_MESSAGE, USER_MESSAGE],
+ "context": [USER_MESSAGE, USER_MESSAGE],
+ "answer": "",
+ "label": "",
+ },
+ id="payload-field-reused-by-several-argilla-fields",
+ ),
+ ],
+)
+async def test_argilla_logger_sends_the_record_shaped_by_the_configured_field_mapping(
+ monkeypatch: pytest.MonkeyPatch,
+ respx_mock: respx.MockRouter,
+ transformation_object: dict[str, str],
+ standard_logging_object: dict[str, object],
+ expected_fields: dict[str, object],
+) -> None:
+ monkeypatch.setattr(litellm, "argilla_transformation_object", transformation_object)
+ respx_mock.get(f"{ARGILLA_BASE_URL}/api/v1/me/datasets", params={"name": "support-chats"}).mock(
+ return_value=httpx.Response(200, json={"items": [{"id": "dataset-403"}]})
+ )
+ bulk_route = respx_mock.post(f"{ARGILLA_BASE_URL}/api/v1/datasets/dataset-403/records/bulk").mock(
+ return_value=httpx.Response(200, json={})
+ )
+ logger = ArgillaLogger(
+ argilla_api_key="argilla-key",
+ argilla_dataset_name="support-chats",
+ argilla_base_url=ARGILLA_BASE_URL,
+ batch_size=1,
+ )
+ try:
+ logger.log_success_event({"standard_logging_object": standard_logging_object}, None, None, None)
+ finally:
+ await _cancel_background_tasks()
+
+ sent_request = bulk_route.calls.last.request
+ assert json.loads(sent_request.content) == [{"fields": expected_fields}]
+ assert sent_request.headers["X-Argilla-Api-Key"] == "argilla-key"
+ assert logger.argilla_transformation_object is transformation_object
+
+
+@pytest.mark.parametrize(
+ ("transformation_object", "message"),
+ [
+ pytest.param(
+ None,
+ "'litellm.argilla_transformation_object' is required, to log your payload to Argilla.",
+ id="mapping-not-configured",
+ ),
+ pytest.param(
+ ["messages", "response"],
+ "'argilla_transformation_object' must be a dictionary, to log your payload to Argilla.",
+ id="mapping-is-a-list",
+ ),
+ pytest.param(
+ {"prompt": "messages", "spend": "response_cost"},
+ "All values in argilla_transformation_object must be a key in SUPPORTED_PAYLOAD_FIELDS, "
+ "response_cost is not a valid key.",
+ id="mapping-names-an-unsupported-payload-field",
+ ),
+ ],
+)
+def test_argilla_logger_refuses_to_start_without_a_supported_field_mapping(
+ monkeypatch: pytest.MonkeyPatch, transformation_object: object, message: str
+) -> None:
+ monkeypatch.setattr(litellm, "argilla_transformation_object", transformation_object)
+
+ with pytest.raises(Exception, match=re.escape(message)):
+ ArgillaLogger(argilla_api_key="argilla-key")
diff --git a/tests/unit/integrations/vantage/__init__.py b/tests/unit/integrations/vantage/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/integrations/vantage/test_vantage_logger.py b/tests/unit/integrations/vantage/test_vantage_logger.py
new file mode 100644
index 00000000000..8ff2c7fa046
--- /dev/null
+++ b/tests/unit/integrations/vantage/test_vantage_logger.py
@@ -0,0 +1,92 @@
+import pytest
+
+from litellm.integrations.vantage.vantage_logger import VantageLogger
+from litellm.proxy import proxy_server
+
+FOCUS_LOCK = "cronjob_lock:focus_export_usage_data"
+VANTAGE_LOCK = "cronjob_lock:vantage_export_usage_data"
+
+
+class _InMemoryRedis:
+ def __init__(self, locks: dict[str, str]) -> None:
+ self.locks = locks
+ self.granted: list[tuple[str, str]] = []
+
+ async def async_set_cache(self, key: str, value: str, nx: bool, ttl: int) -> bool:
+ if nx and key in self.locks:
+ return False
+ self.locks[key] = value
+ self.granted.append((key, value))
+ return True
+
+ async def async_get_cache(self, key: str) -> str | None:
+ return self.locks.get(key)
+
+ async def async_delete_cache(self, key: str) -> int:
+ return 1 if self.locks.pop(key, None) is not None else 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("proxy_logging_registered", [True, False])
+async def test_scheduled_vantage_export_without_a_redis_lock_runs_until_the_database_boundary(
+ monkeypatch: pytest.MonkeyPatch, proxy_logging_registered: bool
+) -> None:
+ monkeypatch.setattr(proxy_server.proxy_logging_obj.db_spend_update_writer.pod_lock_manager, "redis_cache", None)
+ if not proxy_logging_registered:
+ monkeypatch.setattr(proxy_server, "proxy_logging_obj", None)
+ monkeypatch.setattr(proxy_server, "prisma_client", None)
+ logger = VantageLogger(api_key="vantage-key", integration_token="vantage-token")
+
+ with pytest.raises(RuntimeError, match="Database not connected"):
+ await logger.initialize_focus_export_job()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ ("locks_before", "locks_after", "granted"),
+ [
+ pytest.param({}, {}, [(VANTAGE_LOCK, "this-pod")], id="lock-free"),
+ pytest.param({VANTAGE_LOCK: "this-pod"}, {}, [], id="lock-already-held-by-this-pod"),
+ pytest.param(
+ {FOCUS_LOCK: "another-pod"},
+ {FOCUS_LOCK: "another-pod"},
+ [(VANTAGE_LOCK, "this-pod")],
+ id="focus-lock-held-by-another-pod",
+ ),
+ ],
+)
+async def test_scheduled_vantage_export_runs_under_the_vantage_pod_lock_and_releases_it_when_the_export_fails(
+ monkeypatch: pytest.MonkeyPatch,
+ locks_before: dict[str, str],
+ locks_after: dict[str, str],
+ granted: list[tuple[str, str]],
+) -> None:
+ lock_manager = proxy_server.proxy_logging_obj.db_spend_update_writer.pod_lock_manager
+ redis = _InMemoryRedis(dict(locks_before))
+ monkeypatch.setattr(lock_manager, "redis_cache", redis)
+ monkeypatch.setattr(lock_manager, "pod_id", "this-pod")
+ monkeypatch.setattr(proxy_server, "prisma_client", None)
+ logger = VantageLogger(api_key="vantage-key", integration_token="vantage-token")
+
+ with pytest.raises(RuntimeError, match="Database not connected"):
+ await logger.initialize_focus_export_job()
+
+ assert redis.locks == locks_after
+ assert redis.granted == granted
+
+
+@pytest.mark.asyncio
+async def test_scheduled_vantage_export_is_skipped_while_another_pod_holds_the_vantage_lock(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ lock_manager = proxy_server.proxy_logging_obj.db_spend_update_writer.pod_lock_manager
+ redis = _InMemoryRedis({VANTAGE_LOCK: "another-pod"})
+ monkeypatch.setattr(lock_manager, "redis_cache", redis)
+ monkeypatch.setattr(lock_manager, "pod_id", "this-pod")
+ monkeypatch.setattr(proxy_server, "prisma_client", None)
+ logger = VantageLogger(api_key="vantage-key", integration_token="vantage-token")
+
+ await logger.initialize_focus_export_job()
+
+ assert redis.locks == {VANTAGE_LOCK: "another-pod"}
+ assert redis.granted == []
diff --git a/tests/unit/llms/a2a/chat/test_a2a_chat_transformation.py b/tests/unit/llms/a2a/chat/test_a2a_chat_transformation.py
index a419e73a2e1..76bce26e2e0 100644
--- a/tests/unit/llms/a2a/chat/test_a2a_chat_transformation.py
+++ b/tests/unit/llms/a2a/chat/test_a2a_chat_transformation.py
@@ -1,11 +1,15 @@
"""Tests for litellm/llms/a2a/chat/transformation.py response transform."""
+import uuid
from unittest.mock import MagicMock
+import httpx
import pytest
+from pydantic import ValidationError
from litellm.llms.a2a.chat.streaming_iterator import A2AModelResponseIterator
from litellm.llms.a2a.chat.transformation import A2AConfig
+from litellm.llms.a2a.common_utils import A2AError
from litellm.types.utils import ModelResponse
@@ -121,3 +125,104 @@ def test_resolve_agent_config_from_registry_returns_the_explicit_headers_object_
assert resolved[2] is headers
assert resolved[:2] == ("http://explicit.example", "explicit-key")
assert optional_params == {"stream": True}
+
+
+def _transform(payload: object, headers: dict[str, str] | None = None) -> ModelResponse:
+ return A2AConfig().transform_response(
+ model="a2a/test-agent",
+ raw_response=httpx.Response(200, json=payload, headers=headers),
+ model_response=ModelResponse(),
+ logging_obj=None,
+ request_data={},
+ messages=[{"role": "user", "content": "hi there agent"}],
+ optional_params={},
+ litellm_params={},
+ encoding=None,
+ )
+
+
+@pytest.mark.parametrize(
+ ("result", "expected_text"),
+ [
+ ({"kind": "message", "parts": [{"kind": "text", "text": "direct"}]}, "direct"),
+ ({"message": {"parts": [{"kind": "text", "text": "nested"}]}}, "nested"),
+ ({"kind": "task", "status": {"message": {"parts": [{"kind": "text", "text": "status"}]}}}, "status"),
+ ({"kind": "task", "artifacts": [{"parts": [{"kind": "text", "text": "artifact"}]}]}, "artifact"),
+ ({"kind": "task", "status": {"state": "submitted"}}, ""),
+ (None, ""),
+ ("plain text", ""),
+ ([{"kind": "text", "text": "not a result object"}], ""),
+ ],
+)
+def test_transform_response_extracts_the_text_of_each_result_shape(result: object, expected_text: str):
+ response = _transform({"jsonrpc": "2.0", "id": "resp-1", "result": result, "extra": None})
+
+ assert response.choices[0].message.content == expected_text
+ assert response.choices[0].finish_reason == "stop"
+ assert response.id == "resp-1"
+ assert response.model == "a2a/test-agent"
+
+
+@pytest.mark.parametrize("payload", [{}, {"jsonrpc": "2.0", "result": {"parts": [{"kind": "text", "text": "4"}]}}])
+def test_transform_response_generates_an_id_when_the_agent_sends_none(payload: dict[str, object]):
+ response = _transform(payload)
+
+ assert str(uuid.UUID(response.id)) == response.id
+
+
+@pytest.mark.parametrize(
+ ("error", "expected_message"),
+ [
+ ({"code": -32600, "message": "Invalid Request"}, "A2A error: Invalid Request"),
+ ({"code": -32600}, "A2A error: Unknown error"),
+ ({}, "A2A error: Unknown error"),
+ ({"message": None}, "A2A error: None"),
+ ({"message": 429}, "A2A error: 429"),
+ ({"message": {"detail": "nested"}}, "A2A error: {'detail': 'nested'}"),
+ ],
+)
+def test_transform_response_raises_the_agent_error_with_its_message(error: dict[str, object], expected_message: str):
+ payload = {"jsonrpc": "2.0", "id": "1", "error": error, "result": {"parts": [{"kind": "text", "text": "4"}]}}
+
+ with pytest.raises(A2AError) as exc_info:
+ _transform(payload, headers={"x-agent": "billing"})
+
+ assert exc_info.value.message == expected_message
+ assert exc_info.value.status_code == 200
+ assert exc_info.value.headers["x-agent"] == "billing"
+
+
+@pytest.mark.parametrize("error", [None, "429 Too Many Requests", 429, False, ["Request timed out"]])
+def test_transform_response_rejects_an_error_member_that_is_not_an_object(error: object):
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"jsonrpc": "2.0", "id": "1", "error": error, "result": {"parts": []}})
+
+ assert [(detail["type"], detail["loc"]) for detail in exc_info.value.errors()] == [("dict_type", ())]
+ assert "input_value" not in str(exc_info.value)
+
+
+@pytest.mark.parametrize("payload", [[], [{"result": {"parts": []}}], ["rate limit"], "upstream unavailable"])
+def test_transform_response_rejects_a_body_that_is_not_an_object(payload: object):
+ with pytest.raises(ValidationError) as exc_info:
+ _transform(payload)
+
+ assert [(detail["type"], detail["loc"]) for detail in exc_info.value.errors()] == [("dict_type", ())]
+ assert "input_value" not in str(exc_info.value)
+
+
+def test_transform_response_reports_an_undecodable_body_as_an_agent_error():
+ with pytest.raises(A2AError) as exc_info:
+ A2AConfig().transform_response(
+ model="a2a/test-agent",
+ raw_response=httpx.Response(502, content=b"Bad Gateway"),
+ model_response=ModelResponse(),
+ logging_obj=None,
+ request_data={},
+ messages=[{"role": "user", "content": "hi there agent"}],
+ optional_params={},
+ litellm_params={},
+ encoding=None,
+ )
+
+ assert exc_info.value.status_code == 502
+ assert exc_info.value.message.startswith("Failed to parse A2A response: ")
diff --git a/tests/unit/llms/aiml/image_generation/test_aiml_image_generation_transformation.py b/tests/unit/llms/aiml/image_generation/test_aiml_image_generation_transformation.py
index 5ac4c7c4643..ea9ace5e994 100644
--- a/tests/unit/llms/aiml/image_generation/test_aiml_image_generation_transformation.py
+++ b/tests/unit/llms/aiml/image_generation/test_aiml_image_generation_transformation.py
@@ -1,7 +1,9 @@
import os
+from datetime import datetime
+import httpx
import pytest
-
+from pydantic import ValidationError
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
@@ -9,13 +11,11 @@ import litellm
litellm.model_cost = litellm.get_model_cost_map(url="")
-from litellm.llms.aiml.image_generation.cost_calculator import (
- cost_calculator as aiml_cost_calculator,
-)
+from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.aiml.image_generation.transformation import (
AimlImageGenerationConfig,
)
-from litellm.types.utils import ImageObject, ImageResponse
+from litellm.types.utils import ImageResponse
def test_openai_style_model_supports_full_openai_param_surface():
@@ -130,3 +130,126 @@ def test_openai_style_unsupported_param_dropped_with_drop_params():
assert mapped == {}
+def _transform_response(payload: object) -> ImageResponse:
+ return AimlImageGenerationConfig().transform_image_generation_response(
+ model="flux-pro",
+ raw_response=httpx.Response(200, json=payload),
+ model_response=ImageResponse(),
+ logging_obj=Logging(
+ model="flux-pro",
+ messages=[],
+ stream=False,
+ call_type="image_generation",
+ start_time=datetime(2026, 1, 1),
+ litellm_call_id="aiml-image-generation-test",
+ function_id="aiml-image-generation-test",
+ ),
+ request_data={},
+ optional_params={},
+ litellm_params={},
+ encoding=None,
+ )
+
+
+@pytest.mark.parametrize(
+ ("payload", "expected"),
+ [
+ (
+ {
+ "data": [
+ {"url": "https://cdn.aiml.example/a.png", "revised_prompt": "a red fox", "width": 1024},
+ "not-an-image",
+ {"b64_json": "QUJD", "revised_prompt": "a blue fox"},
+ [],
+ {"b64_json": "", "image_base64": "REVG"},
+ {"content_type": "image/png"},
+ ]
+ },
+ [
+ (None, "https://cdn.aiml.example/a.png", "a red fox"),
+ ("QUJD", None, "a blue fox"),
+ ("REVG", None, None),
+ ],
+ ),
+ (
+ {
+ "output": {
+ "choices": [
+ {"image_base64": "QUJD", "url": "https://cdn.aiml.example/ignored.png"},
+ "skipped",
+ {"url": "https://cdn.aiml.example/b.png"},
+ {"finish_reason": "stop"},
+ ]
+ }
+ },
+ [("QUJD", None, None), (None, "https://cdn.aiml.example/b.png", None)],
+ ),
+ (
+ {
+ "images": [
+ {"url": "https://cdn.aiml.example/c.png", "image_base64": "aWdub3JlZA=="},
+ [1, 2],
+ {"image_base64": "REVG"},
+ {"seed": 7},
+ ]
+ },
+ [(None, "https://cdn.aiml.example/c.png", None), ("REVG", None, None)],
+ ),
+ ({"output": {"choices": {"finish_reason": "stop"}}}, []),
+ ({"output": {"choices": "pending"}}, []),
+ ({"images": {"count": 0}}, []),
+ ({"images": ""}, []),
+ ({"data": "pending"}, []),
+ ({"id": "gen-1"}, []),
+ ([], []),
+ ],
+)
+def test_transform_image_generation_response_reads_images_from_every_response_shape(
+ payload: object, expected: list[tuple[str | None, str | None, str | None]]
+):
+ result = _transform_response(payload)
+
+ assert [(image.b64_json, image.url, image.revised_prompt) for image in result.data] == expected
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {"data": [7]},
+ {"data": [None]},
+ {"data": ["aiml-secret-url-token"]},
+ {"data": [["b64_json", "aiml-secret-url-token"]]},
+ {"output": {"choices": [False]}},
+ {"output": {"choices": ["aiml-secret image_base64"]}},
+ {"output": {"choices": None}},
+ {"images": [1.5]},
+ {"images": ["aiml-secret-url-token"]},
+ {"images": {"url": "aiml-secret-url-token"}},
+ {"images": 3},
+ ],
+)
+def test_transform_image_generation_response_rejects_malformed_image_entries_without_echoing_them(payload: object):
+ with pytest.raises(ValidationError) as exc_info:
+ _transform_response(payload)
+
+ assert "aiml-secret" not in str(exc_info.value)
+
+
+def _rejection_text(payload: object) -> str:
+ with pytest.raises(ValidationError) as exc_info:
+ _transform_response(payload)
+
+ return str(exc_info.value)
+
+
+@pytest.mark.parametrize("position", [403, 429])
+@pytest.mark.parametrize("malformed_entry", [None, 7, "url"])
+def test_transform_image_generation_response_rejection_text_is_the_same_wherever_the_malformed_entry_sits(
+ position: int, malformed_entry: object
+):
+ first = [malformed_entry]
+ late = [*[{"seed": 1}] * position, malformed_entry]
+
+ assert _rejection_text({"data": late}) == _rejection_text({"data": first})
+ assert _rejection_text({"output": {"choices": late}}) == _rejection_text({"output": {"choices": first}})
+ assert _rejection_text({"images": late}) == _rejection_text({"images": first})
diff --git a/tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py b/tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py
index e9736d47537..b5c7b337806 100644
--- a/tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py
+++ b/tests/unit/llms/anthropic/files/test_anthropic_files_transformation.py
@@ -6,7 +6,9 @@ OpenAI-compatible file operations and Anthropic's Files API format.
"""
import asyncio
+import datetime
import io
+import json
import threading
import time
@@ -16,6 +18,7 @@ from openai.types.file_deleted import FileDeleted
from pydantic import ValidationError
from unittest.mock import Mock, patch
+from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.anthropic.files.transformation import (
AnthropicFilesConfig,
ANTHROPIC_FILES_API_BASE,
@@ -663,3 +666,151 @@ class TestProviderConfigRegistration:
)
assert config is not None
assert isinstance(config, AnthropicFilesConfig)
+
+
+def _listed_files(body: bytes) -> list[OpenAIFileObject]:
+ logging_obj = Logging(
+ model="",
+ messages=[],
+ stream=False,
+ call_type="file_list",
+ start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+ litellm_call_id="call-id",
+ function_id="function-id",
+ )
+ return AnthropicFilesConfig().transform_list_files_response(
+ raw_response=httpx.Response(200, content=body), logging_obj=logging_obj, litellm_params={}
+ )
+
+
+def _anthropic_file(index: int) -> dict[str, object]:
+ return {
+ "id": f"file-{index}",
+ "type": "file",
+ "filename": f"doc-{index}.pdf",
+ "mime_type": "application/pdf",
+ "size_bytes": index,
+ "created_at": "2025-01-01T00:00:00Z",
+ "downloadable": False,
+ }
+
+
+def test_list_files_response_maps_every_anthropic_file_in_provider_order() -> None:
+ result = _listed_files(
+ json.dumps(
+ {
+ "data": [
+ _anthropic_file(100),
+ {"id": "file-b", "filename": "b.jsonl", "bytes": 200, "created_at": "2025-01-02T00:00:00Z"},
+ {"id": "file-c", "filename": "c.jsonl", "size_bytes": "300", "purpose": "batch", "created_at": ""},
+ ],
+ "has_more": False,
+ "first_id": "file-100",
+ "last_id": "file-c",
+ }
+ ).encode()
+ )
+
+ assert [(listed.id, listed.filename, listed.bytes, listed.purpose) for listed in result] == [
+ ("file-100", "doc-100.pdf", 100, "messages"),
+ ("file-b", "b.jsonl", 200, "messages"),
+ ("file-c", "c.jsonl", 300, "batch"),
+ ]
+ assert [type(listed.bytes) for listed in result] == [int, int, int]
+ assert [listed.created_at for listed in result[:2]] == [1735689600, 1735776000]
+ assert result[0] == OpenAIFileObject(
+ id="file-100",
+ bytes=100,
+ created_at=1735689600,
+ filename="doc-100.pdf",
+ object="file",
+ purpose="messages",
+ status="uploaded",
+ status_details=None,
+ )
+
+
+@pytest.mark.parametrize(
+ "body",
+ [
+ b"{}",
+ b'{"has_more": false}',
+ b'{"data": []}',
+ b'{"data": ""}',
+ b'{"data": {}}',
+ b'{"data": [], "unknown": [1, {"nested": null}]}',
+ ],
+)
+def test_list_files_response_without_any_file_is_an_empty_list(body: bytes) -> None:
+ assert _listed_files(body) == []
+
+
+@pytest.mark.parametrize("body", [b'["file-abc123"]', b'"file-abc123"', b"null", b"7"])
+def test_list_files_response_rejects_a_body_that_is_not_a_json_object(body: bytes) -> None:
+ with pytest.raises(ValidationError) as exc_info:
+ _listed_files(body)
+
+ assert [error["type"] for error in exc_info.value.errors()] == ["dict_type"]
+
+
+@pytest.mark.parametrize("body", [b'{"data": null}', b'{"data": 7}', b'{"data": 0}', b'{"data": true}'])
+def test_list_files_response_rejects_data_that_cannot_be_iterated(body: bytes) -> None:
+ with pytest.raises(ValidationError) as exc_info:
+ _listed_files(body)
+
+ assert [error["type"] for error in exc_info.value.errors()] == ["iterable_type"]
+
+
+@pytest.mark.parametrize(
+ "data",
+ [
+ ["file-abc123"],
+ [None],
+ [7],
+ [["file-abc123"]],
+ [_anthropic_file(0), "file-abc123"],
+ "file-abc123",
+ {"file-abc123": _anthropic_file(0)},
+ ],
+)
+def test_list_files_response_rejects_a_listed_file_that_is_not_a_json_object(data: object) -> None:
+ with pytest.raises(ValidationError) as exc_info:
+ _listed_files(json.dumps({"data": data}).encode())
+
+ assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [("dict_type", ())]
+
+
+def test_list_files_response_error_for_a_late_malformed_file_does_not_name_its_position() -> None:
+ files = [_anthropic_file(index) for index in range(403)]
+
+ with pytest.raises(ValidationError) as exc_info:
+ _listed_files(json.dumps({"data": [*files, "file-abc123"]}).encode())
+
+ assert "403" not in str(exc_info.value)
+ assert "429" not in str(exc_info.value)
+ assert "file-abc123" not in str(exc_info.value)
+
+
+def test_list_files_response_reports_every_invalid_field_of_a_listed_file() -> None:
+ with pytest.raises(ValidationError) as exc_info:
+ _listed_files(b'{"data": [{"id": null, "filename": 7, "size_bytes": "many", "purpose": "unknown"}]}')
+
+ assert [error["loc"] for error in exc_info.value.errors()] == [("id",), ("bytes",), ("filename",), ("purpose",)]
+ assert "many" in str(exc_info.value)
+
+
+@pytest.mark.parametrize(
+ "data, expected_errors",
+ [
+ ([{"id": None}, "file-abc123"], [("string_type", ("id",))]),
+ (["file-abc123", {"id": None}], [("dict_type", ())]),
+ ([{"id": None}, {"filename": 7}], [("string_type", ("id",))]),
+ ],
+)
+def test_list_files_response_reports_only_the_first_bad_file_in_provider_order(
+ data: object, expected_errors: object
+) -> None:
+ with pytest.raises(ValidationError) as exc_info:
+ _listed_files(json.dumps({"data": data}).encode())
+
+ assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == expected_errors
diff --git a/tests/unit/llms/azure_ai/anthropic/__init__.py b/tests/unit/llms/azure_ai/anthropic/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/__init__.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py
new file mode 100644
index 00000000000..29380c6fc53
--- /dev/null
+++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py
@@ -0,0 +1,188 @@
+import json
+import operator
+from typing import Final
+
+import httpx
+import pytest
+import respx
+from pydantic import ValidationError
+
+import litellm
+from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import AzureAIAnthropicTokenCounter
+from litellm.types.utils import TokenCountResponse
+
+API_BASE: Final = "https://my-resource.services.ai.azure.com"
+COUNT_URL: Final = f"{API_BASE}/anthropic/v1/messages/count_tokens"
+
+
+async def _count(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> TokenCountResponse | None:
+ monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+ with respx.mock:
+ respx.post(COUNT_URL).mock(return_value=response)
+ return await AzureAIAnthropicTokenCounter().count_tokens(
+ model_to_use="claude-sonnet-4-5",
+ messages=[{"role": "user", "content": "hi"}],
+ contents=None,
+ deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}},
+ request_model="azure_ai/claude-sonnet-4-5",
+ )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ ("payload", "expected_total"),
+ [
+ ({"input_tokens": 15}, 15),
+ ({"input_tokens": "15"}, 15),
+ ({"input_tokens": 15.0}, 15),
+ ({"input_tokens": True}, 1),
+ ({"input_tokens": -3}, -3),
+ ({}, 0),
+ ({"type": "error", "error": {"message": "429 rate limit"}}, 0),
+ ({"input_tokens": 7, "context_management": {"original_input_tokens": 9}, "nested": {"a": [1, None]}}, 7),
+ ],
+)
+async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload(
+ monkeypatch: pytest.MonkeyPatch, payload: dict[str, object], expected_total: int
+) -> None:
+ result: Final = await _count(monkeypatch, httpx.Response(200, json=payload))
+
+ assert result is not None
+ assert repr(result.model_dump()) == repr(
+ {
+ "total_tokens": expected_total,
+ "request_model": "azure_ai/claude-sonnet-4-5",
+ "model_used": "claude-sonnet-4-5",
+ "tokenizer_type": "azure_ai_anthropic_api",
+ "original_response": payload,
+ "error": False,
+ "error_message": None,
+ "status_code": None,
+ }
+ )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "rejected_count",
+ [
+ None,
+ 15.5,
+ "Request timed out",
+ "429 Too Many Requests",
+ [15],
+ [*range(403), "Request timed out"],
+ {"value": 429},
+ ],
+)
+async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_validation_message(
+ monkeypatch: pytest.MonkeyPatch, rejected_count: object
+) -> None:
+ with pytest.raises(ValidationError) as rejection:
+ TokenCountResponse.model_validate(
+ {
+ "total_tokens": rejected_count,
+ "request_model": "azure_ai/claude-sonnet-4-5",
+ "model_used": "claude-sonnet-4-5",
+ "tokenizer_type": "azure_ai_anthropic_api",
+ }
+ )
+
+ result: Final = await _count(monkeypatch, httpx.Response(200, json={"input_tokens": rejected_count}))
+
+ assert result is not None
+ assert [error["loc"] for error in rejection.value.errors()] == [("total_tokens",)]
+ assert repr(result.model_dump()) == repr(
+ {
+ "total_tokens": 0,
+ "request_model": "azure_ai/claude-sonnet-4-5",
+ "model_used": "claude-sonnet-4-5",
+ "tokenizer_type": "azure_ai_anthropic_api",
+ "original_response": None,
+ "error": True,
+ "error_message": str(rejection.value),
+ "status_code": 500,
+ }
+ )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "payload",
+ [[*({"input_tokens": position} for position in range(403)), "429"], "input_tokens", 15, 0, False],
+)
+async def test_count_payload_that_is_not_an_object_becomes_an_error_response_naming_the_failed_lookup(
+ monkeypatch: pytest.MonkeyPatch, payload: object
+) -> None:
+ with pytest.raises(AttributeError) as failed_lookup:
+ operator.attrgetter("get")(payload)
+
+ result: Final = await _count(monkeypatch, httpx.Response(200, json=payload))
+
+ assert result is not None
+ assert repr(result.model_dump()) == repr(
+ {
+ "total_tokens": 0,
+ "request_model": "azure_ai/claude-sonnet-4-5",
+ "model_used": "claude-sonnet-4-5",
+ "tokenizer_type": "azure_ai_anthropic_api",
+ "original_response": None,
+ "error": True,
+ "error_message": str(failed_lookup.value),
+ "status_code": 500,
+ }
+ )
+
+
+@pytest.mark.asyncio
+async def test_null_count_payload_falls_back_to_local_counting(monkeypatch: pytest.MonkeyPatch) -> None:
+ assert await _count(monkeypatch, httpx.Response(200, content=b"null")) is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ ("upstream_status", "upstream_text"),
+ [(429, "rate limit exceeded"), (403, "forbidden"), (408, "Request timed out")],
+)
+async def test_upstream_error_status_becomes_an_error_response_with_that_status_and_text(
+ monkeypatch: pytest.MonkeyPatch, upstream_status: int, upstream_text: str
+) -> None:
+ result: Final = await _count(monkeypatch, httpx.Response(upstream_status, text=upstream_text))
+
+ assert result is not None
+ assert repr(result.model_dump()) == repr(
+ {
+ "total_tokens": 0,
+ "request_model": "azure_ai/claude-sonnet-4-5",
+ "model_used": "claude-sonnet-4-5",
+ "tokenizer_type": "azure_ai_anthropic_api",
+ "original_response": None,
+ "error": True,
+ "error_message": upstream_text,
+ "status_code": upstream_status,
+ }
+ )
+
+
+@pytest.mark.asyncio
+async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ with pytest.raises(json.JSONDecodeError) as undecodable:
+ json.loads("not json")
+
+ result: Final = await _count(monkeypatch, httpx.Response(200, text="not json"))
+
+ assert result is not None
+ assert repr(result.model_dump()) == repr(
+ {
+ "total_tokens": 0,
+ "request_model": "azure_ai/claude-sonnet-4-5",
+ "model_used": "claude-sonnet-4-5",
+ "tokenizer_type": "azure_ai_anthropic_api",
+ "original_response": None,
+ "error": True,
+ "error_message": f"CountTokens processing error: {undecodable.value}",
+ "status_code": 500,
+ }
+ )
diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py
index 4db786668b8..6f326bfec44 100644
--- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py
+++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py
@@ -1,8 +1,8 @@
-import asyncio
-import json
from unittest.mock import Mock
+import httpx
import pytest
+from pydantic import ValidationError
# Ensure the project root is on the import path so `litellm` can be imported when
# tests are executed from any working directory.
@@ -13,6 +13,21 @@ from litellm.llms.bedrock.chat.invoke_transformations.amazon_qwen2_transformatio
from litellm.types.utils import ModelResponse
+def _transform(body: object) -> ModelResponse:
+ return AmazonQwen2Config().transform_response(
+ model="qwen2/test-model",
+ raw_response=httpx.Response(200, json=body),
+ model_response=ModelResponse(),
+ logging_obj=Mock(),
+ request_data={},
+ messages=[{"role": "user", "content": "Hello!"}],
+ optional_params={},
+ litellm_params={},
+ encoding=None,
+ api_key="test-key",
+ )
+
+
def test_qwen2_get_supported_params():
"""Test that Qwen2 config returns correct supported parameters"""
config = AmazonQwen2Config()
@@ -253,10 +268,89 @@ def test_qwen2_transform_response_without_usage():
assert result.choices[0]["finish_reason"] == "stop"
+@pytest.mark.parametrize(
+ ("usage", "expected_counts"),
+ [
+ pytest.param({}, (0, 0, 0), id="no-counts"),
+ pytest.param(
+ {"prompt_tokens": None, "completion_tokens": None, "total_tokens": None}, (0, 0, 0), id="null-counts"
+ ),
+ pytest.param({"prompt_tokens": 7}, (7, 0, 0), id="only-prompt-count"),
+ pytest.param(
+ {"prompt_tokens": "7", "completion_tokens": 3.0, "total_tokens": ""}, (7, 3, 0), id="loosely-typed-counts"
+ ),
+ pytest.param(
+ {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3, "latency_ms": [12, None]},
+ (1, 2, 3),
+ id="unmodelled-field",
+ ),
+ ],
+)
+def test_qwen2_transform_response_reads_usage_counts_that_are_missing_null_or_loosely_typed(
+ usage: dict[str, object], expected_counts: tuple[int, int, int]
+) -> None:
+ result = _transform({"text": "Hi", "usage": usage})
+
+ reported = result.usage
+ assert (reported.prompt_tokens, reported.completion_tokens, reported.total_tokens) == expected_counts
+ assert result.choices[0].message.content == "Hi"
+
+
+def test_qwen2_transform_response_yields_empty_content_when_the_body_carries_no_generated_text() -> None:
+ result = _transform({"stop_reason": "length", "metrics": {"latency_ms": 12}})
+
+ assert result.choices[0].message.content == ""
+ assert result.choices[0].finish_reason == "stop"
+
+
+@pytest.mark.parametrize(
+ "body",
+ [
+ pytest.param({"text": ["secret-output"]}, id="generated-text-is-a-list"),
+ pytest.param({"text": None, "debug": "secret-output"}, id="generated-text-is-null"),
+ pytest.param({"text": "Hi", "usage": ["secret-output"]}, id="usage-is-a-list"),
+ pytest.param(["secret-output"], id="body-is-a-list"),
+ ],
+)
+def test_qwen2_transform_response_rejects_a_malformed_body_without_echoing_it(body: object) -> None:
+ with pytest.raises(ValidationError) as exc_info:
+ _transform(body)
+
+ assert "secret-output" not in str(exc_info.value)
+
+
+def test_qwen2_transform_response_reports_every_count_that_is_not_a_number_with_its_value() -> None:
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"text": "Hi", "usage": {"prompt_tokens": "many", "total_tokens": [4, 2]}})
+
+ assert exc_info.value.title == "Usage"
+ assert {error["loc"]: error["input"] for error in exc_info.value.errors()} == {
+ ("prompt_tokens",): "many",
+ ("total_tokens",): [4, 2],
+ }
+ assert "input_value='many'" in str(exc_info.value)
+ assert "input_value=[4, 2]" in str(exc_info.value)
+
+
+@pytest.mark.parametrize(
+ ("body", "expected_content"),
+ [
+ pytest.param({"generation": "", "text": "from text"}, "from text", id="empty-generation"),
+ pytest.param({"generation": None, "text": "from text"}, "from text", id="null-generation"),
+ pytest.param({"generation": "from generation", "text": None}, "from generation", id="null-text"),
+ pytest.param({"generation": None, "text": ""}, "", id="both-empty"),
+ ],
+)
+def test_qwen2_transform_response_falls_back_to_text_when_generation_is_empty_or_null(
+ body: dict[str, object], expected_content: str
+) -> None:
+ assert _transform(body).choices[0].message.content == expected_content
+
+
def test_qwen2_provider_detection():
"""Test that Qwen2 provider is correctly detected from model names"""
- from litellm.utils import ProviderConfigManager
from litellm.types.utils import LlmProviders
+ from litellm.utils import ProviderConfigManager
# Test with qwen2/ prefix
config = ProviderConfigManager.get_provider_chat_config(
diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py
index e011b1fca2b..a147dbb2c9b 100644
--- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py
+++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_qwen3_transformation.py
@@ -1,8 +1,8 @@
-import asyncio
-import json
from unittest.mock import Mock
+import httpx
import pytest
+from pydantic import ValidationError
# Ensure the project root is on the import path so `litellm` can be imported when
# tests are executed from any working directory.
@@ -13,6 +13,21 @@ from litellm.llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformatio
from litellm.types.utils import ModelResponse
+def _transform(body: object) -> ModelResponse:
+ return AmazonQwen3Config().transform_response(
+ model="qwen3/test-model",
+ raw_response=httpx.Response(200, json=body),
+ model_response=ModelResponse(),
+ logging_obj=Mock(),
+ request_data={},
+ messages=[{"role": "user", "content": "Hello!"}],
+ optional_params={},
+ litellm_params={},
+ encoding=None,
+ api_key="test-key",
+ )
+
+
def test_qwen3_get_supported_params():
"""Test that Qwen3 config returns correct supported parameters"""
config = AmazonQwen3Config()
@@ -181,10 +196,74 @@ def test_qwen3_transform_response_without_usage():
assert result.choices[0]["finish_reason"] == "stop"
+@pytest.mark.parametrize(
+ ("usage", "expected_counts"),
+ [
+ pytest.param({}, (0, 0, 0), id="no-counts"),
+ pytest.param(
+ {"prompt_tokens": None, "completion_tokens": None, "total_tokens": None}, (0, 0, 0), id="null-counts"
+ ),
+ pytest.param({"prompt_tokens": 7}, (7, 0, 0), id="only-prompt-count"),
+ pytest.param(
+ {"prompt_tokens": "7", "completion_tokens": 3.0, "total_tokens": ""}, (7, 3, 0), id="loosely-typed-counts"
+ ),
+ pytest.param(
+ {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3, "latency_ms": [12, None]},
+ (1, 2, 3),
+ id="unmodelled-field",
+ ),
+ ],
+)
+def test_qwen3_transform_response_reads_usage_counts_that_are_missing_null_or_loosely_typed(
+ usage: dict[str, object], expected_counts: tuple[int, int, int]
+) -> None:
+ result = _transform({"generation": "Hi", "usage": usage})
+
+ reported = result.usage
+ assert (reported.prompt_tokens, reported.completion_tokens, reported.total_tokens) == expected_counts
+ assert result.choices[0].message.content == "Hi"
+
+
+def test_qwen3_transform_response_yields_empty_content_when_the_body_carries_no_generated_text() -> None:
+ result = _transform({"stop_reason": "length", "metrics": {"latency_ms": 12}})
+
+ assert result.choices[0].message.content == ""
+ assert result.choices[0].finish_reason == "stop"
+
+
+@pytest.mark.parametrize(
+ "body",
+ [
+ pytest.param({"generation": ["secret-output"]}, id="generated-text-is-a-list"),
+ pytest.param({"generation": None, "debug": "secret-output"}, id="generated-text-is-null"),
+ pytest.param({"generation": "Hi", "usage": ["secret-output"]}, id="usage-is-a-list"),
+ pytest.param(["secret-output"], id="body-is-a-list"),
+ ],
+)
+def test_qwen3_transform_response_rejects_a_malformed_body_without_echoing_it(body: object) -> None:
+ with pytest.raises(ValidationError) as exc_info:
+ _transform(body)
+
+ assert "secret-output" not in str(exc_info.value)
+
+
+def test_qwen3_transform_response_reports_every_count_that_is_not_a_number_with_its_value() -> None:
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"generation": "Hi", "usage": {"prompt_tokens": "many", "total_tokens": [4, 2]}})
+
+ assert exc_info.value.title == "Usage"
+ assert {error["loc"]: error["input"] for error in exc_info.value.errors()} == {
+ ("prompt_tokens",): "many",
+ ("total_tokens",): [4, 2],
+ }
+ assert "input_value='many'" in str(exc_info.value)
+ assert "input_value=[4, 2]" in str(exc_info.value)
+
+
def test_qwen3_provider_detection():
"""Test that Qwen3 provider is correctly detected from model names"""
- from litellm.utils import ProviderConfigManager
from litellm.types.utils import LlmProviders
+ from litellm.utils import ProviderConfigManager
# Test with qwen3/ prefix
config = ProviderConfigManager.get_provider_chat_config(
diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py
index edf93e53dc2..d9b01c4e75a 100644
--- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py
+++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py
@@ -401,6 +401,27 @@ def test_router_deployment_with_a_non_numeric_stream_chunk_size_gets_a_400_befor
send.assert_not_called()
+@pytest.mark.parametrize(
+ ("model", "expected_provider"),
+ [
+ pytest.param("us.anthropic.claude-3-5-sonnet-20240620-v1:0", "anthropic", id="region-prefixed-model-id"),
+ pytest.param(
+ "arn:aws:bedrock:us-east-1:123:inference-profile/us.meta.llama3-2-90b-instruct-v1:0",
+ "meta",
+ id="inference-profile-arn",
+ ),
+ pytest.param("my-anthropic-and-cohere-model", "cohere", id="two-providers-named-earliest-listed-wins"),
+ pytest.param("us.twelvelabs.pegasus-1-2-v1:0", "twelvelabs", id="provider-listed-late"),
+ pytest.param("unknown-model", None, id="no-provider-named"),
+ pytest.param("", None, id="empty-model"),
+ ],
+)
+def test_get_bedrock_invoke_provider_finds_a_provider_named_anywhere_in_the_model(
+ model: str, expected_provider: str | None
+) -> None:
+ assert AmazonInvokeConfig.get_bedrock_invoke_provider(model) == expected_provider
+
+
@pytest.mark.respx(assert_all_called=False)
def test_transform_request_meta_llama(monkeypatch, respx_mock):
monkeypatch.setattr(litellm, "known_tokenizer_config", dict(litellm.known_tokenizer_config))
diff --git a/tests/unit/llms/brave/__init__.py b/tests/unit/llms/brave/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/brave/search/__init__.py b/tests/unit/llms/brave/search/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/brave/search/test_brave_search_transformation.py b/tests/unit/llms/brave/search/test_brave_search_transformation.py
new file mode 100644
index 00000000000..8523c7c462d
--- /dev/null
+++ b/tests/unit/llms/brave/search/test_brave_search_transformation.py
@@ -0,0 +1,123 @@
+import httpx
+import pytest
+from pydantic import ValidationError
+
+from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
+from litellm.llms.brave.search.transformation import BraveSearchConfig
+
+
+def _transform(payload: object, query: str = "q=litellm") -> SearchResponse:
+ request = httpx.Request("GET", f"https://api.search.brave.com/res/v1/web/search?{query}")
+ return BraveSearchConfig().transform_search_response(
+ raw_response=httpx.Response(200, json=payload, request=request), logging_obj=None
+ )
+
+
+def _entry(index: int) -> dict[str, str]:
+ return {"title": f"Title {index}", "url": f"https://example.com/{index}", "description": f"Snippet {index}"}
+
+
+def test_brave_results_are_mapped_section_by_section_with_normalized_dates():
+ response = _transform(
+ {
+ "type": "search",
+ "query": {"original": "litellm"},
+ "web": {
+ "type": "search",
+ "results": [
+ {
+ "title": "LiteLLM",
+ "url": "https://example.com/litellm",
+ "description": "Call every LLM API",
+ "page_age": "2024-03-02T10:00:00",
+ "age": "January 1, 2020",
+ "fetched_content_timestamp": 1700000000,
+ "meta_url": {"hostname": "example.com"},
+ },
+ {"title": "Docs", "url": "https://example.com/docs", "description": "Docs", "page_age": None},
+ ],
+ },
+ "news": {
+ "results": [
+ {"title": "News", "url": "https://example.com/news", "description": "Fresh", "age": "March 5, 2024"}
+ ]
+ },
+ "videos": {"results": [{"title": "Video", "url": "https://example.com/video"}]},
+ }
+ )
+
+ assert response.object == "search"
+ assert response.results == [
+ SearchResult(
+ title="LiteLLM",
+ url="https://example.com/litellm",
+ snippet="Call every LLM API",
+ date="2024-03-02",
+ last_updated="2023-11-14",
+ ),
+ SearchResult(title="Docs", url="https://example.com/docs", snippet="Docs"),
+ SearchResult(title="News", url="https://example.com/news", snippet="Fresh", date="2024-03-05"),
+ SearchResult(title="Video", url="https://example.com/video", snippet=""),
+ ]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {},
+ {"type": "search", "query": {"original": "litellm"}},
+ {"web": {}},
+ {"web": {"results": []}, "news": {"results": []}},
+ {"web": {"results": ""}},
+ {"web": {"results": {}}},
+ ],
+)
+def test_brave_response_without_results_is_empty(payload):
+ assert _transform(payload).results == []
+
+
+def test_brave_result_missing_optional_fields_defaults_to_empty_strings_and_no_dates():
+ assert _transform({"web": {"results": [{"profile": {"name": "Example"}}]}}).results == [
+ SearchResult(title="", url="", snippet="", date=None, last_updated=None)
+ ]
+
+
+def test_brave_count_limits_results_across_sections_without_reading_entries_past_the_limit():
+ response = _transform(
+ {
+ "web": {"results": [_entry(1), _entry(2), "not a result object"]},
+ "news": {"results": ["not a result object"]},
+ },
+ query="q=litellm&count=2",
+ )
+
+ assert [result.title for result in response.results] == ["Title 1", "Title 2"]
+
+
+def test_brave_result_filter_restricts_the_sections_that_are_read():
+ response = _transform(
+ {"web": {"results": [_entry(1)]}, "news": {"results": [_entry(2)]}, "videos": None},
+ query="q=litellm&result_filter=news",
+ )
+
+ assert [result.title for result in response.results] == ["Title 2"]
+
+
+@pytest.mark.parametrize("field", ["title", "url", "description"])
+@pytest.mark.parametrize("value", [None, 7, ["text"]])
+def test_brave_result_with_non_string_field_is_rejected(field, value):
+ with pytest.raises(ValidationError):
+ _transform({"web": {"results": [{**_entry(1), field: value}]}})
+
+
+def test_brave_result_with_several_non_string_fields_reports_every_field():
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"web": {"results": [{"title": 7, "url": None, "description": ["text"]}]}})
+
+ assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",)]
+
+
+@pytest.mark.parametrize("payload", [{"web": {"results": [_entry(1)]}}, {}, [], "plain text body"])
+def test_brave_non_integer_count_is_reported_whatever_the_response_body_is(payload):
+ with pytest.raises(ValueError, match="invalid literal for int"):
+ _transform(payload, query="q=litellm&count=abc")
diff --git a/tests/unit/llms/cohere/embed/test_transformation.py b/tests/unit/llms/cohere/embed/test_transformation.py
index d7a2e604fc6..c0f355cfb57 100644
--- a/tests/unit/llms/cohere/embed/test_transformation.py
+++ b/tests/unit/llms/cohere/embed/test_transformation.py
@@ -6,6 +6,7 @@ from typing import Final
import httpx
import pytest
import respx
+from pydantic import ValidationError
import litellm
from litellm.caching.llm_caching_handler import LLMClientCache
@@ -13,6 +14,144 @@ from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.cohere.embed.transformation import CohereEmbeddingConfig
from litellm.types.utils import EmbeddingResponse
+FLOAT_VECTORS = [[0.1, 0.2], [0.3, 0.4]]
+INT8_VECTORS = [[1, -2], [3, 4]]
+
+
+def _transform(payload: object) -> EmbeddingResponse:
+ logging_obj = Logging(
+ model="embed-english-v3.0",
+ messages=[],
+ stream=False,
+ call_type="embedding",
+ start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+ litellm_call_id="call-id",
+ function_id="function-id",
+ )
+ logging_obj.model_call_details["input"] = ["hello", "world"]
+ return CohereEmbeddingConfig().transform_embedding_response(
+ model="embed-english-v3.0",
+ raw_response=httpx.Response(200, json=payload),
+ model_response=EmbeddingResponse(),
+ logging_obj=logging_obj,
+ api_key="sk-test",
+ request_data={"model": "embed-english-v3.0", "texts": ["hello", "world"]},
+ optional_params={},
+ litellm_params={},
+ )
+
+
+def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]:
+ assert response.usage is not None
+ return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens)
+
+
+def test_cohere_embedding_response_lists_every_embedding_type_in_provider_order():
+ response = _transform(
+ {
+ "id": "embed-id",
+ "embeddings": {"float": FLOAT_VECTORS, "int8": INT8_VECTORS},
+ "meta": {"billed_units": {"input_tokens": 5}},
+ }
+ )
+
+ assert response.object == "list"
+ assert response.model == "embed-english-v3.0"
+ assert response.data == [
+ {"object": "embedding", "index": 0, "embedding": [0.1, 0.2]},
+ {"object": "embedding", "index": 1, "embedding": [0.3, 0.4]},
+ {"object": "embedding", "index": 0, "embedding": [1, -2]},
+ {"object": "embedding", "index": 1, "embedding": [3, 4]},
+ ]
+ assert _token_counts(response) == (5, 0, 5)
+
+
+@pytest.mark.parametrize(
+ ("meta", "expected_counts", "expected_details"),
+ [
+ ({"billed_units": {"input_tokens": 5}}, (5, 0, 5), (5, None)),
+ ({"billed_units": {"images": 3}}, (3, 0, 3), (None, 3)),
+ ({"billed_units": {"input_tokens": 5, "images": 3}}, (8, 0, 8), (5, 3)),
+ ({"billed_units": {"input_tokens": 0}}, (0, 0, 0), (0, None)),
+ (
+ {"api_version": {"version": "2"}, "billed_units": {"input_tokens": 2, "output_tokens": 9}},
+ (2, 0, 2),
+ (2, None),
+ ),
+ ],
+)
+def test_cohere_embedding_usage_comes_from_billed_units(meta, expected_counts, expected_details):
+ response = _transform({"embeddings": {"float": FLOAT_VECTORS}, "meta": meta})
+
+ assert response.usage is not None
+ details = response.usage.prompt_tokens_details
+ assert details is not None
+ assert _token_counts(response) == expected_counts
+ assert (details.text_tokens, details.image_tokens) == expected_details
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {"embeddings": {"float": FLOAT_VECTORS}},
+ {"embeddings": {"float": FLOAT_VECTORS}, "meta": {}},
+ {"embeddings": {"float": FLOAT_VECTORS}, "meta": {"billed_units": {}}},
+ {"embeddings": {"float": FLOAT_VECTORS}, "meta": {"billed_units": {"input_tokens": None, "images": None}}},
+ ],
+)
+def test_cohere_embedding_usage_counts_input_tokens_when_billed_units_are_absent(payload):
+ response = _transform(payload)
+
+ assert response.usage is not None
+ assert response.usage.prompt_tokens_details is None
+ assert _token_counts(response) == (2, 0, 2)
+
+
+@pytest.mark.parametrize(
+ ("embeddings", "expected_vectors"),
+ [
+ ({}, []),
+ ({"float": []}, []),
+ ({"float": [], "int8": INT8_VECTORS}, INT8_VECTORS),
+ ({"base64": ["AAAA", "BBBB"]}, ["AAAA", "BBBB"]),
+ ({"float": [[0.1], "x", None, 3]}, [[0.1], "x", None, 3]),
+ ],
+)
+def test_cohere_embedding_rows_pass_through_whatever_each_type_lists(embeddings, expected_vectors):
+ response = _transform({"embeddings": embeddings, "meta": {"billed_units": {"input_tokens": 1}}})
+
+ assert [row["embedding"] for row in response.data] == expected_vectors
+
+
+@pytest.mark.parametrize("embeddings", [FLOAT_VECTORS, [], None, "vectors", 3])
+def test_cohere_embeddings_that_are_not_keyed_by_type_are_rejected(embeddings):
+ with pytest.raises(ValidationError):
+ _transform({"embeddings": embeddings, "meta": {"billed_units": {"input_tokens": 1}}})
+
+
+@pytest.mark.parametrize("vectors", [None, 3, False, 1.5])
+def test_cohere_embedding_type_without_a_list_of_vectors_is_rejected(vectors):
+ with pytest.raises(ValidationError):
+ _transform({"embeddings": {"float": vectors}, "meta": {"billed_units": {"input_tokens": 1}}})
+
+
+@pytest.mark.parametrize("meta", [None, [], "meta", 3])
+def test_cohere_embedding_meta_that_is_not_an_object_is_rejected(meta):
+ with pytest.raises(ValidationError):
+ _transform({"embeddings": {"float": FLOAT_VECTORS}, "meta": meta})
+
+
+@pytest.mark.parametrize("payload", [[1], "text", 3, True])
+def test_cohere_embedding_body_that_is_not_an_object_is_rejected(payload):
+ with pytest.raises(ValidationError):
+ _transform(payload)
+
+
+def test_cohere_embedding_body_without_embeddings_is_rejected():
+ with pytest.raises(KeyError):
+ _transform({"meta": {"billed_units": {"input_tokens": 1}}})
+
+
COHERE_V2_EMBED_URL: Final = "https://api.cohere.ai/v2/embed"
diff --git a/tests/unit/llms/cometapi/embed/__init__.py b/tests/unit/llms/cometapi/embed/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/cometapi/embed/test_cometapi_embed_transformation.py b/tests/unit/llms/cometapi/embed/test_cometapi_embed_transformation.py
new file mode 100644
index 00000000000..d59f6b4e153
--- /dev/null
+++ b/tests/unit/llms/cometapi/embed/test_cometapi_embed_transformation.py
@@ -0,0 +1,89 @@
+import datetime
+
+import httpx
+import pytest
+from pydantic import ValidationError
+
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.cometapi.embed.transformation import CometAPIEmbeddingConfig
+from litellm.types.utils import EmbeddingResponse
+
+EMBEDDING_ROW = {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}
+
+
+def _transform(payload: object) -> EmbeddingResponse:
+ logging_obj = Logging(
+ model="cometapi/text-embedding-3-small",
+ messages=[],
+ stream=False,
+ call_type="embedding",
+ start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+ litellm_call_id="call-id",
+ function_id="function-id",
+ )
+ return CometAPIEmbeddingConfig().transform_embedding_response(
+ model="text-embedding-3-small",
+ raw_response=httpx.Response(200, json=payload),
+ model_response=EmbeddingResponse(),
+ logging_obj=logging_obj,
+ api_key="sk-test",
+ request_data={},
+ optional_params={},
+ litellm_params={},
+ )
+
+
+def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]:
+ assert response.usage is not None
+ return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens)
+
+
+def test_cometapi_embedding_response_carries_provider_fields_and_usage():
+ response = _transform(
+ {
+ "object": "list",
+ "model": "text-embedding-3-small",
+ "data": [EMBEDDING_ROW],
+ "usage": {"prompt_tokens": 8, "total_tokens": 9},
+ }
+ )
+
+ assert response.model == "text-embedding-3-small"
+ assert response.object == "list"
+ assert response.data == [EMBEDDING_ROW]
+ assert _token_counts(response) == (8, 0, 9)
+
+
+@pytest.mark.parametrize(
+ ("payload", "expected"),
+ [
+ ({"data": [EMBEDDING_ROW]}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 4}}, (4, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 6}}, (0, 0, 6)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": None, "total_tokens": None}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": "", "total_tokens": []}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": "12", "total_tokens": 12.0}}, (12, 0, 12)),
+ (
+ {"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 3, "total_tokens": 5, "completion_tokens": 2}},
+ (3, 0, 5),
+ ),
+ ],
+)
+def test_cometapi_embedding_usage_tolerates_missing_null_and_loosely_typed_counts(payload, expected):
+ response = _transform(payload)
+
+ assert response.data == [EMBEDDING_ROW]
+ assert _token_counts(response) == expected
+
+
+@pytest.mark.parametrize("usage", [None, [], "tokens", 3])
+def test_cometapi_embedding_usage_that_is_not_an_object_is_rejected(usage):
+ with pytest.raises(ValidationError):
+ _transform({"data": [EMBEDDING_ROW], "usage": usage})
+
+
+@pytest.mark.parametrize("usage", [{"prompt_tokens": "many"}, {"total_tokens": 1.5}, {"prompt_tokens": [1]}])
+def test_cometapi_embedding_token_count_that_is_not_an_integer_is_rejected(usage):
+ with pytest.raises(ValidationError):
+ _transform({"data": [EMBEDDING_ROW], "usage": usage})
diff --git a/tests/unit/llms/dataforseo/search/test_dataforseo_search_transformation.py b/tests/unit/llms/dataforseo/search/test_dataforseo_search_transformation.py
index 0f20c5b1108..d94f6a53c89 100644
--- a/tests/unit/llms/dataforseo/search/test_dataforseo_search_transformation.py
+++ b/tests/unit/llms/dataforseo/search/test_dataforseo_search_transformation.py
@@ -2,13 +2,18 @@
Unit tests for DataForSEO Search functionality.
"""
+import datetime
import os
from unittest.mock import AsyncMock, patch
+import httpx
import pytest
+from pydantic import ValidationError
import litellm
+from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
+from litellm.llms.dataforseo.search.transformation import DataForSEOSearchConfig
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@@ -54,6 +59,119 @@ async def test_dataforseo_search_basic():
assert len(response.results[0].snippet) > 0
+def _transform(payload: object) -> SearchResponse:
+ logging_obj = Logging(
+ model="dataforseo/search",
+ messages=[],
+ stream=False,
+ call_type="search",
+ start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+ litellm_call_id="call-id",
+ function_id="function-id",
+ )
+ return DataForSEOSearchConfig().transform_search_response(
+ raw_response=httpx.Response(200, json=payload), logging_obj=logging_obj
+ )
+
+
+def _successful_task(result: object) -> dict[str, object]:
+ return {"status_code": 20000, "status_message": "Ok.", "tasks": [{"status_code": 20000, "result": result}]}
+
+
+def test_dataforseo_keeps_only_organic_items_of_the_first_result_in_provider_order():
+ response = _transform(
+ _successful_task(
+ [
+ {
+ "keyword": "litellm",
+ "items_count": 4,
+ "items": [
+ {
+ "type": "organic",
+ "rank_absolute": 1,
+ "title": "LiteLLM",
+ "url": "https://example.com/litellm",
+ "description": "Call every LLM API",
+ "links": None,
+ },
+ {"type": "people_also_ask", "items": [{"title": "What is LiteLLM?"}]},
+ {"type": "paid", "title": "Ad", "url": "https://example.com/ad", "description": "Sponsored"},
+ {
+ "type": "organic",
+ "rank_absolute": 4,
+ "title": "Docs",
+ "url": "https://example.com/docs",
+ "description": "Docs",
+ },
+ ],
+ },
+ "a second result entry is never read",
+ ]
+ )
+ )
+
+ assert response.object == "search"
+ assert response.results == [
+ SearchResult(title="LiteLLM", url="https://example.com/litellm", snippet="Call every LLM API"),
+ SearchResult(title="Docs", url="https://example.com/docs", snippet="Docs"),
+ ]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {},
+ {"status_code": 20000, "tasks": []},
+ {"tasks": [{"status_code": 40501, "status_message": "Invalid Field.", "result": None}]},
+ {"tasks": [{"status_code": 20000}]},
+ _successful_task([]),
+ _successful_task(""),
+ _successful_task([{"keyword": "litellm", "items_count": 0}]),
+ _successful_task([{"items": []}]),
+ _successful_task([{"items": ""}]),
+ _successful_task([{"items": {}}]),
+ _successful_task([{"items": [{"type": "paid", "title": "Ad"}, {"title": "no type"}]}]),
+ [],
+ "plain text body",
+ ],
+)
+def test_dataforseo_response_without_organic_items_is_empty(payload):
+ assert _transform(payload).results == []
+
+
+def test_dataforseo_organic_item_missing_optional_fields_defaults_to_empty_strings():
+ assert _transform(_successful_task([{"items": [{"type": "organic"}]}])).results == [
+ SearchResult(title="", url="", snippet="", date=None, last_updated=None)
+ ]
+
+
+@pytest.mark.parametrize("field", ["title", "url", "description"])
+@pytest.mark.parametrize("value", [None, 7, ["text"]])
+def test_dataforseo_organic_item_with_non_string_field_is_rejected(field, value):
+ item = {"type": "organic", "title": "T", "url": "https://example.com", "description": "D", field: value}
+
+ with pytest.raises(ValidationError):
+ _transform(_successful_task([{"items": [item]}]))
+
+
+def test_dataforseo_organic_item_with_several_non_string_fields_reports_every_field():
+ item = {"type": "organic", "title": 7, "url": None, "description": ["text"]}
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform(_successful_task([{"items": [item]}]))
+
+ assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",)]
+
+
+def test_dataforseo_malformed_item_is_rejected_without_naming_its_position():
+ items = [{}] * 429 + ["not-an-object"]
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform(_successful_task([{"items": items}]))
+
+ assert "429" not in str(exc_info.value)
+
+
@pytest.fixture(autouse=True)
def _vcr_outcome_gate(request, vcr):
install_live_call_probe(request, vcr)
diff --git a/tests/unit/llms/duckduckgo/search/test_duckduckgo_search_transformation.py b/tests/unit/llms/duckduckgo/search/test_duckduckgo_search_transformation.py
index 68419891166..484a724ad7c 100644
--- a/tests/unit/llms/duckduckgo/search/test_duckduckgo_search_transformation.py
+++ b/tests/unit/llms/duckduckgo/search/test_duckduckgo_search_transformation.py
@@ -5,8 +5,11 @@ from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import respx
+from pydantic import ValidationError
import litellm
+from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
+from litellm.llms.duckduckgo.search.transformation import DuckDuckGoSearchConfig
class TestDuckDuckGoSearchMocked:
@@ -297,3 +300,133 @@ def test_duckduckgo_sync_search_returns_typed_results_without_a_limit(respx_mock
for result in response.results
)
assert response.results[-1].url == "https://example.com/9"
+
+
+def _transform(payload: object, query: str = "q=python&format=json") -> SearchResponse:
+ request = httpx.Request("GET", f"https://api.duckduckgo.com/?{query}")
+ return DuckDuckGoSearchConfig().transform_search_response(
+ raw_response=httpx.Response(200, json=payload, request=request), logging_obj=None
+ )
+
+
+def _topic(name: str) -> dict[str, object]:
+ return {
+ "FirstURL": f"https://duckduckgo.com/{name}",
+ "Icon": {"Height": "", "URL": "", "Width": ""},
+ "Result": f'{name}',
+ "Text": f"{name} - About {name}.",
+ }
+
+
+def test_duckduckgo_abstract_comes_first_then_related_and_nested_topics_in_provider_order():
+ long_text = "word " * 12
+ response = _transform(
+ {
+ "Abstract": "",
+ "AbstractSource": "Wikipedia",
+ "AbstractText": "Python is a high-level programming language.",
+ "AbstractURL": "https://en.wikipedia.org/wiki/Python",
+ "Answer": "",
+ "Heading": "Python",
+ "ImageHeight": 0,
+ "Infobox": "",
+ "RelatedTopics": [
+ _topic("Python_programming"),
+ {"FirstURL": "https://duckduckgo.com/Long", "Text": long_text},
+ {"Name": "Related Topics", "Topics": [_topic("Indus")]},
+ ],
+ "Results": [],
+ "Type": "A",
+ "meta": {"attribution": None, "src_options": {"is_fanon": 0}},
+ }
+ )
+
+ assert response.object == "search"
+ assert response.results == [
+ SearchResult(
+ title="Python",
+ url="https://en.wikipedia.org/wiki/Python",
+ snippet="Python is a high-level programming language.",
+ ),
+ SearchResult(
+ title="Python_programming",
+ url="https://duckduckgo.com/Python_programming",
+ snippet="About Python_programming.",
+ ),
+ SearchResult(title=long_text[:50] + "...", url="https://duckduckgo.com/Long", snippet=long_text),
+ SearchResult(title="Indus", url="https://duckduckgo.com/Indus", snippet="About Indus."),
+ ]
+
+
+def test_duckduckgo_abstract_without_heading_gets_an_empty_title():
+ assert _transform({"AbstractURL": "https://example.com", "AbstractText": "Summary"}).results == [
+ SearchResult(title="", url="https://example.com", snippet="Summary")
+ ]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {},
+ {"Abstract": "", "AbstractText": "", "AbstractURL": "", "Heading": "", "RelatedTopics": [], "Results": []},
+ {"AbstractURL": "https://example.com", "AbstractText": "", "Heading": "Disambiguation"},
+ {"AbstractURL": "", "AbstractText": "Summary without a link", "Heading": None},
+ {"RelatedTopics": ""},
+ {"RelatedTopics": {}},
+ {"RelatedTopics": ["not a topic", None, 3, ["nested"]]},
+ {"RelatedTopics": [{"Name": "Group without topics"}, {"FirstURL": "https://example.com"}, {"Text": "no url"}]},
+ {"RelatedTopics": [{"Name": "Empty group", "Topics": []}]},
+ ],
+)
+def test_duckduckgo_response_without_abstract_or_linked_topics_is_empty(payload):
+ assert _transform(payload).results == []
+
+
+def test_duckduckgo_non_object_related_topics_are_skipped_and_the_rest_are_kept():
+ response = _transform({"RelatedTopics": ["not a topic", _topic("First"), None, _topic("Second")]})
+
+ assert [result.title for result in response.results] == ["First", "Second"]
+
+
+@pytest.mark.parametrize(
+ ("query", "expected_titles"),
+ [
+ ("q=python&_max_results=1", ["Python"]),
+ ("q=python&_max_results=2", ["Python", "First"]),
+ ("q=python&_max_results=3", ["Python", "First", "Nested"]),
+ ("q=python&_max_results=not-a-number", ["Python", "First", "Nested", "Second"]),
+ ("q=python", ["Python", "First", "Nested", "Second"]),
+ ],
+)
+def test_duckduckgo_max_results_query_param_limits_the_results_including_the_abstract(query, expected_titles):
+ response = _transform(
+ {
+ "AbstractURL": "https://en.wikipedia.org/wiki/Python",
+ "AbstractText": "Python is a language.",
+ "Heading": "Python",
+ "RelatedTopics": [_topic("First"), {"Topics": [_topic("Nested")]}, _topic("Second")],
+ },
+ query=query,
+ )
+
+ assert [result.title for result in response.results] == expected_titles
+
+
+@pytest.mark.parametrize(
+ ("field", "value"),
+ [("Heading", None), ("Heading", 7), ("AbstractURL", 7), ("AbstractText", ["Summary"])],
+)
+def test_duckduckgo_abstract_with_non_string_field_is_rejected(field, value):
+ payload = {"AbstractURL": "https://example.com", "AbstractText": "Summary", "Heading": "Title", field: value}
+
+ with pytest.raises(ValidationError):
+ _transform(payload)
+
+
+def test_duckduckgo_abstract_with_several_non_string_fields_reports_every_field():
+ payload = {"Heading": 7, "AbstractURL": ["https://example.com"], "AbstractText": {"text": "Summary"}}
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform(payload)
+
+ assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",)]
diff --git a/tests/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py
index 6b556ff241f..abcd33869b7 100644
--- a/tests/unit/llms/exa_ai/search/test_transformation.py
+++ b/tests/unit/llms/exa_ai/search/test_transformation.py
@@ -5,8 +5,10 @@ from unittest.mock import Mock
import httpx
import pytest
import respx
+from pydantic import ValidationError
import litellm
+from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
from litellm.llms.exa_ai.search.transformation import ExaAISearchConfig
EXA_SEARCH_URL: Final = "https://api.exa.ai/search"
@@ -88,3 +90,110 @@ def test_search_sends_max_results_as_num_results(respx_mock: respx.MockRouter) -
}
assert route.calls.last.request.headers["x-api-key"] == "test-exa-key"
assert [result.url for result in response.results] == ["https://example.com/ml"]
+
+
+def _transform(payload: object) -> SearchResponse:
+ return ExaAISearchConfig().transform_search_response(httpx.Response(200, json=payload), logging_obj=Mock())
+
+
+def test_transform_search_response_maps_every_exa_result_in_provider_order():
+ response: Final = _transform(
+ {
+ "requestId": "b5947044c4b78efa9552a7c89b306d95",
+ "resolvedSearchType": "neural",
+ "results": [
+ {
+ "id": "https://example.com/first",
+ "title": "First",
+ "url": "https://example.com/first",
+ "publishedDate": "2024-01-05T00:00:00.000Z",
+ "author": "Ada",
+ "score": 0.42,
+ "text": "First text",
+ "highlights": ["ignored when text is present"],
+ "highlightScores": [0.5],
+ "image": None,
+ },
+ {"title": "Second", "url": "https://example.com/second", "highlights": ["one", "two"]},
+ {"title": "Third", "url": "https://example.com/third", "summary": "Summary", "publishedDate": None},
+ ],
+ "costDollars": {"total": 0.005},
+ }
+ )
+
+ assert response.object == "search"
+ assert response.results == [
+ SearchResult(
+ title="First", url="https://example.com/first", snippet="First text", date="2024-01-05T00:00:00.000Z"
+ ),
+ SearchResult(title="Second", url="https://example.com/second", snippet="one\n\ntwo"),
+ SearchResult(title="Third", url="https://example.com/third", snippet="Summary"),
+ ]
+
+
+def test_transform_search_response_result_without_optional_fields_gets_empty_defaults():
+ assert _transform({"results": [{}]}).results == [SearchResult(title="", url="", snippet="")]
+
+
+@pytest.mark.parametrize(
+ ("highlights", "expected_snippet"),
+ [
+ (None, "a summary"),
+ ("", "a summary"),
+ ({}, "a summary"),
+ ([""], "a summary"),
+ (["", ""], "\n\n"),
+ ("abc", "a\n\nb\n\nc"),
+ ({"first": 1, "second": 2}, "first\n\nsecond"),
+ ],
+)
+def test_transform_search_response_joins_whatever_the_highlights_value_iterates_over(
+ highlights: object, expected_snippet: str
+):
+ response: Final = _transform(
+ {
+ "results": [
+ {"title": "Title", "url": "https://example.com", "highlights": highlights, "summary": "a summary"}
+ ]
+ }
+ )
+
+ assert response.results[0].snippet == expected_snippet
+
+
+@pytest.mark.parametrize("payload", [{}, {"results": []}, {"results": ""}, {"results": {}}, {"requestId": "abc"}])
+def test_transform_search_response_without_results_is_empty(payload: dict[str, object]):
+ assert _transform(payload).results == []
+
+
+@pytest.mark.parametrize(
+ "result",
+ [
+ {"title": None, "url": "https://example.com"},
+ {"title": "Title", "url": 7},
+ {"title": "Title", "url": "https://example.com", "text": 7},
+ {"title": "Title", "url": "https://example.com", "summary": ["a summary"]},
+ {"title": "Title", "url": "https://example.com", "publishedDate": 1704412800},
+ ],
+)
+def test_transform_search_response_rejects_result_with_wrongly_typed_field(result: dict[str, object]):
+ with pytest.raises(ValidationError):
+ _transform({"results": [result]})
+
+
+def test_transform_search_response_reports_every_wrongly_typed_field_of_a_result():
+ result: Final = {"title": 7, "url": None, "text": ["text"], "publishedDate": 1704412800}
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"results": [result]})
+
+ assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",), ("date",)]
+
+
+def test_transform_search_response_rejects_malformed_result_without_naming_its_position():
+ results = [{}] * 429 + ["not-an-object"]
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"results": results})
+
+ assert "429" not in str(exc_info.value)
diff --git a/tests/unit/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py b/tests/unit/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py
index 6c55760b625..11078f64a9f 100644
--- a/tests/unit/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py
+++ b/tests/unit/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py
@@ -6,7 +6,9 @@ from pathlib import Path
import httpx
import pytest
+from pydantic import ValidationError
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.fal_ai.image_edit import FalAIImageEditConfig
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.router import GenericLiteLLMParams
@@ -156,3 +158,69 @@ def test_transform_request_requires_an_image(image):
litellm_params=GenericLiteLLMParams(),
headers={},
)
+
+
+def _transform_response(payload: object) -> ImageResponse:
+ return FalAIImageEditConfig().transform_image_edit_response(
+ model="openai/gpt-image-2.5/flare/edit", raw_response=httpx.Response(200, json=payload), logging_obj=None
+ )
+
+
+def test_transform_response_maps_image_objects_and_bare_urls_in_provider_order():
+ response = _transform_response(
+ {
+ "images": [
+ "https://fal.media/a.png",
+ {"b64_json": "QUJD", "width": 0, "height": "768", "content_type": 5},
+ 7,
+ None,
+ {"url": "https://fal.media/b.png", "width": 512},
+ ],
+ "seed": 42,
+ }
+ )
+
+ assert [(image.url, image.b64_json, image.provider_specific_fields) for image in response.data] == [
+ ("https://fal.media/a.png", None, None),
+ (None, "QUJD", None),
+ ("https://fal.media/b.png", None, {"width": 512}),
+ ]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {},
+ {"detail": "rate limit exceeded"},
+ {"images": None},
+ {"images": "https://fal.media/a.png"},
+ {"images": {"url": "https://fal.media/a.png"}},
+ {"images": 429},
+ ],
+)
+def test_transform_response_without_an_image_list_is_empty(payload):
+ assert _transform_response(payload).data == []
+
+
+@pytest.mark.parametrize(
+ "payload", [7, True, "https://fal.media/a.png", [{"url": "https://fal.media/a.png"}], [{"images": []}]]
+)
+def test_transform_response_rejects_non_object_bodies_without_echoing_them(payload):
+ with pytest.raises(ValidationError) as exc_info:
+ _transform_response(payload)
+
+ assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [("dict_type", ())]
+ assert "input_value" not in str(exc_info.value)
+ assert "fal.media" not in str(exc_info.value)
+
+
+def test_transform_response_reports_a_body_that_is_not_json_as_a_provider_error():
+ with pytest.raises(BaseLLMException) as exc_info:
+ FalAIImageEditConfig().transform_image_edit_response(
+ model="openai/gpt-image-2.5/flare/edit",
+ raw_response=httpx.Response(502, content=b"502 Bad Gateway"),
+ logging_obj=None,
+ )
+
+ assert exc_info.value.status_code == 502
+ assert str(exc_info.value).startswith("Error parsing Fal AI image edit response: ")
diff --git a/tests/unit/llms/fal_ai/image_generation/test_bria_transformation.py b/tests/unit/llms/fal_ai/image_generation/test_bria_transformation.py
new file mode 100644
index 00000000000..76e9b514805
--- /dev/null
+++ b/tests/unit/llms/fal_ai/image_generation/test_bria_transformation.py
@@ -0,0 +1,93 @@
+import httpx
+import pytest
+from pydantic import ValidationError
+
+from litellm.llms.fal_ai.image_generation.bria_transformation import FalAIBriaConfig
+from litellm.types.utils import ImageObject, ImageResponse
+
+
+def _transform(payload: object, model_response: ImageResponse | None = None) -> ImageResponse:
+ return FalAIBriaConfig().transform_image_generation_response(
+ model="bria/text-to-image/3.2",
+ raw_response=httpx.Response(200, json=payload),
+ model_response=model_response or ImageResponse(),
+ logging_obj=None,
+ request_data={},
+ optional_params={},
+ litellm_params={},
+ encoding=None,
+ )
+
+
+def test_transform_image_generation_response_maps_the_single_image():
+ response = _transform(
+ {
+ "image": {
+ "url": "https://fal.media/a.png",
+ "content_type": "image/png",
+ "file_name": "a.png",
+ "file_size": 123456,
+ "width": 1024,
+ "height": 1024,
+ },
+ "images": [{"url": "https://fal.media/ignored.png"}],
+ }
+ )
+
+ assert [(image.url, image.b64_json, image.provider_specific_fields) for image in response.data] == [
+ ("https://fal.media/a.png", None, None)
+ ]
+
+
+def test_transform_image_generation_response_keeps_an_image_without_a_url():
+ response = _transform({"image": {"content_type": "image/png", "url": None}})
+
+ assert [(image.url, image.b64_json) for image in response.data] == [(None, None)]
+
+
+def test_transform_image_generation_response_appends_to_images_already_on_the_response():
+ model_response = ImageResponse(data=[ImageObject(url="https://fal.media/existing.png")])
+
+ response = _transform({"image": {"url": "https://fal.media/new.png"}}, model_response)
+
+ assert response is model_response
+ assert [image.url for image in response.data] == ["https://fal.media/existing.png", "https://fal.media/new.png"]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {},
+ {"detail": "rate limit exceeded"},
+ {"image": None},
+ {"image": {}},
+ {"image": "https://fal.media/a.png"},
+ {"image": ["https://fal.media/a.png"]},
+ {"image": 429},
+ {"images": [{"url": "https://fal.media/a.png"}]},
+ ],
+)
+def test_transform_image_generation_response_without_an_image_object_is_empty(payload: dict[str, object]):
+ assert _transform(payload).data == []
+
+
+@pytest.mark.parametrize(
+ "payload", [7, True, "https://fal.media/a.png", [{"image": {"url": "https://fal.media/a.png"}}]]
+)
+def test_transform_image_generation_response_rejects_non_object_bodies_without_echoing_them(payload: object):
+ with pytest.raises(ValidationError) as exc_info:
+ _transform(payload)
+
+ assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [("dict_type", ())]
+ assert "input_value" not in str(exc_info.value)
+ assert "fal.media" not in str(exc_info.value)
+
+
+@pytest.mark.parametrize("url", [7, ["Request timed out"], {"code": 429}])
+def test_transform_image_generation_response_reports_a_url_of_the_wrong_type_with_its_value(url: object):
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"image": {"url": url}})
+
+ assert [(error["type"], error["loc"], error["input"]) for error in exc_info.value.errors()] == [
+ ("string_type", ("url",), url)
+ ]
diff --git a/tests/unit/llms/fal_ai/image_generation/test_fal_ai_image_generation_transformation.py b/tests/unit/llms/fal_ai/image_generation/test_fal_ai_image_generation_transformation.py
new file mode 100644
index 00000000000..a7313c32c57
--- /dev/null
+++ b/tests/unit/llms/fal_ai/image_generation/test_fal_ai_image_generation_transformation.py
@@ -0,0 +1,97 @@
+import httpx
+import pytest
+from pydantic import ValidationError
+
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
+from litellm.llms.fal_ai.image_generation.transformation import FalAIImageGenerationConfig
+from litellm.types.utils import ImageObject, ImageResponse
+
+
+def _transform(payload: object, model_response: ImageResponse | None = None) -> ImageResponse:
+ return FalAIImageGenerationConfig().transform_image_generation_response(
+ model="fal-ai/some-new-model",
+ raw_response=httpx.Response(200, json=payload),
+ model_response=model_response or ImageResponse(),
+ logging_obj=None,
+ request_data={},
+ optional_params={},
+ litellm_params={},
+ encoding=None,
+ )
+
+
+def test_transform_image_generation_response_maps_image_objects_and_bare_urls():
+ response = _transform(
+ {
+ "images": [
+ {"url": "https://fal.media/a.png", "width": 1024, "height": 768, "content_type": "image/png"},
+ "https://fal.media/b.png",
+ {"b64_json": "QUJD", "width": 0, "height": "768", "content_type": 5},
+ 7,
+ None,
+ ["https://fal.media/ignored.png"],
+ ],
+ "seed": 42,
+ "timings": {"inference": 2.5},
+ }
+ )
+
+ assert [(image.url, image.b64_json, image.provider_specific_fields) for image in response.data] == [
+ ("https://fal.media/a.png", None, {"width": 1024, "height": 768, "content_type": "image/png"}),
+ ("https://fal.media/b.png", None, None),
+ (None, "QUJD", None),
+ ]
+
+
+def test_transform_image_generation_response_appends_to_images_already_on_the_response():
+ model_response = ImageResponse(data=[ImageObject(url="https://fal.media/existing.png")])
+
+ response = _transform({"images": ["https://fal.media/new.png"]}, model_response)
+
+ assert response is model_response
+ assert [image.url for image in response.data] == ["https://fal.media/existing.png", "https://fal.media/new.png"]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {},
+ {"detail": "rate limit exceeded"},
+ {"images": None},
+ {"images": "https://fal.media/a.png"},
+ {"images": {"url": "https://fal.media/a.png"}},
+ {"images": 429},
+ {"image": {"url": "https://fal.media/a.png"}},
+ ],
+)
+def test_transform_image_generation_response_without_an_image_list_is_empty(payload: dict[str, object]):
+ assert _transform(payload).data == []
+
+
+@pytest.mark.parametrize(
+ "payload", [7, True, "https://fal.media/a.png", [{"url": "https://fal.media/a.png"}], [{"images": []}]]
+)
+def test_transform_image_generation_response_rejects_non_object_bodies_without_echoing_them(payload: object):
+ with pytest.raises(ValidationError) as exc_info:
+ _transform(payload)
+
+ assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [("dict_type", ())]
+ assert "input_value" not in str(exc_info.value)
+ assert "fal.media" not in str(exc_info.value)
+
+
+def test_transform_image_generation_response_reports_a_body_that_is_not_json_as_a_provider_error():
+ with pytest.raises(BaseLLMException) as exc_info:
+ FalAIImageGenerationConfig().transform_image_generation_response(
+ model="fal-ai/some-new-model",
+ raw_response=httpx.Response(502, content=b"502 Bad Gateway"),
+ model_response=ImageResponse(),
+ logging_obj=None,
+ request_data={},
+ optional_params={},
+ litellm_params={},
+ encoding=None,
+ )
+
+ assert exc_info.value.status_code == 502
+ assert str(exc_info.value).startswith("Error transforming image generation response: ")
diff --git a/tests/unit/llms/fal_ai/image_generation/test_imagen4_transformation.py b/tests/unit/llms/fal_ai/image_generation/test_imagen4_transformation.py
new file mode 100644
index 00000000000..4866016a62d
--- /dev/null
+++ b/tests/unit/llms/fal_ai/image_generation/test_imagen4_transformation.py
@@ -0,0 +1,112 @@
+import httpx
+import pytest
+from pydantic import ValidationError
+
+from litellm.llms.fal_ai.image_generation.imagen4_transformation import FalAIImagen4Config
+from litellm.types.utils import ImageObject, ImageResponse
+
+
+def _transform(payload: object, model_response: ImageResponse | None = None) -> ImageResponse:
+ return FalAIImagen4Config().transform_image_generation_response(
+ model="fal-ai/imagen4/preview",
+ raw_response=httpx.Response(200, json=payload),
+ model_response=model_response or ImageResponse(),
+ logging_obj=None,
+ request_data={},
+ optional_params={},
+ litellm_params={},
+ encoding=None,
+ )
+
+
+def test_transform_image_generation_response_maps_image_objects_bare_urls_and_seed():
+ response = _transform(
+ {
+ "images": [
+ {
+ "url": "https://fal.media/a.png",
+ "content_type": "image/png",
+ "file_name": "z9RV14K95DvU.png",
+ "file_size": 4404019,
+ },
+ "https://fal.media/b.png",
+ {"content_type": "image/png"},
+ 7,
+ None,
+ ["https://fal.media/ignored.png"],
+ ],
+ "seed": 42,
+ }
+ )
+
+ assert [(image.url, image.b64_json, image.provider_specific_fields) for image in response.data] == [
+ ("https://fal.media/a.png", None, None),
+ ("https://fal.media/b.png", None, None),
+ (None, None, None),
+ ]
+ assert repr(response._hidden_params["seed"]) == "42"
+
+
+@pytest.mark.parametrize("seed", [0, None, "42", 1.0, [1, 2], {"value": 42}])
+def test_transform_image_generation_response_stores_the_seed_exactly_as_received(seed: object):
+ response = _transform({"images": [], "seed": seed})
+
+ assert repr(response._hidden_params["seed"]) == repr(seed)
+
+
+def test_transform_image_generation_response_appends_to_images_already_on_the_response():
+ model_response = ImageResponse(data=[ImageObject(url="https://fal.media/existing.png")])
+
+ response = _transform({"images": ["https://fal.media/new.png"]}, model_response)
+
+ assert response is model_response
+ assert [image.url for image in response.data] == ["https://fal.media/existing.png", "https://fal.media/new.png"]
+ assert "seed" not in response._hidden_params
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {},
+ {"detail": "rate limit exceeded"},
+ {"images": None},
+ {"images": "https://fal.media/a.png"},
+ {"images": {"url": "https://fal.media/a.png"}},
+ {"images": 429},
+ {"image": {"url": "https://fal.media/a.png"}},
+ ],
+)
+def test_transform_image_generation_response_without_an_image_list_is_empty(payload: dict[str, object]):
+ response = _transform(payload)
+
+ assert response.data == []
+ assert "seed" not in response._hidden_params
+
+
+@pytest.mark.parametrize(
+ "payload", [7, True, "https://fal.media/a.png", [{"url": "https://fal.media/a.png"}], [{"images": [], "seed": 1}]]
+)
+def test_transform_image_generation_response_rejects_non_object_bodies_without_echoing_them(payload: object):
+ with pytest.raises(ValidationError) as exc_info:
+ _transform(payload)
+
+ assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [("dict_type", ())]
+ assert "input_value" not in str(exc_info.value)
+ assert "fal.media" not in str(exc_info.value)
+
+
+@pytest.mark.parametrize("position", [0, 403])
+def test_transform_image_generation_response_keeps_earlier_images_when_a_url_has_the_wrong_type(position: int):
+ model_response = ImageResponse()
+ images = [{"url": f"https://fal.media/{index}.png"} for index in range(position)] + [{"url": 7}]
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"images": images, "seed": 42}, model_response)
+
+ assert [(error["type"], error["loc"], error["input"]) for error in exc_info.value.errors()] == [
+ ("string_type", ("url",), 7)
+ ]
+ assert [image.url for image in model_response.data] == [
+ f"https://fal.media/{index}.png" for index in range(position)
+ ]
+ assert "seed" not in model_response._hidden_params
diff --git a/tests/unit/llms/fastcrw/search/test_transformation.py b/tests/unit/llms/fastcrw/search/test_transformation.py
index adf8fec087c..a8becec7840 100644
--- a/tests/unit/llms/fastcrw/search/test_transformation.py
+++ b/tests/unit/llms/fastcrw/search/test_transformation.py
@@ -1,9 +1,12 @@
import os
from unittest.mock import Mock, patch
+import httpx
import pytest
+from pydantic import ValidationError
import litellm
+from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
from litellm.llms.fastcrw.search.transformation import FastCRWSearchConfig
@@ -180,3 +183,102 @@ def test_transform_search_response_non_list_data():
_resp({"success": True, "data": {"unexpected": "shape"}}), logging_obj=Mock()
)
assert resp.results == []
+
+
+def _transform(payload: object) -> SearchResponse:
+ return _config().transform_search_response(
+ raw_response=httpx.Response(200, json=payload), logging_obj=None
+ )
+
+
+def _entry(index: int) -> dict[str, str]:
+ return {
+ "title": f"Title {index}",
+ "url": f"https://example.com/{index}",
+ "description": f"Description {index}",
+ }
+
+
+def test_transform_search_response_keeps_provider_order_and_snippet_fallbacks():
+ response = _transform(
+ {
+ "success": True,
+ "warning": None,
+ "data": [
+ {**_entry(1), "markdown": "# Page 1", "metadata": {"statusCode": 200}},
+ {**_entry(2), "markdown": ""},
+ {**_entry(3), "markdown": None},
+ {"title": "Title 4", "url": "https://example.com/4"},
+ {"metadata": {"statusCode": 200}},
+ ],
+ }
+ )
+
+ assert response.object == "search"
+ assert response.results == [
+ SearchResult(title="Title 1", url="https://example.com/1", snippet="# Page 1"),
+ SearchResult(
+ title="Title 2", url="https://example.com/2", snippet="Description 2"
+ ),
+ SearchResult(
+ title="Title 3", url="https://example.com/3", snippet="Description 3"
+ ),
+ SearchResult(title="Title 4", url="https://example.com/4", snippet=""),
+ SearchResult(title="", url="", snippet="", date=None, last_updated=None),
+ ]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {},
+ {"success": False, "error": "rate limit"},
+ {"success": True, "data": None},
+ {"success": True, "data": "rate limit"},
+ {"success": True, "data": 429},
+ ],
+)
+def test_transform_search_response_without_a_data_list_is_empty(payload):
+ assert _transform(payload).results == []
+
+
+@pytest.mark.parametrize("payload", [[_entry(1)], "upstream unavailable", True, 429])
+def test_transform_search_response_rejects_a_body_that_is_not_an_object(payload):
+ with pytest.raises(ValidationError) as exc_info:
+ _transform(payload)
+
+ assert [error["type"] for error in exc_info.value.errors()] == ["dict_type"]
+ assert "input_value" not in str(exc_info.value)
+
+
+@pytest.mark.parametrize("position", [0, 1, 403])
+@pytest.mark.parametrize("entry", ["rate limit", None, 429, [_entry(1)]])
+def test_transform_search_response_rejects_an_entry_that_is_not_an_object(
+ position, entry
+):
+ entries = [_entry(index) for index in range(position)] + [entry]
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"success": True, "data": entries})
+
+ assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [
+ ("dict_type", ())
+ ]
+ assert "input_value" not in str(exc_info.value)
+
+
+@pytest.mark.parametrize("position", [0, 403])
+def test_transform_search_response_reports_every_field_of_wrong_type(position):
+ entries = [_entry(index) for index in range(position)] + [
+ {"title": 7, "url": None, "markdown": ["Request timed out"]}
+ ]
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"success": True, "data": entries})
+
+ assert [(error["loc"], error["input"]) for error in exc_info.value.errors()] == [
+ (("title",), 7),
+ (("url",), None),
+ (("snippet",), ["Request timed out"]),
+ ]
+ assert "Request timed out" in str(exc_info.value)
diff --git a/tests/unit/llms/firecrawl/search/test_firecrawl_search_transformation.py b/tests/unit/llms/firecrawl/search/test_firecrawl_search_transformation.py
new file mode 100644
index 00000000000..571466ac86f
--- /dev/null
+++ b/tests/unit/llms/firecrawl/search/test_firecrawl_search_transformation.py
@@ -0,0 +1,160 @@
+import httpx
+import pytest
+from pydantic import ValidationError
+
+from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
+from litellm.llms.firecrawl.search.transformation import FirecrawlSearchConfig
+
+
+def _transform(payload: object) -> SearchResponse:
+ return FirecrawlSearchConfig().transform_search_response(
+ raw_response=httpx.Response(200, json=payload), logging_obj=None
+ )
+
+
+def _entry(index: int) -> dict[str, str]:
+ return {"title": f"Title {index}", "url": f"https://example.com/{index}", "description": f"Description {index}"}
+
+
+def _news(index: int) -> dict[str, str]:
+ return {"title": f"News {index}", "url": f"https://example.com/news/{index}", "snippet": f"Snippet {index}"}
+
+
+def test_firecrawl_cloud_response_lists_web_results_before_news_results():
+ response = _transform(
+ {
+ "success": True,
+ "data": {
+ "web": [
+ {**_entry(1), "markdown": "# Page 1", "metadata": {"statusCode": 200}},
+ {**_entry(2), "markdown": ""},
+ {**_entry(3), "markdown": None, "snippet": "ignored for web results"},
+ {"position": 4},
+ ],
+ "images": [{"title": "Logo", "imageUrl": "https://example.com/logo.png"}],
+ "news": [
+ {**_news(1), "date": "2024-03-05", "markdown": "# News 1"},
+ {**_news(2), "date": None, "description": "ignored for news results"},
+ {"title": "News 3", "url": "https://example.com/news/3", "description": "ignored for news results"},
+ ],
+ },
+ "creditsUsed": 2,
+ }
+ )
+
+ assert response.object == "search"
+ assert response.results == [
+ SearchResult(title="Title 1", url="https://example.com/1", snippet="# Page 1"),
+ SearchResult(title="Title 2", url="https://example.com/2", snippet="Description 2"),
+ SearchResult(title="Title 3", url="https://example.com/3", snippet="Description 3"),
+ SearchResult(title="", url="", snippet=""),
+ SearchResult(title="News 1", url="https://example.com/news/1", snippet="# News 1", date="2024-03-05"),
+ SearchResult(title="News 2", url="https://example.com/news/2", snippet="Snippet 2"),
+ SearchResult(title="News 3", url="https://example.com/news/3", snippet="", date=None, last_updated=None),
+ ]
+
+
+def test_firecrawl_self_hosted_response_maps_the_flat_data_list():
+ response = _transform(
+ {
+ "success": True,
+ "data": [
+ {**_entry(1), "markdown": "# Page 1"},
+ {**_entry(2), "date": "2024-03-05"},
+ {"metadata": {"statusCode": 200}},
+ ],
+ }
+ )
+
+ assert response.results == [
+ SearchResult(title="Title 1", url="https://example.com/1", snippet="# Page 1"),
+ SearchResult(title="Title 2", url="https://example.com/2", snippet="Description 2", date=None),
+ SearchResult(title="", url="", snippet=""),
+ ]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {},
+ {"success": False, "error": "rate limit"},
+ {"data": None},
+ {"data": "rate limit"},
+ {"data": 429},
+ {"data": []},
+ {"data": {}},
+ {"data": {"images": [{"imageUrl": "https://example.com/logo.png"}]}},
+ {"data": {"web": [], "news": []}},
+ {"data": {"web": "", "news": {}}},
+ ],
+)
+def test_firecrawl_response_without_web_or_news_results_is_empty(payload):
+ assert _transform(payload).results == []
+
+
+@pytest.mark.parametrize("payload", [[_entry(1)], "upstream unavailable", True, 429])
+def test_firecrawl_body_that_is_not_an_object_is_rejected_without_echoing_it(payload):
+ with pytest.raises(ValidationError) as exc_info:
+ _transform(payload)
+
+ assert [error["type"] for error in exc_info.value.errors()] == ["dict_type"]
+ assert "input_value" not in str(exc_info.value)
+
+
+@pytest.mark.parametrize("section", ["web", "news"])
+@pytest.mark.parametrize("value", [None, 429, True])
+def test_firecrawl_section_that_cannot_be_iterated_is_rejected_without_echoing_it(section, value):
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"data": {"web": [_entry(1)], "news": [_news(1)], section: value}})
+
+ assert [error["type"] for error in exc_info.value.errors()] == ["iterable_type"]
+ assert "input_value" not in str(exc_info.value)
+
+
+def _payload_with_entries(shape: str, entries: list[object]) -> dict[str, object]:
+ if shape == "self-hosted":
+ return {"data": entries}
+ if shape == "web":
+ return {"data": {"web": entries}}
+ return {"data": {"web": [_entry(0)], "news": entries}}
+
+
+@pytest.mark.parametrize("shape", ["self-hosted", "web", "news"])
+@pytest.mark.parametrize("position", [0, 403])
+@pytest.mark.parametrize("entry", ["rate limit", None, 429, [_entry(1)]])
+def test_firecrawl_entry_that_is_not_an_object_is_rejected_without_its_position(shape, position, entry):
+ entries = [_entry(index) for index in range(position)] + [entry]
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform(_payload_with_entries(shape, entries))
+
+ assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [("dict_type", ())]
+ assert "input_value" not in str(exc_info.value)
+
+
+@pytest.mark.parametrize("shape", ["self-hosted", "web", "news"])
+@pytest.mark.parametrize("position", [0, 403])
+def test_firecrawl_entry_with_several_fields_of_wrong_type_reports_every_field(shape, position):
+ entries = [_entry(index) for index in range(position)] + [
+ {"title": 7, "url": None, "markdown": ["Request timed out"]}
+ ]
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform(_payload_with_entries(shape, entries))
+
+ assert [(error["loc"], error["input"]) for error in exc_info.value.errors()] == [
+ (("title",), 7),
+ (("url",), None),
+ (("snippet",), ["Request timed out"]),
+ ]
+ assert "Request timed out" in str(exc_info.value)
+
+
+def test_firecrawl_news_entry_with_non_string_date_reports_the_date_next_to_the_other_fields():
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"data": {"news": [{**_news(1), "title": 7, "date": 1709596800}]}})
+
+ assert [(error["loc"], error["input"]) for error in exc_info.value.errors()] == [
+ (("title",), 7),
+ (("date",), 1709596800),
+ ]
diff --git a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
index 7b25484ea90..16e2221dc84 100644
--- a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
+++ b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
@@ -5,10 +5,11 @@ from unittest.mock import MagicMock, patch
import httpx
import pytest
import respx
+from pydantic import ValidationError
import litellm
-from litellm.litellm_core_utils.get_supported_openai_params import get_supported_openai_params
from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
+from litellm.litellm_core_utils.get_supported_openai_params import get_supported_openai_params
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig
from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id
@@ -2011,6 +2012,121 @@ VISION_MODEL = next(
)
+_CHAT_REPLY: Final = {
+ "id": "chat-1",
+ "object": "chat.completion",
+ "created": 1,
+ "model": "accounts/fireworks/models/llama-v3p1-8b-instruct",
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": "pong"}, "finish_reason": "stop"}],
+ "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
+}
+_FIREWORKS_BASE: Final = "https://api.fireworks.ai/inference/v1"
+_MODELS_URL: Final = f"{_FIREWORKS_BASE}/accounts/acct/models"
+
+
+def _complete_against_reply(reply: object) -> ModelResponse:
+ return litellm.completion(
+ model="fireworks_ai/accounts/fireworks/models/llama-v3p1-8b-instruct",
+ messages=[{"role": "user", "content": "ping"}],
+ api_key="fw-test-key",
+ client=HTTPHandler(
+ client=httpx.Client(
+ transport=httpx.MockTransport(lambda request: httpx.Response(200, json=reply, request=request))
+ )
+ ),
+ )
+
+
+def test_completion_keeps_unmodelled_reply_fields_and_nested_fireworks_metrics() -> None:
+ perf_metrics: Final = {"ttft": 0.1, "per_stage": [{"name": "prefill", "ms": None}]}
+
+ response: Final = _complete_against_reply(
+ {
+ **_CHAT_REPLY,
+ "system_fingerprint": None,
+ "service_tier": "default",
+ "x-trace": {"spans": [1, 2]},
+ "perf_metrics": perf_metrics,
+ "prompt_token_ids": [7, 8, 9],
+ }
+ )
+
+ assert response.model == "fireworks_ai/accounts/fireworks/models/llama-v3p1-8b-instruct"
+ assert response.choices[0].message.content == "pong"
+ assert response.model_dump()["service_tier"] == "default"
+ assert response.model_dump()["x-trace"] == {"spans": [1, 2]}
+ assert response._hidden_params["fireworks_perf_metrics"] == perf_metrics
+ assert response._hidden_params["fireworks_prompt_token_ids"] == [7, 8, 9]
+
+
+@pytest.mark.parametrize(
+ "reply",
+ [pytest.param([_CHAT_REPLY], id="list"), pytest.param("pong", id="string"), pytest.param(7, id="number")],
+)
+def test_completion_with_non_object_reply_raises_api_connection_error(reply: object) -> None:
+ with pytest.raises(litellm.APIConnectionError):
+ _complete_against_reply(reply)
+
+
+@pytest.mark.parametrize(
+ ("body", "expected"),
+ [
+ pytest.param(
+ {"models": [{"name": "accounts/acct/models/a", "state": {"ready": [True]}}, {"name": "b"}], "totalSize": 2},
+ ["fireworks_ai/accounts/acct/models/a", "fireworks_ai/b"],
+ id="listed-models",
+ ),
+ pytest.param({"models": []}, [], id="empty-list"),
+ pytest.param({"models": ""}, [], id="models-empty-string"),
+ pytest.param({"models": {}}, [], id="models-empty-object"),
+ pytest.param({"models": [{"name": ""}]}, ["fireworks_ai/"], id="empty-name"),
+ ],
+)
+def test_get_models_prefixes_each_listed_name(
+ respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, body: dict[str, object], expected: list[str]
+) -> None:
+ monkeypatch.setenv("FIREWORKS_ACCOUNT_ID", "acct")
+ respx_mock.get(_MODELS_URL).respond(200, json=body)
+
+ assert FireworksAIConfig().get_models(api_key="fw-test-key", api_base=_FIREWORKS_BASE) == expected
+
+
+@pytest.mark.parametrize(
+ ("body", "missing_key"),
+ [
+ pytest.param({"data": []}, "models", id="no-models-key"),
+ pytest.param({"models": [{"name": "a"}, {"id": "b"}]}, "name", id="entry-without-name"),
+ ],
+)
+def test_get_models_raises_key_error_for_missing_field(
+ respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, body: dict[str, object], missing_key: str
+) -> None:
+ monkeypatch.setenv("FIREWORKS_ACCOUNT_ID", "acct")
+ respx_mock.get(_MODELS_URL).respond(200, json=body)
+
+ with pytest.raises(KeyError, match=missing_key):
+ FireworksAIConfig().get_models(api_key="fw-test-key", api_base=_FIREWORKS_BASE)
+
+
+@pytest.mark.parametrize(
+ "body",
+ [
+ pytest.param([{"name": "a"}], id="top-level-list"),
+ pytest.param({"models": None}, id="models-null"),
+ pytest.param({"models": ["a"]}, id="entry-not-an-object"),
+ pytest.param({"models": [{"name": 7}]}, id="name-not-a-string"),
+ ],
+)
+def test_get_models_rejects_malformed_listing(
+ respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, body: object
+) -> None:
+ monkeypatch.setenv("FIREWORKS_ACCOUNT_ID", "acct")
+ respx_mock.get(_MODELS_URL).respond(200, json=body)
+
+ with pytest.raises(ValidationError):
+ FireworksAIConfig().get_models(api_key="fw-test-key", api_base=_FIREWORKS_BASE)
+
+
def _fireworks_chat_client() -> MagicMock:
body: Final = {
"id": "chat-user-attribution",
diff --git a/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py b/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py
index 2a69926414a..138aafcbfcd 100644
--- a/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py
+++ b/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py
@@ -460,6 +460,8 @@ def test_imagen_generation_response_keeps_predictions_without_image_bytes():
("gemini/imagen-4.0-generate-001", {"predictions": [{"bytesBase64Encoded": "a"}, "not an object"]}),
("gemini-3.1-flash-image-preview", {"usageMetadata": "not an object"}),
("gemini-3.1-flash-image-preview", {"usageMetadata": ["not an object"]}),
+ ("gemini/imagen-4.0-generate-001", ["not an object"]),
+ ("gemini/imagen-4.0-generate-001", "not an object"),
],
)
def test_image_generation_response_rejects_malformed_payloads_without_echoing_them(model: str, payload: object):
@@ -467,3 +469,117 @@ def test_image_generation_response_rejects_malformed_payloads_without_echoing_th
_transform_response(model, payload)
assert "input_value" not in str(exc_info.value)
+
+
+_GEMINI_IMAGE_MODEL = "gemini-3.1-flash-image-preview"
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {},
+ {"candidates": []},
+ {"candidates": ""},
+ {"candidates": {}},
+ {"candidates": [{}]},
+ {"candidates": [{"finishReason": "SAFETY"}]},
+ {"candidates": [{"content": {}}]},
+ {"candidates": [{"content": {"parts": ""}}]},
+ {"candidates": [{"content": {"parts": "plain text"}}]},
+ {"candidates": [{"content": {"parts": {"text": "only text"}}}]},
+ {"candidates": [{"content": {"parts": [{"text": "only text"}, "plain text", ["text"], [], ""]}}]},
+ {"candidates": [{"content": {"parts": [{"inlineData": {"mimeType": "image/png"}}]}}]},
+ {"candidates": [{"content": {"parts": [{"inlineData": "no image"}, {"inlineData": ["mimeType"]}]}}]},
+ {"candidates": [{"content": {"parts": [{"inlineData": ""}, {"inlineData": []}, {"inlineData": {}}]}}]},
+ {"candidates": [{"content": {"parts": [{"inline_data": {"data": "snake-case-is-not-read"}}]}}]},
+ ],
+)
+def test_gemini_image_generation_response_without_inline_image_data_has_no_images(payload: object):
+ assert _transform_response(_GEMINI_IMAGE_MODEL, payload).data == []
+
+
+def test_gemini_image_generation_response_keeps_image_order_and_thought_signatures():
+ result = _transform_response(
+ _GEMINI_IMAGE_MODEL,
+ {
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {"text": "caption"},
+ {"inlineData": {"mimeType": "image/png", "data": "first"}, "thoughtSignature": "sig-1"},
+ "plain text",
+ ]
+ }
+ },
+ {
+ "content": {
+ "parts": [
+ {"inlineData": {"data": "second"}},
+ {"inlineData": {"data": None}, "thoughtSignature": ""},
+ {"inlineData": {"data": "fourth"}, "thoughtSignature": {"nested": ["sig"]}},
+ ]
+ }
+ },
+ ]
+ },
+ )
+
+ assert [(image.b64_json, image.provider_specific_fields) for image in result.data or []] == [
+ ("first", {"thought_signature": "sig-1"}),
+ ("second", None),
+ (None, None),
+ ("fourth", {"thought_signature": {"nested": ["sig"]}}),
+ ]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ ["gemini-secret"],
+ "gemini-secret",
+ {"candidates": None},
+ {"candidates": 7},
+ {"candidates": "gemini-secret"},
+ {"candidates": ["gemini-secret"]},
+ {"candidates": [{"content": None}]},
+ {"candidates": [{"content": "gemini-secret"}]},
+ {"candidates": [{"content": ["gemini-secret"]}]},
+ {"candidates": [{"content": {"parts": None}}]},
+ {"candidates": [{"content": {"parts": [None]}}]},
+ {"candidates": [{"content": {"parts": [7]}}]},
+ {"candidates": [{"content": {"parts": ["inlineData gemini-secret"]}}]},
+ {"candidates": [{"content": {"parts": [["inlineData", "gemini-secret"]]}}]},
+ {"candidates": [{"content": {"parts": {"inlineData": "gemini-secret"}}}]},
+ {"candidates": [{"content": {"parts": [{"inlineData": None}]}}]},
+ {"candidates": [{"content": {"parts": [{"inlineData": 7}]}}]},
+ {"candidates": [{"content": {"parts": [{"inlineData": "data gemini-secret"}]}}]},
+ {"candidates": [{"content": {"parts": [{"inlineData": ["data", "gemini-secret"]}]}}]},
+ ],
+)
+def test_gemini_image_generation_response_rejects_malformed_candidates_without_echoing_them(payload: object):
+ with pytest.raises(ValidationError) as exc_info:
+ _transform_response(_GEMINI_IMAGE_MODEL, payload)
+
+ assert "gemini-secret" not in str(exc_info.value)
+
+
+def _rejection_text(payload: object) -> str:
+ with pytest.raises(ValidationError) as exc_info:
+ _transform_response(_GEMINI_IMAGE_MODEL, payload)
+
+ return str(exc_info.value)
+
+
+@pytest.mark.parametrize("position", [403, 429])
+@pytest.mark.parametrize("malformed_entry", [None, 7, "inlineData"])
+def test_gemini_image_generation_response_rejection_text_is_the_same_wherever_the_malformed_entry_sits(
+ position: int, malformed_entry: object
+):
+ late_parts = [*[{"text": "caption"}] * position, malformed_entry]
+ late_candidates = [*[{"finishReason": "STOP"}] * position, malformed_entry]
+
+ assert _rejection_text({"candidates": [{"content": {"parts": late_parts}}]}) == _rejection_text(
+ {"candidates": [{"content": {"parts": [malformed_entry]}}]}
+ )
+ assert _rejection_text({"candidates": late_candidates}) == _rejection_text({"candidates": [malformed_entry]})
diff --git a/tests/unit/llms/google_pse/__init__.py b/tests/unit/llms/google_pse/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/google_pse/search/__init__.py b/tests/unit/llms/google_pse/search/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/google_pse/search/test_google_pse_search_transformation.py b/tests/unit/llms/google_pse/search/test_google_pse_search_transformation.py
new file mode 100644
index 00000000000..9234e46d5b7
--- /dev/null
+++ b/tests/unit/llms/google_pse/search/test_google_pse_search_transformation.py
@@ -0,0 +1,110 @@
+import datetime
+
+import httpx
+import pytest
+from pydantic import ValidationError
+
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
+from litellm.llms.google_pse.search.transformation import GooglePSESearchConfig
+
+
+def _transform(payload: object) -> SearchResponse:
+ logging_obj = Logging(
+ model="google_pse/search",
+ messages=[],
+ stream=False,
+ call_type="search",
+ start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+ litellm_call_id="call-id",
+ function_id="function-id",
+ )
+ return GooglePSESearchConfig().transform_search_response(
+ raw_response=httpx.Response(200, json=payload), logging_obj=logging_obj
+ )
+
+
+def test_google_pse_items_keep_provider_order_and_ignore_search_metadata():
+ response = _transform(
+ {
+ "kind": "customsearch#search",
+ "url": {"type": "application/json", "template": "https://www.googleapis.com/customsearch/v1?q={q}"},
+ "queries": {"request": [{"title": "Google Custom Search - litellm", "totalResults": "2", "count": 2}]},
+ "searchInformation": {"searchTime": 0.21, "formattedSearchTime": "0.21", "totalResults": "2"},
+ "items": [
+ {
+ "kind": "customsearch#result",
+ "title": "LiteLLM",
+ "htmlTitle": "LiteLLM",
+ "link": "https://example.com/litellm",
+ "displayLink": "example.com",
+ "snippet": "Call every LLM API",
+ "htmlSnippet": "Call every LLM API",
+ "pagemap": {"metatags": [{"og:title": "LiteLLM"}]},
+ },
+ {"title": "Docs", "link": "https://example.com/docs", "snippet": "Docs"},
+ ],
+ }
+ )
+
+ assert response.object == "search"
+ assert response.results == [
+ SearchResult(title="LiteLLM", url="https://example.com/litellm", snippet="Call every LLM API"),
+ SearchResult(title="Docs", url="https://example.com/docs", snippet="Docs"),
+ ]
+
+
+def test_google_pse_item_without_title_link_or_snippet_gets_empty_strings():
+ assert _transform({"items": [{"kind": "customsearch#result"}]}).results == [
+ SearchResult(title="", url="", snippet="", date=None, last_updated=None)
+ ]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {},
+ {"kind": "customsearch#search", "searchInformation": {"totalResults": "0"}},
+ {"items": []},
+ {"items": ""},
+ {"items": {}},
+ ],
+)
+def test_google_pse_response_without_items_is_empty(payload):
+ assert _transform(payload).results == []
+
+
+@pytest.mark.parametrize(
+ ("field", "value"),
+ [
+ ("title", None),
+ ("title", 7),
+ ("link", None),
+ ("link", ["https://example.com"]),
+ ("snippet", None),
+ ("snippet", {"text": "snippet"}),
+ ],
+)
+def test_google_pse_item_with_field_of_wrong_type_is_rejected(field, value):
+ item = {"title": "T", "link": "https://example.com", "snippet": "S", field: value}
+
+ with pytest.raises(ValidationError):
+ _transform({"items": [item]})
+
+
+def test_google_pse_item_with_several_fields_of_wrong_type_reports_every_field():
+ item = {"title": 7, "link": None, "snippet": ["snippet"]}
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"items": [item]})
+
+ assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",)]
+
+
+def test_google_pse_malformed_item_is_rejected_without_naming_its_position():
+ items = [{}] * 429 + ["not-an-object"]
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"items": items})
+
+ assert "429" not in str(exc_info.value)
diff --git a/tests/unit/llms/infinity/embedding/__init__.py b/tests/unit/llms/infinity/embedding/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/infinity/embedding/test_infinity_embedding_transformation.py b/tests/unit/llms/infinity/embedding/test_infinity_embedding_transformation.py
new file mode 100644
index 00000000000..e890360d725
--- /dev/null
+++ b/tests/unit/llms/infinity/embedding/test_infinity_embedding_transformation.py
@@ -0,0 +1,89 @@
+import datetime
+
+import httpx
+import pytest
+from pydantic import ValidationError
+
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.infinity.embedding.transformation import InfinityEmbeddingConfig
+from litellm.types.utils import EmbeddingResponse
+
+EMBEDDING_ROW = {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}
+
+
+def _transform(payload: object) -> EmbeddingResponse:
+ logging_obj = Logging(
+ model="infinity/bge-small",
+ messages=[],
+ stream=False,
+ call_type="embedding",
+ start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+ litellm_call_id="call-id",
+ function_id="function-id",
+ )
+ return InfinityEmbeddingConfig().transform_embedding_response(
+ model="bge-small",
+ raw_response=httpx.Response(200, json=payload),
+ model_response=EmbeddingResponse(),
+ logging_obj=logging_obj,
+ api_key="sk-test",
+ request_data={},
+ optional_params={},
+ litellm_params={},
+ )
+
+
+def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]:
+ assert response.usage is not None
+ return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens)
+
+
+def test_infinity_embedding_response_carries_provider_fields_and_usage():
+ response = _transform(
+ {
+ "object": "list",
+ "model": "bge-small",
+ "data": [EMBEDDING_ROW],
+ "usage": {"prompt_tokens": 8, "total_tokens": 9},
+ }
+ )
+
+ assert response.model == "bge-small"
+ assert response.object == "list"
+ assert response.data == [EMBEDDING_ROW]
+ assert _token_counts(response) == (8, 0, 9)
+
+
+@pytest.mark.parametrize(
+ ("payload", "expected"),
+ [
+ ({"data": [EMBEDDING_ROW]}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 4}}, (4, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 6}}, (0, 0, 6)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": None, "total_tokens": None}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": "", "total_tokens": []}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": "12", "total_tokens": 12.0}}, (12, 0, 12)),
+ (
+ {"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 3, "total_tokens": 5, "completion_tokens": 2}},
+ (3, 0, 5),
+ ),
+ ],
+)
+def test_infinity_embedding_usage_tolerates_missing_null_and_loosely_typed_counts(payload, expected):
+ response = _transform(payload)
+
+ assert response.data == [EMBEDDING_ROW]
+ assert _token_counts(response) == expected
+
+
+@pytest.mark.parametrize("usage", [None, [], "tokens", 3])
+def test_infinity_embedding_usage_that_is_not_an_object_is_rejected(usage):
+ with pytest.raises(ValidationError):
+ _transform({"data": [EMBEDDING_ROW], "usage": usage})
+
+
+@pytest.mark.parametrize("usage", [{"prompt_tokens": "many"}, {"total_tokens": 1.5}, {"prompt_tokens": [1]}])
+def test_infinity_embedding_token_count_that_is_not_an_integer_is_rejected(usage):
+ with pytest.raises(ValidationError):
+ _transform({"data": [EMBEDDING_ROW], "usage": usage})
diff --git a/tests/unit/llms/langgraph/chat/test_langgraph_transformation.py b/tests/unit/llms/langgraph/chat/test_langgraph_transformation.py
index 881da2ce15b..238f7dd2f26 100644
--- a/tests/unit/llms/langgraph/chat/test_langgraph_transformation.py
+++ b/tests/unit/llms/langgraph/chat/test_langgraph_transformation.py
@@ -20,11 +20,14 @@ Non-streaming:
import asyncio
import importlib
+import json
import os
+import httpx
import pytest
import litellm
+from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@@ -232,3 +235,36 @@ _SCALAR_DEFAULTS = {
"api_key": getattr(litellm, "api_key", None),
"cohere_key": getattr(litellm, "cohere_key", None),
}
+
+
+@pytest.mark.asyncio
+async def test_langgraph_streaming_acompletion_without_a_client_posts_through_the_cached_langgraph_client(monkeypatch):
+ sent_requests = []
+
+ async def events():
+ yield b'data: ["messages", [[{"type": "ai", "content": "100"}, {}]]]\n\n'
+ yield b'data: ["metadata", {"run_id": "run-1"}]\n\n'
+
+ def respond(request):
+ sent_requests.append(request)
+ return httpx.Response(200, content=events())
+
+ cached_client = get_async_httpx_client(llm_provider="langgraph", params={})
+ monkeypatch.setattr(cached_client, "client", httpx.AsyncClient(transport=httpx.MockTransport(respond)))
+
+ response = await litellm.acompletion(
+ model="langgraph/agent",
+ messages=[{"role": "user", "content": "What is 25 * 4?"}],
+ api_base="http://langgraph.test",
+ stream=True,
+ )
+ streamed = [(chunk.choices[0].delta.content, chunk.choices[0].finish_reason) async for chunk in response]
+
+ assert streamed == [("100", None), (None, "stop")]
+ assert [(request.method, str(request.url)) for request in sent_requests] == [
+ ("POST", "http://langgraph.test/runs/stream")
+ ]
+ assert json.loads(sent_requests[0].content) == {
+ "assistant_id": "agent",
+ "input": {"messages": [{"role": "human", "content": "What is 25 * 4?"}]},
+ }
diff --git a/tests/unit/llms/lemonade/chat/__init__.py b/tests/unit/llms/lemonade/chat/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/lemonade/chat/test_lemonade_chat_transformation.py b/tests/unit/llms/lemonade/chat/test_lemonade_chat_transformation.py
new file mode 100644
index 00000000000..6f177d9bf8b
--- /dev/null
+++ b/tests/unit/llms/lemonade/chat/test_lemonade_chat_transformation.py
@@ -0,0 +1,172 @@
+import pytest
+import respx
+from pydantic import ValidationError
+
+from litellm.llms.lemonade.chat.transformation import LemonadeChatConfig
+
+LEMONADE_BASE = "http://lemonade.test/api/v1"
+MODEL = "Qwen3-4B-GGUF"
+
+
+def _expected_model_info(**overrides: object) -> dict[str, object]:
+ return {
+ "key": f"lemonade/{MODEL}",
+ "litellm_provider": "lemonade",
+ "mode": "chat",
+ "input_cost_per_token": 0.0,
+ "output_cost_per_token": 0.0,
+ "max_tokens": None,
+ "max_input_tokens": None,
+ "max_output_tokens": None,
+ **overrides,
+ }
+
+
+@pytest.mark.parametrize(
+ ("body", "expected"),
+ [
+ pytest.param(
+ {"object": "list", "data": [{"id": MODEL, "recipe_options": {"ctx_size": [4096]}}, {"id": "nomic-embed"}]},
+ [f"lemonade/{MODEL}", "lemonade/nomic-embed"],
+ id="listed-models",
+ ),
+ pytest.param({"data": []}, [], id="empty-list"),
+ pytest.param({"object": "list"}, [], id="data-key-missing"),
+ pytest.param({"data": ""}, [], id="data-empty-string"),
+ pytest.param({"data": {}}, [], id="data-empty-object"),
+ pytest.param({"data": [{"id": ""}]}, ["lemonade/"], id="empty-id"),
+ ],
+)
+def test_get_models_prefixes_each_listed_id(
+ respx_mock: respx.MockRouter, body: dict[str, object], expected: list[str]
+) -> None:
+ respx_mock.get(f"{LEMONADE_BASE}/models").respond(200, json=body)
+
+ assert LemonadeChatConfig().get_models(api_base=LEMONADE_BASE) == expected
+
+
+def test_get_models_raises_key_error_for_entry_without_id(respx_mock: respx.MockRouter) -> None:
+ respx_mock.get(f"{LEMONADE_BASE}/models").respond(200, json={"data": [{"id": MODEL}, {"name": "no-id"}]})
+
+ with pytest.raises(KeyError, match="id"):
+ LemonadeChatConfig().get_models(api_base=LEMONADE_BASE)
+
+
+@pytest.mark.parametrize(
+ "body",
+ [
+ pytest.param([{"id": MODEL}], id="top-level-list"),
+ pytest.param({"data": None}, id="data-null"),
+ pytest.param({"data": [MODEL]}, id="entry-not-an-object"),
+ pytest.param({"data": [{"id": 7}]}, id="id-not-a-string"),
+ ],
+)
+def test_get_models_rejects_malformed_listing(respx_mock: respx.MockRouter, body: object) -> None:
+ respx_mock.get(f"{LEMONADE_BASE}/models").respond(200, json=body)
+
+ with pytest.raises(ValidationError):
+ LemonadeChatConfig().get_models(api_base=LEMONADE_BASE)
+
+
+@pytest.mark.parametrize(
+ ("body", "expected"),
+ [
+ pytest.param({}, _expected_model_info(), id="empty-object"),
+ pytest.param(
+ {"max_tokens": "2048", "max_output_tokens": 4096, "max_input_tokens": 8192},
+ _expected_model_info(max_tokens=2048, max_output_tokens=4096, max_input_tokens=8192),
+ id="numeric-string-and-ints",
+ ),
+ pytest.param(
+ {"max_output_tokens": 4096},
+ _expected_model_info(max_tokens=4096, max_output_tokens=4096),
+ id="max-tokens-falls-back-to-output",
+ ),
+ pytest.param(
+ {"max_tokens": True, "max_output_tokens": 0, "max_input_tokens": 12.5},
+ _expected_model_info(),
+ id="bool-zero-and-float-ignored",
+ ),
+ pytest.param(
+ {"max_tokens": None, "max_output_tokens": "many", "recipe_options": [1], "provider_specific_entry": "x"},
+ _expected_model_info(provider_specific_entry={"recipe_options": [1]}),
+ id="null-text-and-non-object-options-ignored",
+ ),
+ pytest.param(
+ {"provider_specific_entry": {"recipe_options": {"ctx_size": "32768"}, "labels": ["hot"]}},
+ _expected_model_info(
+ max_input_tokens=32768,
+ provider_specific_entry={"recipe_options": {"ctx_size": "32768"}, "labels": ["hot"]},
+ ),
+ id="nested-entry-kept",
+ ),
+ ],
+)
+def test_get_model_info_reads_token_limits_from_server_payload(
+ respx_mock: respx.MockRouter, body: dict[str, object], expected: dict[str, object]
+) -> None:
+ respx_mock.get(f"{LEMONADE_BASE}/models/{MODEL}").respond(200, json=body)
+
+ assert LemonadeChatConfig().get_model_info(model=f"lemonade/{MODEL}", api_base=LEMONADE_BASE) == expected
+
+
+@pytest.mark.parametrize("body", [pytest.param([{"max_tokens": 5}], id="list"), pytest.param("loaded", id="string")])
+def test_get_model_info_raises_for_non_object_payload_instead_of_defaulting(
+ respx_mock: respx.MockRouter, body: object
+) -> None:
+ respx_mock.get(f"{LEMONADE_BASE}/models/{MODEL}").respond(200, json=body)
+
+ with pytest.raises(ValidationError):
+ LemonadeChatConfig().get_model_info(model=f"lemonade/{MODEL}", api_base=LEMONADE_BASE)
+
+
+def test_get_model_info_defaults_when_server_errors(respx_mock: respx.MockRouter) -> None:
+ respx_mock.get(f"{LEMONADE_BASE}/models/{MODEL}").respond(503, json={"max_tokens": 5})
+
+ assert (
+ LemonadeChatConfig().get_model_info(model=f"lemonade/{MODEL}", api_base=LEMONADE_BASE) == _expected_model_info()
+ )
+
+
+@pytest.mark.parametrize(
+ ("supplied", "stored"),
+ [
+ pytest.param({}, {}, id="nothing-supplied"),
+ pytest.param(
+ {"repeat_penalty": 0.0, "n": 0, "stop": "", "tools": []},
+ {"repeat_penalty": 0.0, "n": 0, "stop": "", "tools": []},
+ id="falsy-values-are-kept",
+ ),
+ pytest.param({"top_k": 40, "top_p": None, "max_tokens": None}, {"top_k": 40}, id="none-is-skipped"),
+ ],
+)
+def test_init_stores_each_supplied_value_except_none_as_class_level_config(
+ supplied: dict[str, object], stored: dict[str, object]
+) -> None:
+ class ScopedLemonadeConfig(LemonadeChatConfig):
+ pass
+
+ ScopedLemonadeConfig(**supplied)
+
+ assert ScopedLemonadeConfig.get_config() == stored
+
+
+def test_init_stores_the_supplied_object_itself_rather_than_a_copy() -> None:
+ class ScopedLemonadeConfig(LemonadeChatConfig):
+ pass
+
+ logit_bias = {"50256": -100}
+
+ ScopedLemonadeConfig(logit_bias=logit_bias)
+
+ assert ScopedLemonadeConfig.get_config()["logit_bias"] is logit_bias
+
+
+def test_init_called_again_with_none_keeps_the_value_stored_earlier() -> None:
+ class ScopedLemonadeConfig(LemonadeChatConfig):
+ pass
+
+ ScopedLemonadeConfig(top_k=40)
+ ScopedLemonadeConfig(top_k=None, n=2)
+
+ assert ScopedLemonadeConfig.get_config() == {"top_k": 40, "n": 2}
diff --git a/tests/unit/llms/linkup/search/test_linkup_search_transformation.py b/tests/unit/llms/linkup/search/test_linkup_search_transformation.py
index c98990b5a55..80967c6e539 100644
--- a/tests/unit/llms/linkup/search/test_linkup_search_transformation.py
+++ b/tests/unit/llms/linkup/search/test_linkup_search_transformation.py
@@ -2,12 +2,18 @@
Tests for Linkup Search API integration.
"""
+import datetime
import os
from unittest.mock import Mock, patch
+import httpx
import pytest
+from pydantic import ValidationError
import litellm
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
+from litellm.llms.linkup.search.transformation import LinkupSearchConfig
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@@ -99,6 +105,100 @@ class TestLinkupSearchTransformation:
assert "Microsoft Cloud revenue" in first_result.snippet
+def _transform(payload: object) -> SearchResponse:
+ logging_obj = Logging(
+ model="linkup/search",
+ messages=[],
+ stream=False,
+ call_type="search",
+ start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+ litellm_call_id="call-id",
+ function_id="function-id",
+ )
+ return LinkupSearchConfig().transform_search_response(
+ raw_response=httpx.Response(200, json=payload), logging_obj=logging_obj
+ )
+
+
+def test_linkup_text_and_image_results_keep_provider_order_and_fields():
+ response = _transform(
+ {
+ "results": [
+ {"type": "text", "name": "Annual Report", "url": "https://example.com/report", "content": "Revenue up"},
+ {"type": "image", "name": "Chart", "url": "https://example.com/chart.png", "content": "A chart"},
+ {"type": "image", "url": "https://example.com/logo.png"},
+ {"name": "Untyped", "url": "https://example.com/untyped", "content": "Defaults to text"},
+ ]
+ }
+ )
+
+ assert response.results == [
+ SearchResult(title="Annual Report", url="https://example.com/report", snippet="Revenue up"),
+ SearchResult(title="Chart", url="https://example.com/chart.png", snippet="A chart"),
+ SearchResult(title="https://example.com/logo.png", url="https://example.com/logo.png", snippet=""),
+ SearchResult(title="Untyped", url="https://example.com/untyped", snippet="Defaults to text"),
+ ]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {},
+ {"answer": "A sourced answer without a results array"},
+ {"results": []},
+ {"results": ""},
+ {"results": {}},
+ {"results": [{"type": "video", "name": "Clip", "url": "https://example.com/clip"}]},
+ ],
+)
+def test_linkup_response_without_text_or_image_results_is_empty(payload):
+ response = _transform(payload)
+
+ assert response.object == "search"
+ assert response.results == []
+
+
+def test_linkup_result_missing_optional_fields_defaults_to_empty_strings():
+ response = _transform(
+ {"results": [{"type": "text", "favicon": "https://example.com/favicon.ico"}, {"type": "image"}]}
+ )
+
+ assert response.results == [
+ SearchResult(title="", url="", snippet=""),
+ SearchResult(title="", url="", snippet=""),
+ ]
+
+
+@pytest.mark.parametrize(
+ "result",
+ [
+ {"type": "text", "name": None, "url": "https://example.com", "content": "c"},
+ {"type": "text", "name": "Title", "url": 7, "content": "c"},
+ {"type": "image", "name": "Title", "url": "https://example.com", "content": None},
+ ],
+)
+def test_linkup_result_with_non_string_field_is_rejected(result):
+ with pytest.raises(ValidationError):
+ _transform({"results": [result]})
+
+
+@pytest.mark.parametrize("result_type", ["text", "image"])
+def test_linkup_result_with_several_non_string_fields_reports_every_field(result_type):
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"results": [{"type": result_type, "name": 7, "url": None, "content": ["content"]}]})
+
+ assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",)]
+
+
+def test_linkup_malformed_result_is_rejected_without_naming_its_position():
+ results = [{}] * 429 + ["not-an-object"]
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"results": results})
+
+ assert "429" not in str(exc_info.value)
+
+
@pytest.fixture(autouse=True)
def _vcr_outcome_gate(request, vcr):
install_live_call_probe(request, vcr)
diff --git a/tests/unit/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py b/tests/unit/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py
index d819a79cef1..87e2c165515 100644
--- a/tests/unit/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py
+++ b/tests/unit/llms/mistral/audio_speech/test_mistral_audio_speech_transformation.py
@@ -1,9 +1,11 @@
import base64
+import json
from typing import Final
from unittest.mock import MagicMock
import httpx
import pytest
+from pydantic import ValidationError
import litellm
from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig
@@ -195,3 +197,113 @@ def test_transform_response_invalid_base64_raises():
raw_response=raw_response,
logging_obj=MagicMock(),
)
+
+
+_SPEECH_AUDIO: Final = b"ID3-fake-mp3-bytes"
+_SPEECH_AUDIO_B64: Final = base64.b64encode(_SPEECH_AUDIO).decode()
+
+
+@pytest.mark.parametrize(
+ ("request_body", "expected_content_type"),
+ [
+ pytest.param(b'{"input": "hi", "response_format": "opus"}', "audio/ogg", id="known-format"),
+ pytest.param(b'{"input": "hi", "response_format": "aac"}', "audio/mpeg", id="unknown-format"),
+ pytest.param(b'{"input": "hi", "response_format": null}', "audio/mpeg", id="null-format"),
+ pytest.param(b'{"input": "hi", "response_format": ["wav"]}', "audio/mpeg", id="non-string-format"),
+ pytest.param(b'{"input": "hi", "ref_audio": "bXktdm9pY2Utc2FtcGxl"}', "audio/mpeg", id="format-omitted"),
+ pytest.param(b"", "audio/mpeg", id="request-without-body"),
+ ],
+)
+def test_transform_response_labels_audio_with_the_requested_format(request_body: bytes, expected_content_type: str):
+ raw_response: Final = httpx.Response(
+ 200,
+ json={"id": None, "audio_data": _SPEECH_AUDIO_B64, "usage": {"characters": 2}, "voices": ["en_paul_neutral"]},
+ request=httpx.Request("POST", SPEECH_URL, content=request_body),
+ )
+ result: Final = MistralTextToSpeechConfig().transform_text_to_speech_response(
+ model="voxtral-mini-tts-2603",
+ raw_response=raw_response,
+ logging_obj=MagicMock(),
+ )
+ assert result.content == _SPEECH_AUDIO
+ assert result.response.headers["content-type"] == expected_content_type
+
+
+@pytest.mark.parametrize(
+ ("payload", "expected_keys"),
+ [
+ pytest.param({}, "()", id="empty-object"),
+ pytest.param({"audio_data": None}, "('audio_data',)", id="null-audio"),
+ pytest.param({"audio_data": ""}, "('audio_data',)", id="empty-audio"),
+ pytest.param({"id": "tts-1", "audio_data": 7}, "('id', 'audio_data')", id="numeric-audio"),
+ pytest.param({"audio_data": [_SPEECH_AUDIO_B64], "id": "tts-1"}, "('audio_data', 'id')", id="list-audio"),
+ ],
+)
+def test_transform_response_without_usable_audio_reports_the_response_keys(
+ payload: dict[str, object], expected_keys: str
+):
+ raw_response: Final = httpx.Response(200, json=payload, request=httpx.Request("POST", SPEECH_URL))
+ with pytest.raises(MistralTextToSpeechException) as exc_info:
+ MistralTextToSpeechConfig().transform_text_to_speech_response(
+ model="voxtral-mini-tts-2603",
+ raw_response=raw_response,
+ logging_obj=MagicMock(),
+ )
+ assert exc_info.value.status_code == 500
+ assert exc_info.value.message == f"No audio_data in Mistral speech response. Response keys: {expected_keys}"
+
+
+def test_transform_response_non_json_body_raises_provider_error_with_upstream_status():
+ raw_response: Final = httpx.Response(
+ 502,
+ text="bad gateway",
+ request=httpx.Request("POST", SPEECH_URL),
+ )
+ with pytest.raises(MistralTextToSpeechException) as exc_info:
+ MistralTextToSpeechConfig().transform_text_to_speech_response(
+ model="voxtral-mini-tts-2603",
+ raw_response=raw_response,
+ logging_obj=MagicMock(),
+ )
+ assert exc_info.value.status_code == 502
+ assert exc_info.value.message == "Non-JSON response from Mistral speech API: bad gateway"
+
+
+@pytest.mark.parametrize(
+ "response_body",
+ [b'["secret-audio-payload"]', b'"secret-audio-payload"', b"7", b"null", b"true"],
+)
+def test_transform_response_json_body_that_is_not_an_object_is_rejected_without_echoing_it(response_body: bytes):
+ raw_response: Final = httpx.Response(200, content=response_body, request=httpx.Request("POST", SPEECH_URL))
+ with pytest.raises(ValidationError) as exc_info:
+ MistralTextToSpeechConfig().transform_text_to_speech_response(
+ model="voxtral-mini-tts-2603",
+ raw_response=raw_response,
+ logging_obj=MagicMock(),
+ )
+ assert "secret-audio-payload" not in str(exc_info.value)
+ assert "input_value" not in str(exc_info.value)
+
+
+@pytest.mark.parametrize(
+ ("request_body", "expected_error"),
+ [
+ pytest.param(b'["wav"]', ValidationError, id="json-array-body"),
+ pytest.param(b'"wav"', ValidationError, id="json-string-body"),
+ pytest.param(b"response_format=wav", json.JSONDecodeError, id="form-encoded-body"),
+ ],
+)
+def test_transform_response_rejects_a_request_body_that_is_not_a_json_object(
+ request_body: bytes, expected_error: type[Exception]
+):
+ raw_response: Final = httpx.Response(
+ 200,
+ json={"audio_data": _SPEECH_AUDIO_B64},
+ request=httpx.Request("POST", SPEECH_URL, content=request_body),
+ )
+ with pytest.raises(expected_error):
+ MistralTextToSpeechConfig().transform_text_to_speech_response(
+ model="voxtral-mini-tts-2603",
+ raw_response=raw_response,
+ logging_obj=MagicMock(),
+ )
diff --git a/tests/unit/llms/nlp_cloud/__init__.py b/tests/unit/llms/nlp_cloud/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/nlp_cloud/chat/__init__.py b/tests/unit/llms/nlp_cloud/chat/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/nlp_cloud/chat/test_nlp_cloud_chat_transformation.py b/tests/unit/llms/nlp_cloud/chat/test_nlp_cloud_chat_transformation.py
new file mode 100644
index 00000000000..35f8d933cee
--- /dev/null
+++ b/tests/unit/llms/nlp_cloud/chat/test_nlp_cloud_chat_transformation.py
@@ -0,0 +1,47 @@
+import pytest
+
+from litellm.llms.nlp_cloud.chat.transformation import NLPCloudConfig
+
+
+@pytest.mark.parametrize(
+ ("supplied", "stored"),
+ [
+ pytest.param({}, {}, id="nothing-supplied"),
+ pytest.param(
+ {"max_length": 0, "length_no_input": False, "end_sequence": "", "temperature": 0.0},
+ {"max_length": 0, "length_no_input": False, "end_sequence": "", "temperature": 0.0},
+ id="falsy-values-are-kept",
+ ),
+ pytest.param({"top_k": 40, "top_p": None, "num_beams": None}, {"top_k": 40}, id="none-is-skipped"),
+ ],
+)
+def test_init_stores_each_supplied_value_except_none_as_class_level_config(
+ supplied: dict[str, object], stored: dict[str, object]
+) -> None:
+ class ScopedNLPCloudConfig(NLPCloudConfig):
+ pass
+
+ ScopedNLPCloudConfig(**supplied)
+
+ assert ScopedNLPCloudConfig.get_config() == stored
+
+
+def test_init_stores_the_supplied_object_itself_rather_than_a_copy() -> None:
+ class ScopedNLPCloudConfig(NLPCloudConfig):
+ pass
+
+ bad_words = ["darn", "heck"]
+
+ ScopedNLPCloudConfig(bad_words=bad_words)
+
+ assert ScopedNLPCloudConfig.get_config()["bad_words"] is bad_words
+
+
+def test_init_called_again_with_none_keeps_the_value_stored_earlier() -> None:
+ class ScopedNLPCloudConfig(NLPCloudConfig):
+ pass
+
+ ScopedNLPCloudConfig(top_k=40)
+ ScopedNLPCloudConfig(top_k=None, top_p=0.5)
+
+ assert ScopedNLPCloudConfig.get_config() == {"top_k": 40, "top_p": 0.5}
diff --git a/tests/unit/llms/ollama/test_ollama_completion_transformation.py b/tests/unit/llms/ollama/test_ollama_completion_transformation.py
index 7be377247af..7e6dafe030c 100644
--- a/tests/unit/llms/ollama/test_ollama_completion_transformation.py
+++ b/tests/unit/llms/ollama/test_ollama_completion_transformation.py
@@ -889,6 +889,59 @@ def test_transform_request_leaves_unreadable_images_untouched(payload: str) -> N
assert data["images"] == [payload]
+async def _complete_in_json_mode(response_text: str) -> ModelResponse:
+ def respond(request: httpx.Request) -> httpx.Response:
+ return httpx.Response(200, json={"model": "llama3", "response": response_text, "done": True})
+
+ return await litellm.acompletion(
+ model="ollama/llama3",
+ messages=[{"role": "user", "content": "hi"}],
+ api_base="http://ollama.example:11434",
+ format="json",
+ client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)),
+ )
+
+
+@pytest.mark.parametrize(
+ ("response_text", "expected_content"),
+ [
+ ('[{"name": "f", "arguments": {}}]', '[{"name": "f", "arguments": {}}]'),
+ ('["name", "arguments"]', '["name", "arguments"]'),
+ ('"name arguments"', '"name arguments"'),
+ ('{"name": "f"}', '{"name": "f"}'),
+ ('{"b": 1,"a": 2}', '{"b": 1, "a": 2}'),
+ ("1.50", "1.5"),
+ ("true", "true"),
+ ("null", "null"),
+ ],
+ ids=["list_of_calls", "list_of_key_names", "text_with_key_names", "name_only", "object", "number", "bool", "null"],
+)
+async def test_ollama_json_mode_completion_reserializes_json_that_is_not_a_function_call(
+ response_text: str, expected_content: str
+) -> None:
+ response = await _complete_in_json_mode(response_text)
+
+ assert response.choices[0].message.content == expected_content
+ assert response.choices[0].message.tool_calls is None
+ assert response.choices[0].finish_reason == "stop"
+
+
+@pytest.mark.parametrize(
+ ("arguments_json", "expected_arguments"),
+ [('{"city":"Paris"}', '{"city": "Paris"}'), ("[1,2]", "[1, 2]"), ('"{}"', '"{}"'), ("null", "null")],
+ ids=["object", "list", "text", "null"],
+)
+async def test_ollama_json_mode_completion_turns_a_name_and_arguments_object_into_a_tool_call(
+ arguments_json: str, expected_arguments: str
+) -> None:
+ response = await _complete_in_json_mode(f'{{"name": "get_weather", "arguments": {arguments_json}}}')
+
+ (tool_call,) = response.choices[0].message.tool_calls
+ assert (tool_call.function.name, tool_call.function.arguments) == ("get_weather", expected_arguments)
+ assert response.choices[0].message.content is None
+ assert response.choices[0].finish_reason == "tool_calls"
+
+
def test_ollama_vision_model_sends_the_image_beside_the_user_prompt(respx_mock: respx.MockRouter) -> None:
image_base64: Final = _image_base64("PNG")
route: Final = respx_mock.post("http://ollama.test:11434/api/generate").respond(
diff --git a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py
new file mode 100644
index 00000000000..2e0fc55cda4
--- /dev/null
+++ b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py
@@ -0,0 +1,187 @@
+import json
+import operator
+from typing import Final
+
+import httpx
+import pytest
+import respx
+from pydantic import ValidationError
+
+import litellm
+from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter
+from litellm.types.utils import TokenCountResponse
+
+COUNT_URL: Final = "https://api.openai.com/v1/responses/input_tokens"
+
+
+async def _count(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> TokenCountResponse | None:
+ monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+ with respx.mock:
+ respx.post(COUNT_URL).mock(return_value=response)
+ return await OpenAITokenCounter().count_tokens(
+ model_to_use="gpt-4o",
+ messages=[{"role": "user", "content": "hi"}],
+ contents=None,
+ deployment={"litellm_params": {"api_key": "sk-test"}},
+ request_model="openai/gpt-4o",
+ )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ ("payload", "expected_total"),
+ [
+ ({"input_tokens": 15}, 15),
+ ({"input_tokens": "15"}, 15),
+ ({"input_tokens": 15.0}, 15),
+ ({"input_tokens": True}, 1),
+ ({"input_tokens": -3}, -3),
+ ({}, 0),
+ ({"object": "response.input_tokens", "error": {"message": "429 rate limit"}}, 0),
+ ({"input_tokens": 7, "object": "response.input_tokens", "nested": {"a": [1, None]}}, 7),
+ ],
+)
+async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload(
+ monkeypatch: pytest.MonkeyPatch, payload: dict[str, object], expected_total: int
+) -> None:
+ result: Final = await _count(monkeypatch, httpx.Response(200, json=payload))
+
+ assert result is not None
+ assert repr(result.model_dump()) == repr(
+ {
+ "total_tokens": expected_total,
+ "request_model": "openai/gpt-4o",
+ "model_used": "gpt-4o",
+ "tokenizer_type": "openai_api",
+ "original_response": payload,
+ "error": False,
+ "error_message": None,
+ "status_code": None,
+ }
+ )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "rejected_count",
+ [
+ None,
+ 15.5,
+ "Request timed out",
+ "429 Too Many Requests",
+ [15],
+ [*range(403), "Request timed out"],
+ {"value": 429},
+ ],
+)
+async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_validation_message(
+ monkeypatch: pytest.MonkeyPatch, rejected_count: object
+) -> None:
+ with pytest.raises(ValidationError) as rejection:
+ TokenCountResponse.model_validate(
+ {
+ "total_tokens": rejected_count,
+ "request_model": "openai/gpt-4o",
+ "model_used": "gpt-4o",
+ "tokenizer_type": "openai_api",
+ }
+ )
+
+ result: Final = await _count(monkeypatch, httpx.Response(200, json={"input_tokens": rejected_count}))
+
+ assert result is not None
+ assert [error["loc"] for error in rejection.value.errors()] == [("total_tokens",)]
+ assert repr(result.model_dump()) == repr(
+ {
+ "total_tokens": 0,
+ "request_model": "openai/gpt-4o",
+ "model_used": "gpt-4o",
+ "tokenizer_type": "openai_api",
+ "original_response": None,
+ "error": True,
+ "error_message": str(rejection.value),
+ "status_code": 500,
+ }
+ )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "payload",
+ [[*({"input_tokens": position} for position in range(403)), "429"], "input_tokens", 15, 0, False],
+)
+async def test_count_payload_that_is_not_an_object_becomes_an_error_response_naming_the_failed_lookup(
+ monkeypatch: pytest.MonkeyPatch, payload: object
+) -> None:
+ with pytest.raises(AttributeError) as failed_lookup:
+ operator.attrgetter("get")(payload)
+
+ result: Final = await _count(monkeypatch, httpx.Response(200, json=payload))
+
+ assert result is not None
+ assert repr(result.model_dump()) == repr(
+ {
+ "total_tokens": 0,
+ "request_model": "openai/gpt-4o",
+ "model_used": "gpt-4o",
+ "tokenizer_type": "openai_api",
+ "original_response": None,
+ "error": True,
+ "error_message": str(failed_lookup.value),
+ "status_code": 500,
+ }
+ )
+
+
+@pytest.mark.asyncio
+async def test_null_count_payload_falls_back_to_local_counting(monkeypatch: pytest.MonkeyPatch) -> None:
+ assert await _count(monkeypatch, httpx.Response(200, content=b"null")) is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ ("upstream_status", "upstream_text"),
+ [(429, "rate limit exceeded"), (403, "forbidden"), (408, "Request timed out")],
+)
+async def test_upstream_error_status_becomes_an_error_response_with_that_status_and_text(
+ monkeypatch: pytest.MonkeyPatch, upstream_status: int, upstream_text: str
+) -> None:
+ result: Final = await _count(monkeypatch, httpx.Response(upstream_status, text=upstream_text))
+
+ assert result is not None
+ assert repr(result.model_dump()) == repr(
+ {
+ "total_tokens": 0,
+ "request_model": "openai/gpt-4o",
+ "model_used": "gpt-4o",
+ "tokenizer_type": "openai_api",
+ "original_response": None,
+ "error": True,
+ "error_message": upstream_text,
+ "status_code": upstream_status,
+ }
+ )
+
+
+@pytest.mark.asyncio
+async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ with pytest.raises(json.JSONDecodeError) as undecodable:
+ json.loads("not json")
+
+ result: Final = await _count(monkeypatch, httpx.Response(200, text="not json"))
+
+ assert result is not None
+ assert repr(result.model_dump()) == repr(
+ {
+ "total_tokens": 0,
+ "request_model": "openai/gpt-4o",
+ "model_used": "gpt-4o",
+ "tokenizer_type": "openai_api",
+ "original_response": None,
+ "error": True,
+ "error_message": f"CountTokens processing error: {undecodable.value}",
+ "status_code": 500,
+ }
+ )
diff --git a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py
index 32b47c1c680..59355dfabcb 100644
--- a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py
+++ b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py
@@ -7,8 +7,7 @@ import pytest
import respx
import litellm
-
-
+from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
from litellm.llms.openrouter.chat.transformation import (
OpenRouterChatCompletionStreamingHandler,
@@ -628,6 +627,88 @@ def test_openrouter_reasoning_effort_high_passes_through():
assert result["reasoning_effort"] == "high"
+OPENROUTER_COST_HEADER = "llm_provider-x-litellm-response-cost"
+
+
+def _complete_with_openrouter_usage(usage):
+ def respond(request: httpx.Request) -> httpx.Response:
+ return httpx.Response(
+ 200,
+ json={
+ "id": "gen-1",
+ "object": "chat.completion",
+ "created": 1700000000,
+ "model": "openai/gpt-4o",
+ "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}],
+ "usage": usage,
+ },
+ )
+
+ return litellm.completion(
+ model="openrouter/openai/gpt-4o",
+ messages=[{"role": "user", "content": "hi"}],
+ api_key="sk-test",
+ client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))),
+ )
+
+
+@pytest.mark.parametrize(
+ ("cost", "expected"),
+ [
+ pytest.param(0.00015, 0.00015, id="float"),
+ pytest.param(3, 3.0, id="int"),
+ pytest.param(0, 0.0, id="zero"),
+ pytest.param("0.5", 0.5, id="numeric-string"),
+ pytest.param(" 1.5 ", 1.5, id="padded-numeric-string"),
+ pytest.param(True, 1.0, id="bool"),
+ ],
+)
+def test_openrouter_completion_stores_usage_cost_as_float_cost_header(cost, expected):
+ response = _complete_with_openrouter_usage(
+ {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30, "cost": cost, "is_byok": False}
+ )
+
+ stored = response._hidden_params["additional_headers"][OPENROUTER_COST_HEADER]
+ assert type(stored) is float
+ assert stored == expected
+ assert response._hidden_params["response_cost"] == expected
+
+
+@pytest.mark.parametrize(
+ "usage",
+ [
+ pytest.param({"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, id="cost-missing"),
+ pytest.param({"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30, "cost": None}, id="cost-null"),
+ pytest.param({"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30, "cost": "abc"}, id="cost-text"),
+ pytest.param({"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30, "cost": [1]}, id="cost-list"),
+ pytest.param(
+ {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30, "cost": {"upstream": 1.0}},
+ id="cost-object",
+ ),
+ pytest.param({}, id="usage-empty"),
+ pytest.param(None, id="usage-null"),
+ ],
+)
+def test_openrouter_completion_without_usable_cost_returns_response_without_cost_header(usage):
+ response = _complete_with_openrouter_usage(usage)
+
+ assert response.choices[0].message.content == "hi"
+ assert OPENROUTER_COST_HEADER not in response._hidden_params["additional_headers"]
+
+
+@pytest.mark.parametrize(
+ "usage",
+ [
+ pytest.param([{"cost": 1.0}], id="usage-list"),
+ pytest.param("cost", id="usage-text"),
+ pytest.param(5, id="usage-number"),
+ ],
+)
+def test_openrouter_completion_with_non_object_usage_raises_api_connection_error(usage):
+ with pytest.raises(litellm.APIConnectionError):
+ _complete_with_openrouter_usage(usage)
+
+
@respx.mock
def test_openrouter_streaming_completion_assembles_text_and_finish_reason() -> None:
chunks: Final = (
diff --git a/tests/unit/llms/opensandbox/__init__.py b/tests/unit/llms/opensandbox/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/opensandbox/sandbox/__init__.py b/tests/unit/llms/opensandbox/sandbox/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/opensandbox/sandbox/test_transformation.py b/tests/unit/llms/opensandbox/sandbox/test_transformation.py
new file mode 100644
index 00000000000..ec2a47cb92f
--- /dev/null
+++ b/tests/unit/llms/opensandbox/sandbox/test_transformation.py
@@ -0,0 +1,153 @@
+import json
+from collections.abc import Sequence
+
+import httpx
+import pytest
+from pydantic import ValidationError
+
+import litellm
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+from litellm.llms.opensandbox.sandbox.transformation import OpenSandboxSandboxConfig
+
+API_BASE = "https://sandbox.test/v1"
+PENDING = {"id": "osb_1", "status": {"state": "Pending"}}
+RUNNING = {"id": "osb_1", "status": {"state": "Running"}}
+EXECD_ENDPOINT = {"endpoint": "execd.local:44772", "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}}
+
+
+def _sandbox_api(created: object, polled: Sequence[object] = ()) -> tuple[AsyncHTTPHandler, list[str]]:
+ remaining = list(polled)
+ requested: list[str] = []
+
+ def respond(request: httpx.Request) -> httpx.Response:
+ requested.append(f"{request.method} {request.url.path}")
+ if request.method == "POST":
+ return httpx.Response(200, content=json.dumps(created).encode())
+ if "/endpoints/" in request.url.path:
+ return httpx.Response(200, content=json.dumps(EXECD_ENDPOINT).encode())
+ return httpx.Response(200, content=json.dumps(remaining.pop(0)).encode())
+
+ return AsyncHTTPHandler(transport=httpx.MockTransport(respond)), requested
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ ("created", "expected_id"),
+ [
+ pytest.param(
+ {**RUNNING, "createdAt": "2026-01-01T00:00:00Z", "metadata": {"team": ["a"]}}, "osb_1", id="text-id"
+ ),
+ pytest.param({"id": 123, "status": {"state": "Running"}}, "123", id="numeric-id"),
+ ],
+)
+async def test_acreate_sandbox_builds_the_handle_from_a_running_create_response(created: object, expected_id: str):
+ client, requested = _sandbox_api(created)
+
+ handle = await OpenSandboxSandboxConfig().acreate_sandbox(api_key="osb_key", api_base=API_BASE, client=client)
+
+ assert (handle.id, handle.provider, handle.domain) == (expected_id, "opensandbox", API_BASE)
+ assert type(handle.id) is str
+ assert handle._hidden_params == {
+ "api_base": API_BASE,
+ "api_key": "osb_key",
+ "execd_endpoint": "execd.local:44772",
+ "execd_headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"},
+ "use_server_proxy": False,
+ }
+ assert requested == ["POST /v1/sandboxes", f"GET /v1/sandboxes/{expected_id}/endpoints/44772"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "created",
+ [
+ pytest.param(PENDING, id="pending"),
+ pytest.param({"id": "osb_1"}, id="status-missing"),
+ pytest.param({"id": "osb_1", "status": None}, id="status-null"),
+ pytest.param({"id": "osb_1", "status": "Running"}, id="status-text"),
+ pytest.param({"id": "osb_1", "status": {"state": None}}, id="state-null"),
+ ],
+)
+async def test_acreate_sandbox_polls_when_the_create_response_does_not_report_running(created: object):
+ client, requested = _sandbox_api(created, polled=[RUNNING])
+
+ handle = await OpenSandboxSandboxConfig().acreate_sandbox(api_key="osb_key", api_base=API_BASE, client=client)
+
+ assert handle.id == "osb_1"
+ assert requested == ["POST /v1/sandboxes", "GET /v1/sandboxes/osb_1", "GET /v1/sandboxes/osb_1/endpoints/44772"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "not_ready",
+ [
+ pytest.param(PENDING, id="pending"),
+ pytest.param([], id="list"),
+ pytest.param("Running", id="text"),
+ pytest.param(None, id="null"),
+ pytest.param(7, id="number"),
+ pytest.param({}, id="empty-object"),
+ pytest.param({"status": ["Running"]}, id="status-list"),
+ pytest.param({"status": {"state": 5}}, id="state-number"),
+ ],
+)
+async def test_acreate_sandbox_keeps_polling_past_a_status_response_without_a_running_state(not_ready: object):
+ client, requested = _sandbox_api(PENDING, polled=[not_ready, RUNNING])
+
+ handle = await OpenSandboxSandboxConfig().acreate_sandbox(
+ api_key="osb_key", api_base=API_BASE, poll_interval=0, client=client
+ )
+
+ assert handle.id == "osb_1"
+ assert requested == [
+ "POST /v1/sandboxes",
+ "GET /v1/sandboxes/osb_1",
+ "GET /v1/sandboxes/osb_1",
+ "GET /v1/sandboxes/osb_1/endpoints/44772",
+ ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("state", ["Failed", "Stopping", "Terminated"])
+async def test_acreate_sandbox_raises_when_the_polled_sandbox_reaches_a_terminal_state(state: str):
+ client, requested = _sandbox_api(PENDING, polled=[[], {"status": {"state": state}}])
+
+ with pytest.raises(ValueError, match=f"OpenSandbox sandbox osb_1 entered {state}"):
+ await OpenSandboxSandboxConfig().acreate_sandbox(
+ api_key="osb_key", api_base=API_BASE, poll_interval=0, client=client
+ )
+
+ assert requested == ["POST /v1/sandboxes", "GET /v1/sandboxes/osb_1", "GET /v1/sandboxes/osb_1"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("created", [["secret-token"], "secret-token", 7, None])
+async def test_acreate_sandbox_rejects_a_create_response_that_is_not_an_object_without_echoing_it(created: object):
+ client, requested = _sandbox_api(created)
+
+ with pytest.raises(ValidationError) as exc_info:
+ await OpenSandboxSandboxConfig().acreate_sandbox(api_key="osb_key", api_base=API_BASE, client=client)
+
+ assert [error["type"] for error in exc_info.value.errors()] == ["dict_type"]
+ assert "secret-token" not in str(exc_info.value)
+ assert requested == ["POST /v1/sandboxes"]
+
+
+@pytest.mark.asyncio
+async def test_public_acreate_sandbox_raises_the_validation_error_for_a_create_response_that_is_not_an_object():
+ client, requested = _sandbox_api([RUNNING])
+
+ with pytest.raises(ValidationError):
+ await litellm.acreate_sandbox(provider="opensandbox", api_key="osb_key", api_base=API_BASE, client=client)
+
+ assert requested == ["POST /v1/sandboxes"]
+
+
+@pytest.mark.asyncio
+async def test_acreate_sandbox_requires_a_sandbox_id():
+ client, requested = _sandbox_api({"status": {"state": "Running"}})
+
+ with pytest.raises(KeyError, match="id"):
+ await OpenSandboxSandboxConfig().acreate_sandbox(api_key="osb_key", api_base=API_BASE, client=client)
+
+ assert requested == ["POST /v1/sandboxes"]
diff --git a/tests/unit/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py b/tests/unit/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py
index ad2a585b536..7dcd910113c 100644
--- a/tests/unit/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py
+++ b/tests/unit/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py
@@ -1,6 +1,14 @@
+import datetime
from unittest.mock import patch
+import httpx
+import pytest
+from pydantic import ValidationError
+
import litellm
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.ovhcloud.embedding.transformation import OVHCloudEmbeddingConfig
+from litellm.types.utils import EmbeddingResponse
model = "ovhcloud/BGE-M3"
@@ -36,3 +44,84 @@ def test_ovhcloud_embeddings():
assert isinstance(response.data[0]["embedding"], list)
assert response.model == model
assert response.object == "embedding"
+
+
+EMBEDDING_ROW = {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}
+
+
+def _transform(payload: object) -> EmbeddingResponse:
+ logging_obj = Logging(
+ model="ovhcloud/BGE-M3",
+ messages=[],
+ stream=False,
+ call_type="embedding",
+ start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+ litellm_call_id="call-id",
+ function_id="function-id",
+ )
+ return OVHCloudEmbeddingConfig().transform_embedding_response(
+ model="BGE-M3",
+ raw_response=httpx.Response(200, json=payload),
+ model_response=EmbeddingResponse(),
+ logging_obj=logging_obj,
+ api_key="sk-test",
+ request_data={},
+ optional_params={},
+ litellm_params={},
+ )
+
+
+def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]:
+ assert response.usage is not None
+ return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens)
+
+
+def test_ovhcloud_embedding_response_carries_provider_fields_and_usage():
+ response = _transform(
+ {
+ "object": "list",
+ "model": "BGE-M3",
+ "data": [EMBEDDING_ROW],
+ "usage": {"prompt_tokens": 8, "total_tokens": 9},
+ }
+ )
+
+ assert response.model == "BGE-M3"
+ assert response.object == "list"
+ assert response.data == [EMBEDDING_ROW]
+ assert _token_counts(response) == (8, 0, 9)
+
+
+@pytest.mark.parametrize(
+ ("payload", "expected"),
+ [
+ ({"data": [EMBEDDING_ROW]}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 4}}, (4, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 6}}, (0, 0, 6)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": None, "total_tokens": None}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": "", "total_tokens": []}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": "12", "total_tokens": 12.0}}, (12, 0, 12)),
+ (
+ {"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 3, "total_tokens": 5, "completion_tokens": 2}},
+ (3, 0, 5),
+ ),
+ ],
+)
+def test_ovhcloud_embedding_usage_tolerates_missing_null_and_loosely_typed_counts(payload, expected):
+ response = _transform(payload)
+
+ assert response.data == [EMBEDDING_ROW]
+ assert _token_counts(response) == expected
+
+
+@pytest.mark.parametrize("usage", [None, [], "tokens", 3])
+def test_ovhcloud_embedding_usage_that_is_not_an_object_is_rejected(usage):
+ with pytest.raises(ValidationError):
+ _transform({"data": [EMBEDDING_ROW], "usage": usage})
+
+
+@pytest.mark.parametrize("usage", [{"prompt_tokens": "many"}, {"total_tokens": 1.5}, {"prompt_tokens": [1]}])
+def test_ovhcloud_embedding_token_count_that_is_not_an_integer_is_rejected(usage):
+ with pytest.raises(ValidationError):
+ _transform({"data": [EMBEDDING_ROW], "usage": usage})
diff --git a/tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py b/tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py
index 29d185b686d..2c95265890d 100644
--- a/tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py
+++ b/tests/unit/llms/perplexity/chat/test_perplexity_chat_transformation.py
@@ -5,13 +5,16 @@ Tests the response transformation to extract citation tokens and search queries
from Perplexity API responses.
"""
+import datetime
from unittest.mock import Mock
+import httpx
import pytest
# Add the project root to Python path
from litellm import ModelResponse
+from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.perplexity.chat.transformation import PerplexityChatConfig
from litellm.types.utils import Usage
@@ -750,3 +753,167 @@ class TestPerplexityChatTransformation:
not hasattr(choice.message, "annotations")
or choice.message.annotations is None
)
+
+
+_PERPLEXITY_COMPLETION = {
+ "id": "cmpl-1",
+ "object": "chat.completion",
+ "created": 1700000000,
+ "model": "sonar",
+ "choices": [
+ {
+ "index": 0,
+ "finish_reason": "stop",
+ "message": {"role": "assistant", "content": "Paris [1] is in France [2][7]."},
+ }
+ ],
+ "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
+}
+_FIRST_CITATION = {
+ "type": "url_citation",
+ "url_citation": {"url": "https://a.example", "title": "", "start_index": 6, "end_index": 9},
+}
+
+
+def _transform_perplexity_response(body: object) -> ModelResponse:
+ logging_obj = Logging(
+ model="sonar",
+ messages=[{"role": "user", "content": "hi"}],
+ stream=False,
+ call_type="completion",
+ start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+ litellm_call_id="perplexity-transform-test",
+ function_id="perplexity-transform-test",
+ )
+ return PerplexityChatConfig().transform_response(
+ model="sonar",
+ raw_response=httpx.Response(200, json=body),
+ model_response=ModelResponse(),
+ logging_obj=logging_obj,
+ request_data={"model": "sonar"},
+ messages=[{"role": "user", "content": "hi"}],
+ optional_params={},
+ litellm_params={},
+ encoding=None,
+ )
+
+
+def test_transform_response_adds_citation_annotations_and_search_usage_from_the_raw_body():
+ response = _transform_perplexity_response(
+ {
+ **_PERPLEXITY_COMPLETION,
+ "citations": ["https://a.example", "https://b.example"],
+ "search_results": [
+ {"url": "https://a.example", "title": "A"},
+ {"url": "https://c.example", "title": "C"},
+ ],
+ "usage": {**_PERPLEXITY_COMPLETION["usage"], "num_search_queries": 3},
+ }
+ )
+
+ assert response.choices[0].message.annotations == [
+ {
+ "type": "url_citation",
+ "url_citation": {"url": "https://a.example", "title": "A", "start_index": 6, "end_index": 9},
+ },
+ {
+ "type": "url_citation",
+ "url_citation": {"url": "https://b.example", "title": "", "start_index": 23, "end_index": 26},
+ },
+ ]
+ assert repr(response.usage.prompt_tokens_details.web_search_requests) == "3"
+ assert repr(response.usage.citation_tokens) == "8"
+ assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (10, 5, 15)
+ assert response.id == "cmpl-1"
+
+
+@pytest.mark.parametrize(
+ "extra",
+ [{}, {"citations": None}, {"citations": []}, {"search_results": []}, {"num_search_queries": 0}],
+)
+def test_transform_response_leaves_a_completion_without_perplexity_fields_unchanged(extra: dict):
+ response = _transform_perplexity_response({**_PERPLEXITY_COMPLETION, **extra})
+
+ assert response.choices[0].message.content == "Paris [1] is in France [2][7]."
+ assert not hasattr(response.choices[0].message, "annotations")
+ assert not hasattr(response.usage, "citation_tokens")
+ assert response.usage.prompt_tokens_details is None
+ assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (10, 5, 15)
+
+
+def test_transform_response_reads_search_queries_from_the_root_of_the_body():
+ response = _transform_perplexity_response({**_PERPLEXITY_COMPLETION, "num_search_queries": 2})
+
+ assert repr(response.usage.prompt_tokens_details.web_search_requests) == "2"
+ assert not hasattr(response.usage, "citation_tokens")
+
+
+def test_transform_response_annotates_and_counts_citations_when_the_provider_sends_no_usage():
+ body = {key: value for key, value in _PERPLEXITY_COMPLETION.items() if key != "usage"}
+
+ response = _transform_perplexity_response({**body, "citations": ["https://a.example"]})
+
+ assert response.choices[0].message.annotations == [_FIRST_CITATION]
+ assert repr(response.usage.citation_tokens) == "4"
+ assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (0, 0, 0)
+
+
+@pytest.mark.parametrize(
+ ("extra", "citation_tokens_recorded"),
+ [
+ ({"usage": None}, False),
+ ({"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "num_search_queries": "3"}}, True),
+ ({"search_results": 3}, True),
+ ({"search_results": None}, True),
+ ],
+)
+def test_transform_response_still_returns_the_completion_when_a_perplexity_field_is_malformed(
+ extra: dict, citation_tokens_recorded: bool
+):
+ response = _transform_perplexity_response({**_PERPLEXITY_COMPLETION, "citations": ["https://a.example"], **extra})
+
+ assert response.choices[0].message.content == "Paris [1] is in France [2][7]."
+ assert not hasattr(response.choices[0].message, "annotations")
+ assert hasattr(response.usage, "citation_tokens") is citation_tokens_recorded
+
+
+@pytest.mark.parametrize("body", [[], [_PERPLEXITY_COMPLETION], "choices", ""])
+def test_transform_response_rejects_a_body_that_is_not_an_object_as_an_invalid_response_object(body: object):
+ with pytest.raises(Exception, match="Invalid response object") as exc_info:
+ _transform_perplexity_response(body)
+
+ assert type(exc_info.value) is Exception
+
+
+@pytest.mark.parametrize("body", [3, 1.5, True, ["error"], "error"])
+def test_transform_response_raises_type_error_for_a_scalar_or_error_bearing_non_object_body(body: object):
+ with pytest.raises(TypeError):
+ _transform_perplexity_response(body)
+
+
+def test_transform_response_keeps_body_values_of_every_json_type_as_the_provider_sent_them():
+ citations = ["https://a.example", 5, None, 1.5, True, {"url": "https://d.example"}, ["https://e.example"]]
+ search_results = [{"url": "https://a.example", "title": "A", "score": 0.5, "tags": None}, "loose", 7]
+
+ response = _transform_perplexity_response(
+ {
+ **_PERPLEXITY_COMPLETION,
+ "citations": citations,
+ "search_results": search_results,
+ "related_score": 1.5,
+ "is_final": True,
+ "trace": {"hops": [1, {"next": None}]},
+ }
+ )
+
+ assert response.choices[0].message.annotations == [
+ {
+ "type": "url_citation",
+ "url_citation": {"url": "https://a.example", "title": "A", "start_index": 6, "end_index": 9},
+ }
+ ]
+ assert response.citations == citations
+ assert [type(citation) for citation in response.citations] == [str, int, type(None), float, bool, dict, list]
+ assert response.search_results == search_results
+ assert [type(result) for result in response.search_results] == [dict, str, int]
+ assert repr(response.usage.citation_tokens) == "18"
diff --git a/tests/unit/llms/perplexity/embedding/test_perplexity_embedding_transformation.py b/tests/unit/llms/perplexity/embedding/test_perplexity_embedding_transformation.py
index 6ca7072e7ab..56afc7a1d67 100644
--- a/tests/unit/llms/perplexity/embedding/test_perplexity_embedding_transformation.py
+++ b/tests/unit/llms/perplexity/embedding/test_perplexity_embedding_transformation.py
@@ -3,12 +3,15 @@ Unit tests for Perplexity embedding transformation logic.
"""
import base64
+import datetime
import struct
from unittest.mock import MagicMock
import httpx
import pytest
+from pydantic import ValidationError
+from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.perplexity.embedding.transformation import (
PerplexityEmbeddingConfig,
PerplexityEmbeddingError,
@@ -297,3 +300,100 @@ class TestPerplexityEmbeddingProviderConfig:
)
assert config is not None
assert isinstance(config, PerplexityEmbeddingConfig)
+
+
+FLOAT_ROW = {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}
+
+
+def _transform(payload: object) -> EmbeddingResponse:
+ logging_obj = Logging(
+ model="perplexity/pplx-embed-v1-0.6b",
+ messages=[],
+ stream=False,
+ call_type="embedding",
+ start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+ litellm_call_id="call-id",
+ function_id="function-id",
+ )
+ return PerplexityEmbeddingConfig().transform_embedding_response(
+ model="pplx-embed-v1-0.6b",
+ raw_response=httpx.Response(200, json=payload),
+ model_response=EmbeddingResponse(),
+ logging_obj=logging_obj,
+ )
+
+
+def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]:
+ assert response.usage is not None
+ return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens)
+
+
+@pytest.mark.parametrize(
+ ("usage", "expected"),
+ [
+ ({"prompt_tokens": 4, "total_tokens": 6}, (4, 0, 6)),
+ ({"total_tokens": 6}, (6, 0, 6)),
+ ({"prompt_tokens": 0, "total_tokens": 6}, (6, 0, 6)),
+ ({"prompt_tokens": None, "total_tokens": 6}, (6, 0, 6)),
+ ({"prompt_tokens": 4}, (4, 0, 0)),
+ ({}, (0, 0, 0)),
+ ({"prompt_tokens": None, "total_tokens": None}, (0, 0, 0)),
+ ({"prompt_tokens": "", "total_tokens": []}, (0, 0, 0)),
+ ({"prompt_tokens": "12", "total_tokens": 12.0}, (12, 0, 12)),
+ ],
+)
+def test_perplexity_embedding_usage_falls_back_to_total_tokens_and_tolerates_loose_counts(usage, expected):
+ response = _transform({"data": [FLOAT_ROW], "usage": usage})
+
+ assert response.data == [FLOAT_ROW]
+ assert _token_counts(response) == expected
+
+
+def test_perplexity_embedding_response_without_usage_reports_zero_tokens():
+ assert _token_counts(_transform({"data": [FLOAT_ROW]})) == (0, 0, 0)
+
+
+def test_perplexity_embedding_rows_keep_their_fields_in_order_around_the_decoded_vector():
+ encoded = base64.b64encode(struct.pack("3b", 127, 0, -127)).decode()
+
+ response = _transform(
+ {"data": [{"z": 1, "embedding": encoded, "index": 0}, {"index": 1, "embedding": [0.5]}, {"index": 2}]}
+ )
+
+ assert [list(row.items()) for row in response.data] == [
+ [("z", 1), ("embedding", [1.0, 0.0, -1.0]), ("index", 0)],
+ [("index", 1), ("embedding", [0.5])],
+ [("index", 2), ("embedding", None)],
+ ]
+
+
+@pytest.mark.parametrize("payload", [{}, {"data": []}, {"data": ""}, {"data": {}}])
+def test_perplexity_embedding_response_without_rows_has_no_data(payload):
+ assert _transform(payload).data == []
+
+
+@pytest.mark.parametrize("row", [None, 3, "ab", [], [["embedding", [1.0]]]])
+def test_perplexity_embedding_row_that_is_not_an_object_is_rejected(row):
+ with pytest.raises(ValidationError):
+ _transform({"data": [FLOAT_ROW, row]})
+
+
+@pytest.mark.parametrize("usage", [None, [], "tokens", 3])
+def test_perplexity_embedding_usage_that_is_not_an_object_is_rejected(usage):
+ with pytest.raises(ValidationError):
+ _transform({"data": [FLOAT_ROW], "usage": usage})
+
+
+@pytest.mark.parametrize(
+ "usage",
+ [
+ {"prompt_tokens": "many"},
+ {"total_tokens": 1.5},
+ {"prompt_tokens": [1]},
+ {"prompt_tokens": 4, "total_tokens": 1.5},
+ {"prompt_tokens": 4, "total_tokens": "many"},
+ ],
+)
+def test_perplexity_embedding_token_count_that_is_not_an_integer_is_rejected(usage):
+ with pytest.raises(ValidationError):
+ _transform({"data": [FLOAT_ROW], "usage": usage})
diff --git a/tests/unit/llms/perplexity/search/test_perplexity_search_transformation.py b/tests/unit/llms/perplexity/search/test_perplexity_search_transformation.py
index 267f8eacb08..af4d8c0c706 100644
--- a/tests/unit/llms/perplexity/search/test_perplexity_search_transformation.py
+++ b/tests/unit/llms/perplexity/search/test_perplexity_search_transformation.py
@@ -1,11 +1,17 @@
+import datetime
import json
from collections.abc import Iterator
from typing import Final
+import httpx
import pytest
import respx
+from pydantic import ValidationError
import litellm
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
+from litellm.llms.perplexity.search.transformation import PerplexitySearchConfig
PERPLEXITY_SEARCH_URL: Final = "https://api.perplexity.ai/search"
@@ -66,6 +72,116 @@ def test_search_sends_max_results_in_request_body(respx_mock: respx.MockRouter)
assert [result.url for result in response.results] == ["https://example.com/ml"]
+def _transform(payload: object) -> SearchResponse:
+ logging_obj = Logging(
+ model="perplexity/search",
+ messages=[],
+ stream=False,
+ call_type="search",
+ start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+ litellm_call_id="call-id",
+ function_id="function-id",
+ )
+ return PerplexitySearchConfig().transform_search_response(
+ raw_response=httpx.Response(200, json=payload), logging_obj=logging_obj
+ )
+
+
+def test_perplexity_results_keep_provider_order_fields_and_dates():
+ response = _transform(
+ {
+ "id": "search-1",
+ "server_time": None,
+ "results": [
+ {
+ "title": "LiteLLM",
+ "url": "https://example.com/litellm",
+ "snippet": "Call every LLM API",
+ "date": "2024-01-05",
+ "last_updated": "2024-02-01",
+ },
+ {
+ "title": "Docs",
+ "url": "https://example.com/docs",
+ "snippet": "Docs",
+ "date": None,
+ "last_updated": None,
+ },
+ ],
+ }
+ )
+
+ assert response.object == "search"
+ assert response.results == [
+ SearchResult(
+ title="LiteLLM",
+ url="https://example.com/litellm",
+ snippet="Call every LLM API",
+ date="2024-01-05",
+ last_updated="2024-02-01",
+ ),
+ SearchResult(title="Docs", url="https://example.com/docs", snippet="Docs", date=None, last_updated=None),
+ ]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [{}, {"id": "search-1"}, {"id": "search-1", "results": []}, {"results": ""}, {"results": {}}],
+)
+def test_perplexity_response_without_results_is_empty(payload):
+ assert _transform(payload).results == []
+
+
+def test_perplexity_result_missing_optional_fields_defaults_to_empty_strings_and_no_dates():
+ assert _transform({"results": [{"score": 0.5}]}).results == [
+ SearchResult(title="", url="", snippet="", date=None, last_updated=None)
+ ]
+
+
+@pytest.mark.parametrize(
+ ("field", "value"),
+ [
+ ("title", None),
+ ("title", 7),
+ ("url", None),
+ ("url", ["https://example.com"]),
+ ("snippet", None),
+ ("snippet", {"text": "snippet"}),
+ ("date", 1704412800),
+ ("last_updated", ["2024-02-01"]),
+ ],
+)
+def test_perplexity_result_with_field_of_wrong_type_is_rejected(field, value):
+ result = {"title": "T", "url": "https://example.com", "snippet": "S", field: value}
+
+ with pytest.raises(ValidationError):
+ _transform({"results": [result]})
+
+
+def test_perplexity_result_with_several_fields_of_wrong_type_reports_every_field():
+ result = {"title": 7, "url": None, "snippet": ["snippet"], "date": 1704412800, "last_updated": ["2024-02-01"]}
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"results": [result]})
+
+ assert [error["loc"] for error in exc_info.value.errors()] == [
+ ("title",),
+ ("url",),
+ ("snippet",),
+ ("date",),
+ ("last_updated",),
+ ]
+
+
+def test_perplexity_malformed_result_is_rejected_without_naming_its_position():
+ results = [{}] * 429 + ["not-an-object"]
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"results": results})
+
+ assert "429" not in str(exc_info.value)
+
+
@pytest.mark.asyncio
async def test_router_search_with_search_tools(
respx_mock: respx.MockRouter, httpx_transport: None
diff --git a/tests/unit/llms/replicate/chat/test_handler.py b/tests/unit/llms/replicate/chat/test_handler.py
new file mode 100644
index 00000000000..756a2a21db3
--- /dev/null
+++ b/tests/unit/llms/replicate/chat/test_handler.py
@@ -0,0 +1,206 @@
+import httpx
+import pytest
+import respx
+from pydantic import ValidationError
+
+import litellm
+from litellm.exceptions import MidStreamFallbackError
+from litellm.llms.custom_httpx.http_handler import HTTPHandler
+from litellm.llms.replicate.chat.handler import handle_prediction_response_streaming
+from litellm.llms.replicate.common_utils import ReplicateError
+
+POLL_URL = "https://api.replicate.com/v1/predictions/test-prediction"
+CREATE_URL = "https://api.replicate.com/v1/models/meta/meta-llama-3-8b-instruct/predictions"
+MANY_TOKENS = [f"token-{index} " for index in range(500)]
+
+
+@pytest.fixture(autouse=True)
+def _poll_without_waiting(monkeypatch: pytest.MonkeyPatch):
+ monkeypatch.setattr("litellm.llms.replicate.chat.handler.REPLICATE_POLLING_DELAY_SECONDS", 0)
+
+
+def _poll(*bodies: object) -> list[dict[str, object]]:
+ with respx.mock:
+ respx.get(POLL_URL).mock(side_effect=[httpx.Response(200, json=body) for body in bodies])
+ return list(
+ handle_prediction_response_streaming(
+ POLL_URL,
+ "replicate-test-key",
+ lambda _message: None,
+ headers={"Authorization": "Bearer replicate-test-key"},
+ http_client=HTTPHandler(),
+ )
+ )
+
+
+def test_streaming_poll_yields_only_the_text_added_since_the_previous_poll():
+ chunks = _poll(
+ {"status": "starting"},
+ {"status": "processing", "output": ["Hel"]},
+ {"status": "processing", "output": ["Hel", "lo"]},
+ {"status": "succeeded", "output": ["Hel", "lo", ", world"], "metrics": {"predict_time": 1.5}},
+ )
+
+ assert chunks == [
+ {"output": "Hel", "status": "processing"},
+ {"output": "lo", "status": "processing"},
+ {"output": ", world", "status": "succeeded"},
+ ]
+
+
+@pytest.mark.parametrize(
+ ("output", "text"),
+ [
+ (["Hello", ", ", "world"], "Hello, world"),
+ ("Hello", "Hello"),
+ ({"Hel": 1, "lo": [2]}, "Hello"),
+ (["12", "12.0"], "1212.0"),
+ ([], ""),
+ ("", ""),
+ ({}, ""),
+ (MANY_TOKENS, "".join(MANY_TOKENS)),
+ ],
+)
+def test_streaming_poll_joins_the_prediction_output_the_way_str_join_does(output: object, text: str):
+ assert _poll({"status": "succeeded", "output": output}) == [{"output": text, "status": "succeeded"}]
+
+
+@pytest.mark.parametrize("status", [None, 1, ["succeeded"], {"state": "succeeded"}])
+def test_streaming_poll_passes_an_unrecognised_status_through_and_keeps_polling(status: object):
+ chunks = _poll({"status": status, "output": ["a"]}, {"status": "succeeded", "output": ["a", "b"]})
+
+ assert chunks == [{"output": "a", "status": status}, {"output": "b", "status": "succeeded"}]
+
+
+@pytest.mark.parametrize(
+ "output",
+ [
+ None,
+ 5,
+ 0,
+ 1.5,
+ True,
+ ["a", 1],
+ ["a", None],
+ [["a"]],
+ [{"a": 1}],
+ [*MANY_TOKENS[:403], 403, *MANY_TOKENS[404:]],
+ ],
+)
+def test_streaming_poll_reports_an_output_that_is_not_text_as_unprocessable(output: object):
+ with pytest.raises(ReplicateError) as exc_info:
+ _poll({"status": "succeeded", "output": output})
+
+ assert exc_info.value.status_code == 422
+ assert exc_info.value.message == f"Unable to parse response. Got={output}"
+
+
+@pytest.mark.parametrize(
+ ("body", "message"),
+ [
+ ({"status": "failed", "error": "boom"}, "Error: boom"),
+ ({"status": "failed", "error": None}, "Error: None"),
+ ({"status": "failed"}, "Error: "),
+ ({"status": "failed", "error": {"detail": "bad input"}}, "Error: {'detail': 'bad input'}"),
+ ],
+)
+def test_streaming_poll_raises_the_provider_error_of_a_failed_prediction(body: object, message: str):
+ with pytest.raises(ReplicateError) as exc_info:
+ _poll(body)
+
+ assert exc_info.value.status_code == 400
+ assert exc_info.value.message == message
+
+
+@pytest.mark.parametrize("body", [{}, {"output": ["a"]}])
+def test_streaming_poll_raises_key_error_when_the_prediction_has_no_status(body: object):
+ with pytest.raises(KeyError) as exc_info:
+ _poll(body)
+
+ assert exc_info.value.args == ("status",)
+
+
+@pytest.mark.parametrize(
+ "body", [[{"status": "succeeded"}], [], "Request timed out", ["secret-output", 429], 3, True, 1.5]
+)
+def test_streaming_poll_rejects_a_body_that_is_not_an_object_without_echoing_it(body: object):
+ with pytest.raises(ValidationError) as exc_info:
+ _poll(body)
+
+ assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [("dict_type", ())]
+ assert "secret-output" not in str(exc_info.value)
+ assert "429" not in str(exc_info.value)
+ assert "Request timed out" not in str(exc_info.value)
+
+
+def test_streaming_poll_keeps_the_chunks_yielded_before_a_malformed_body():
+ chunks = []
+ with respx.mock:
+ respx.get(POLL_URL).mock(
+ side_effect=[
+ httpx.Response(200, json={"status": "processing", "output": ["a"]}),
+ httpx.Response(200, json=["not", "an", "object"]),
+ ]
+ )
+ stream = handle_prediction_response_streaming(
+ POLL_URL, "replicate-test-key", lambda _message: None, headers={}, http_client=HTTPHandler()
+ )
+ chunks.append(next(stream))
+ with pytest.raises(ValidationError):
+ next(stream)
+
+ assert chunks == [{"output": "a", "status": "processing"}]
+
+
+def test_streaming_poll_retries_after_a_non_200_status():
+ with respx.mock:
+ respx.get(POLL_URL).mock(
+ side_effect=[
+ httpx.Response(503, text="busy"),
+ httpx.Response(200, json={"status": "succeeded", "output": ["ok"]}),
+ ]
+ )
+ chunks = list(
+ handle_prediction_response_streaming(
+ POLL_URL, "replicate-test-key", lambda _message: None, headers={}, http_client=HTTPHandler()
+ )
+ )
+
+ assert chunks == [{"output": "ok", "status": "succeeded"}]
+
+
+def _stream_completion(*poll_bodies: object) -> list[str | None]:
+ with respx.mock:
+ respx.post(CREATE_URL).mock(return_value=httpx.Response(201, json={"urls": {"get": POLL_URL}}))
+ respx.get(POLL_URL).mock(side_effect=[httpx.Response(200, json=body) for body in poll_bodies])
+ stream = litellm.completion(
+ model="replicate/meta/meta-llama-3-8b-instruct",
+ messages=[{"role": "user", "content": "hi"}],
+ stream=True,
+ api_key="replicate-test-key",
+ )
+ return [chunk.choices[0].delta.content for chunk in stream]
+
+
+def test_streaming_completion_emits_each_polled_delta_as_content():
+ contents = _stream_completion(
+ {"status": "processing", "output": ["Hel"]}, {"status": "succeeded", "output": ["Hel", "lo"]}
+ )
+
+ assert contents == ["Hel", "lo", None]
+
+
+def test_streaming_completion_maps_a_non_text_output_to_unprocessable_entity():
+ with pytest.raises(litellm.UnprocessableEntityError) as exc_info:
+ _stream_completion({"status": "succeeded", "output": ["a", 1]})
+
+ assert "Unable to parse response. Got=['a', 1]" in str(exc_info.value)
+
+
+@pytest.mark.parametrize("body", [["x"], "Request timed out", ["Request was throttled", 429]])
+def test_streaming_completion_maps_a_non_object_poll_body_to_an_api_error_whatever_text_it_carries(body: object):
+ with pytest.raises(MidStreamFallbackError) as exc_info:
+ _stream_completion(body)
+
+ assert type(exc_info.value.original_exception) is litellm.APIError
+ assert exc_info.value.original_exception.status_code == 500
diff --git a/tests/unit/llms/replicate/chat/test_transformation.py b/tests/unit/llms/replicate/chat/test_transformation.py
index 549e915a6d5..fe6dfbc8653 100644
--- a/tests/unit/llms/replicate/chat/test_transformation.py
+++ b/tests/unit/llms/replicate/chat/test_transformation.py
@@ -3,13 +3,14 @@ from typing import Final
from unittest.mock import AsyncMock, Mock, patch
import httpx
-import litellm
import pytest
from pydantic import ValidationError
+import litellm
from litellm.llms.replicate.chat.handler import async_completion
from litellm.llms.replicate.chat.handler import completion as replicate_completion
from litellm.llms.replicate.chat.transformation import ReplicateConfig
+from litellm.llms.replicate.common_utils import ReplicateError
from litellm.types.utils import ModelResponse
@@ -59,6 +60,106 @@ def test_transform_response_rejects_an_output_that_is_not_made_of_strings_withou
assert "secret-output" not in str(exc_info.value)
+def test_transform_response_reads_a_succeeded_prediction_that_carries_unmodelled_fields() -> None:
+ response = _transform(
+ httpx.Response(
+ 200,
+ json={
+ "id": "p1",
+ "status": "succeeded",
+ "output": ["Hello", " there"],
+ "metrics": {"predict_time": 1.5, "stages": [1, {"name": None}]},
+ "urls": {"get": "https://api.replicate.com/v1/predictions/p1"},
+ "error": None,
+ },
+ )
+ )
+
+ assert response.choices[0].message.content == "Hello there"
+ assert response.model == "replicate/acme/echo-model"
+
+
+@pytest.mark.parametrize(
+ ("body", "reported"),
+ [
+ pytest.param(
+ {"status": "failed", "error": "boom", "output": None},
+ "{'status': 'failed', 'error': 'boom', 'output': None}",
+ id="failed",
+ ),
+ pytest.param(
+ {"status": "processing", "output": ["partial"]},
+ "{'status': 'processing', 'output': ['partial']}",
+ id="still-processing",
+ ),
+ pytest.param(
+ {"detail": {"b": [1, 2.5, True], "a": "z"}},
+ "{'detail': {'b': [1, 2.5, True], 'a': 'z'}}",
+ id="no-status",
+ ),
+ pytest.param({}, "{}", id="empty-object"),
+ ],
+)
+def test_transform_response_reports_an_unsuccessful_prediction_body_as_received(
+ body: dict[str, object], reported: str
+) -> None:
+ with pytest.raises(ReplicateError) as exc_info:
+ _transform(httpx.Response(500, json=body, headers={"retry-after": "3"}))
+
+ assert exc_info.value.status_code == 422
+ assert exc_info.value.message == f"LiteLLM Error - prediction not succeeded - {reported}"
+ assert exc_info.value.headers["retry-after"] == "3"
+
+
+@pytest.mark.parametrize(
+ ("body", "echo"),
+ [
+ pytest.param(["secret-output"], "secret-output", id="list"),
+ pytest.param("secret-output", "secret-output", id="string"),
+ pytest.param(4815162342, "4815162342", id="number"),
+ pytest.param(None, "NoneType", id="null"),
+ ],
+)
+def test_transform_response_rejects_a_body_that_is_not_an_object_without_echoing_it(body: object, echo: str) -> None:
+ with pytest.raises(ValidationError) as exc_info:
+ _transform(httpx.Response(500, content=json.dumps(body).encode()))
+
+ assert echo not in str(exc_info.value)
+
+
+@pytest.mark.parametrize(
+ ("supplied", "stored"),
+ [
+ pytest.param({}, {}, id="nothing-supplied"),
+ pytest.param(
+ {"system_prompt": "", "max_new_tokens": 0, "temperature": 1, "debug": False},
+ {"system_prompt": "", "max_new_tokens": 0, "temperature": 1, "debug": False},
+ id="falsy-values-are-kept",
+ ),
+ pytest.param({"top_k": 40, "seed": None, "stop_sequences": None}, {"top_k": 40}, id="none-is-skipped"),
+ ],
+)
+def test_init_stores_each_supplied_value_except_none_as_class_level_config(
+ supplied: dict[str, object], stored: dict[str, object]
+) -> None:
+ class ScopedReplicateConfig(ReplicateConfig):
+ pass
+
+ ScopedReplicateConfig(**supplied)
+
+ assert ScopedReplicateConfig.get_config() == stored
+
+
+def test_init_called_again_with_none_keeps_the_value_stored_earlier() -> None:
+ class ScopedReplicateConfig(ReplicateConfig):
+ pass
+
+ ScopedReplicateConfig(top_k=40)
+ ScopedReplicateConfig(top_k=None, seed=5)
+
+ assert ScopedReplicateConfig.get_config() == {"top_k": 40, "seed": 5}
+
+
@pytest.mark.asyncio
@patch("litellm.llms.replicate.chat.handler.asyncio.sleep", new_callable=AsyncMock)
@patch("litellm.llms.replicate.chat.handler.get_async_httpx_client")
diff --git a/tests/unit/llms/sambanova/embedding/__init__.py b/tests/unit/llms/sambanova/embedding/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/sambanova/embedding/test_sambanova_embedding_transformation.py b/tests/unit/llms/sambanova/embedding/test_sambanova_embedding_transformation.py
new file mode 100644
index 00000000000..99085214b30
--- /dev/null
+++ b/tests/unit/llms/sambanova/embedding/test_sambanova_embedding_transformation.py
@@ -0,0 +1,89 @@
+import datetime
+
+import httpx
+import pytest
+from pydantic import ValidationError
+
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.sambanova.embedding.transformation import SambaNovaEmbeddingConfig
+from litellm.types.utils import EmbeddingResponse
+
+EMBEDDING_ROW = {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}
+
+
+def _transform(payload: object) -> EmbeddingResponse:
+ logging_obj = Logging(
+ model="sambanova/E5-Mistral-7B-Instruct",
+ messages=[],
+ stream=False,
+ call_type="embedding",
+ start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+ litellm_call_id="call-id",
+ function_id="function-id",
+ )
+ return SambaNovaEmbeddingConfig().transform_embedding_response(
+ model="E5-Mistral-7B-Instruct",
+ raw_response=httpx.Response(200, json=payload),
+ model_response=EmbeddingResponse(),
+ logging_obj=logging_obj,
+ api_key="sk-test",
+ request_data={},
+ optional_params={},
+ litellm_params={},
+ )
+
+
+def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]:
+ assert response.usage is not None
+ return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens)
+
+
+def test_sambanova_embedding_response_carries_provider_fields_and_usage():
+ response = _transform(
+ {
+ "object": "list",
+ "model": "E5-Mistral-7B-Instruct",
+ "data": [EMBEDDING_ROW],
+ "usage": {"prompt_tokens": 8, "total_tokens": 9},
+ }
+ )
+
+ assert response.model == "E5-Mistral-7B-Instruct"
+ assert response.object == "list"
+ assert response.data == [EMBEDDING_ROW]
+ assert _token_counts(response) == (8, 0, 9)
+
+
+@pytest.mark.parametrize(
+ ("payload", "expected"),
+ [
+ ({"data": [EMBEDDING_ROW]}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 4}}, (4, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 6}}, (0, 0, 6)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": None, "total_tokens": None}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": "", "total_tokens": []}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": "12", "total_tokens": 12.0}}, (12, 0, 12)),
+ (
+ {"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 3, "total_tokens": 5, "completion_tokens": 2}},
+ (3, 0, 5),
+ ),
+ ],
+)
+def test_sambanova_embedding_usage_tolerates_missing_null_and_loosely_typed_counts(payload, expected):
+ response = _transform(payload)
+
+ assert response.data == [EMBEDDING_ROW]
+ assert _token_counts(response) == expected
+
+
+@pytest.mark.parametrize("usage", [None, [], "tokens", 3])
+def test_sambanova_embedding_usage_that_is_not_an_object_is_rejected(usage):
+ with pytest.raises(ValidationError):
+ _transform({"data": [EMBEDDING_ROW], "usage": usage})
+
+
+@pytest.mark.parametrize("usage", [{"prompt_tokens": "many"}, {"total_tokens": 1.5}, {"prompt_tokens": [1]}])
+def test_sambanova_embedding_token_count_that_is_not_an_integer_is_rejected(usage):
+ with pytest.raises(ValidationError):
+ _transform({"data": [EMBEDDING_ROW], "usage": usage})
diff --git a/tests/unit/llms/searchapi/search/test_searchapi_search_transformation.py b/tests/unit/llms/searchapi/search/test_searchapi_search_transformation.py
index 117b194b732..8f205e6356e 100644
--- a/tests/unit/llms/searchapi/search/test_searchapi_search_transformation.py
+++ b/tests/unit/llms/searchapi/search/test_searchapi_search_transformation.py
@@ -5,9 +5,10 @@ from unittest.mock import Mock, patch
import httpx
import pytest
import respx
+from pydantic import ValidationError
import litellm
-from litellm.llms.base_llm.search.transformation import SearchResponse
+from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
from litellm.llms.searchapi.search.transformation import SearchAPIConfig
@@ -222,6 +223,97 @@ class TestSearchAPIConfig:
assert "AND" in result
+def _transform(payload: object) -> SearchResponse:
+ return SearchAPIConfig().transform_search_response(raw_response=httpx.Response(200, json=payload), logging_obj=None)
+
+
+def test_searchapi_organic_results_keep_provider_order_fields_and_dates():
+ response = _transform(
+ {
+ "search_metadata": {"id": "search_1", "status": "Success"},
+ "search_information": {"total_results": 2},
+ "organic_results": [
+ {
+ "position": 1,
+ "title": "LiteLLM",
+ "link": "https://example.com/litellm",
+ "snippet": "Call every LLM API",
+ "date": "Jan 5, 2024",
+ "sitelinks": {"inline": [{"title": "Docs", "link": "https://example.com/docs"}]},
+ },
+ {"position": 2, "title": "Docs", "link": "https://example.com/docs", "snippet": "Docs", "date": None},
+ ],
+ "related_searches": [{"query": "litellm proxy"}],
+ }
+ )
+
+ assert response.object == "search"
+ assert response.results == [
+ SearchResult(
+ title="LiteLLM", url="https://example.com/litellm", snippet="Call every LLM API", date="Jan 5, 2024"
+ ),
+ SearchResult(title="Docs", url="https://example.com/docs", snippet="Docs", date=None),
+ ]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {},
+ {"error": "Invalid API key"},
+ {"search_metadata": {"status": "Success"}, "organic_results": []},
+ {"organic_results": ""},
+ {"organic_results": {}},
+ ],
+)
+def test_searchapi_response_without_organic_results_is_empty(payload):
+ assert _transform(payload).results == []
+
+
+def test_searchapi_result_missing_optional_fields_defaults_to_empty_strings_and_no_date():
+ assert _transform({"organic_results": [{"position": 1}]}).results == [
+ SearchResult(title="", url="", snippet="", date=None, last_updated=None)
+ ]
+
+
+@pytest.mark.parametrize(
+ ("field", "value"),
+ [
+ ("title", None),
+ ("title", 7),
+ ("link", None),
+ ("link", ["https://example.com"]),
+ ("snippet", None),
+ ("snippet", {"text": "snippet"}),
+ ("date", 1704412800),
+ ("date", ["Jan 5, 2024"]),
+ ],
+)
+def test_searchapi_result_with_field_of_wrong_type_is_rejected(field, value):
+ result = {"title": "T", "link": "https://example.com", "snippet": "S", "date": "Jan 5, 2024", field: value}
+
+ with pytest.raises(ValidationError):
+ _transform({"organic_results": [result]})
+
+
+def test_searchapi_result_with_several_fields_of_wrong_type_reports_every_field():
+ result = {"title": 7, "link": None, "snippet": ["snippet"], "date": 1704412800}
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"organic_results": [result]})
+
+ assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",), ("date",)]
+
+
+def test_searchapi_malformed_result_is_rejected_without_naming_its_position():
+ results = [{}] * 429 + ["not-an-object"]
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"organic_results": results})
+
+ assert "429" not in str(exc_info.value)
+
+
def test_search_request_and_result_mapping(
respx_mock: respx.MockRouter, httpx_transport: None
) -> None:
diff --git a/tests/unit/llms/searxng/search/test_searxng_search_transformation.py b/tests/unit/llms/searxng/search/test_searxng_search_transformation.py
index 5961cec548d..5ca7976491f 100644
--- a/tests/unit/llms/searxng/search/test_searxng_search_transformation.py
+++ b/tests/unit/llms/searxng/search/test_searxng_search_transformation.py
@@ -10,7 +10,9 @@ from urllib.parse import parse_qs, urlparse
import httpx
import pytest
+from pydantic import ValidationError
+from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
from litellm.llms.searxng.search.transformation import SearXNGSearchConfig
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@@ -322,6 +324,109 @@ class TestSearXNGSearchHeaders:
assert self.config.get_http_method() == "GET"
+def _transform(payload: object) -> SearchResponse:
+ raw_response = httpx.Response(200, json=payload, request=httpx.Request("GET", "https://searxng.example.com/search"))
+ return SearXNGSearchConfig().transform_search_response(raw_response=raw_response, logging_obj=None)
+
+
+def test_searxng_results_keep_provider_order_and_ignore_engine_metadata():
+ response = _transform(
+ {
+ "query": "litellm",
+ "number_of_results": 0,
+ "results": [
+ {
+ "url": "https://example.com/litellm",
+ "title": "LiteLLM",
+ "content": "Call every LLM API",
+ "publishedDate": "2025-01-15T00:00:00",
+ "pubdate": "2025-01-10",
+ "thumbnail": None,
+ "engines": ["google", "bing"],
+ "positions": [1, 2],
+ "score": 4.0,
+ },
+ {"url": "https://example.com/docs", "title": "Docs", "content": "Docs", "publishedDate": None},
+ ],
+ "answers": [],
+ "infoboxes": [],
+ "unresponsive_engines": [["brave", "timeout"]],
+ }
+ )
+
+ assert response.object == "search"
+ assert response.results == [
+ SearchResult(
+ title="LiteLLM",
+ url="https://example.com/litellm",
+ snippet="Call every LLM API",
+ date="2025-01-15T00:00:00",
+ ),
+ SearchResult(title="Docs", url="https://example.com/docs", snippet="Docs", date=None),
+ ]
+
+
+@pytest.mark.parametrize(
+ ("dates", "expected"),
+ [
+ ({"publishedDate": "2025-01-15", "pubdate": "2025-01-10"}, "2025-01-15"),
+ ({"publishedDate": None, "pubdate": "2025-01-10"}, "2025-01-10"),
+ ({"publishedDate": "", "pubdate": "2025-01-10"}, "2025-01-10"),
+ ({"publishedDate": 0, "pubdate": "2025-01-10"}, "2025-01-10"),
+ ({"publishedDate": None, "pubdate": None}, None),
+ ({"publishedDate": "", "pubdate": ""}, ""),
+ ({"publishedDate": 0}, None),
+ ({}, None),
+ ],
+)
+def test_searxng_date_prefers_published_date_and_falls_back_to_pubdate_when_it_is_empty(dates, expected):
+ response = _transform({"results": [{"title": "T", "url": "https://example.com", "content": "C", **dates}]})
+
+ assert response.results == [SearchResult(title="T", url="https://example.com", snippet="C", date=expected)]
+
+
+@pytest.mark.parametrize("payload", [{}, {"results": ""}, {"results": {}}])
+def test_searxng_response_with_empty_non_list_results_is_empty(payload):
+ assert _transform(payload).results == []
+
+
+@pytest.mark.parametrize(
+ ("field", "value"),
+ [
+ ("title", None),
+ ("title", 7),
+ ("url", None),
+ ("content", None),
+ ("content", ["snippet"]),
+ ("publishedDate", 1736899200),
+ ("pubdate", ["2025-01-10"]),
+ ],
+)
+def test_searxng_result_with_field_of_wrong_type_is_rejected(field, value):
+ result = {"title": "T", "url": "https://example.com", "content": "C", field: value}
+
+ with pytest.raises(ValidationError):
+ _transform({"results": [result]})
+
+
+def test_searxng_result_with_several_fields_of_wrong_type_reports_every_field():
+ result = {"title": 7, "url": None, "content": ["snippet"], "publishedDate": 1736899200}
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"results": [result]})
+
+ assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",), ("date",)]
+
+
+def test_searxng_malformed_result_is_rejected_without_naming_its_position():
+ results = [{}] * 429 + ["not-an-object"]
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"results": results})
+
+ assert "429" not in str(exc_info.value)
+
+
@pytest.fixture(autouse=True)
def _vcr_outcome_gate(request, vcr):
install_live_call_probe(request, vcr)
diff --git a/tests/unit/llms/serper/search/test_serper_search_transformation.py b/tests/unit/llms/serper/search/test_serper_search_transformation.py
index acc9c634c47..92b6984f7ca 100644
--- a/tests/unit/llms/serper/search/test_serper_search_transformation.py
+++ b/tests/unit/llms/serper/search/test_serper_search_transformation.py
@@ -5,9 +5,13 @@ Tests for Serper Search API integration.
import os
from unittest.mock import AsyncMock, MagicMock, patch
+import httpx
import pytest
+from pydantic import ValidationError
import litellm
+from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
+from litellm.llms.serper.search.transformation import SerperSearchConfig
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@@ -194,6 +198,85 @@ class TestSerperSearch:
assert len(response.results) == 0
+def _transform(payload: object) -> SearchResponse:
+ return SerperSearchConfig().transform_search_response(
+ raw_response=httpx.Response(200, json=payload), logging_obj=None
+ )
+
+
+def test_serper_organic_results_keep_provider_order_and_ignore_other_sections():
+ response = _transform(
+ {
+ "searchParameters": {"q": "litellm", "type": "search", "engine": "google"},
+ "knowledgeGraph": {"title": "LiteLLM", "attributes": {"Founded": "2023"}},
+ "organic": [
+ {
+ "title": "LiteLLM",
+ "link": "https://example.com/litellm",
+ "snippet": "Call every LLM API",
+ "date": "Jan 15, 2025",
+ "position": 1,
+ "sitelinks": [{"title": "Docs", "link": "https://example.com/docs"}],
+ },
+ {"title": "Docs", "link": "https://example.com/docs", "snippet": "Docs", "position": 2},
+ {"title": "Blog", "link": "https://example.com/blog", "snippet": "Blog", "date": None},
+ ],
+ "peopleAlsoAsk": [{"question": "What is LiteLLM?", "title": "FAQ", "link": "https://example.com/faq"}],
+ "credits": 1,
+ }
+ )
+
+ assert response.object == "search"
+ assert response.results == [
+ SearchResult(
+ title="LiteLLM", url="https://example.com/litellm", snippet="Call every LLM API", date="Jan 15, 2025"
+ ),
+ SearchResult(title="Docs", url="https://example.com/docs", snippet="Docs"),
+ SearchResult(title="Blog", url="https://example.com/blog", snippet="Blog"),
+ ]
+
+
+def test_serper_result_without_optional_fields_gets_empty_defaults():
+ assert _transform({"organic": [{"position": 1}]}).results == [SearchResult(title="", url="", snippet="")]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [{}, {"organic": []}, {"organic": ""}, {"organic": {}}, {"message": "Unauthorized.", "statusCode": 403}],
+)
+def test_serper_response_without_organic_results_is_empty(payload):
+ assert _transform(payload).results == []
+
+
+@pytest.mark.parametrize(
+ ("field", "value"),
+ [("title", None), ("title", 7), ("link", None), ("snippet", ["snippet"]), ("date", 1736899200)],
+)
+def test_serper_result_with_field_of_wrong_type_is_rejected(field, value):
+ result = {"title": "T", "link": "https://example.com", "snippet": "S", field: value}
+
+ with pytest.raises(ValidationError):
+ _transform({"organic": [result]})
+
+
+def test_serper_result_with_several_fields_of_wrong_type_reports_every_field():
+ result = {"title": 7, "link": None, "snippet": ["snippet"], "date": 1736899200}
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"organic": [result]})
+
+ assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",), ("date",)]
+
+
+def test_serper_malformed_result_is_rejected_without_naming_its_position():
+ results = [{}] * 429 + ["not-an-object"]
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"organic": results})
+
+ assert "429" not in str(exc_info.value)
+
+
@pytest.fixture(autouse=True)
def _vcr_outcome_gate(request, vcr):
install_live_call_probe(request, vcr)
diff --git a/tests/unit/llms/tavily/search/test_tavily_search_transformation.py b/tests/unit/llms/tavily/search/test_tavily_search_transformation.py
index bd43588fb96..f66aaffcf48 100644
--- a/tests/unit/llms/tavily/search/test_tavily_search_transformation.py
+++ b/tests/unit/llms/tavily/search/test_tavily_search_transformation.py
@@ -5,9 +5,13 @@ Tests for Tavily Search API integration.
import os
from unittest.mock import AsyncMock, MagicMock, patch
+import httpx
import pytest
+from pydantic import ValidationError
import litellm
+from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
+from litellm.llms.tavily.search.transformation import TavilySearchConfig
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@@ -90,6 +94,93 @@ class TestTavilySearch:
assert first_result.snippet == "This is a test snippet for result 1"
+def _transform(payload: object) -> SearchResponse:
+ return TavilySearchConfig().transform_search_response(
+ raw_response=httpx.Response(200, json=payload), logging_obj=None
+ )
+
+
+def test_tavily_results_keep_provider_order_and_ignore_answer_and_usage_metadata():
+ response = _transform(
+ {
+ "query": "litellm",
+ "answer": "LiteLLM is a gateway.",
+ "follow_up_questions": None,
+ "images": [{"url": "https://example.com/image.png", "description": "logo"}],
+ "results": [
+ {
+ "title": "LiteLLM",
+ "url": "https://example.com/litellm",
+ "content": "Call every LLM API",
+ "score": 0.81,
+ "raw_content": None,
+ "favicon": "https://example.com/favicon.ico",
+ },
+ {"title": "Docs", "url": "https://example.com/docs", "content": "Docs", "raw_content": "Raw docs"},
+ ],
+ "response_time": 1.67,
+ "usage": {"credits": 1},
+ "request_id": "123e4567-e89b-12d3-a456-426614174111",
+ }
+ )
+
+ assert response.object == "search"
+ assert response.results == [
+ SearchResult(title="LiteLLM", url="https://example.com/litellm", snippet="Call every LLM API"),
+ SearchResult(title="Docs", url="https://example.com/docs", snippet="Docs"),
+ ]
+
+
+def test_tavily_result_without_title_url_or_content_gets_empty_strings():
+ assert _transform({"results": [{"score": 0.5, "raw_content": None}]}).results == [
+ SearchResult(title="", url="", snippet="", date=None, last_updated=None)
+ ]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [{}, {"query": "litellm", "answer": None}, {"results": []}, {"results": ""}, {"results": {}}],
+)
+def test_tavily_response_without_results_is_empty(payload):
+ assert _transform(payload).results == []
+
+
+@pytest.mark.parametrize(
+ ("field", "value"),
+ [
+ ("title", None),
+ ("title", 7),
+ ("url", None),
+ ("url", ["https://example.com"]),
+ ("content", None),
+ ("content", {"text": "snippet"}),
+ ],
+)
+def test_tavily_result_with_field_of_wrong_type_is_rejected(field, value):
+ result = {"title": "T", "url": "https://example.com", "content": "S", field: value}
+
+ with pytest.raises(ValidationError):
+ _transform({"results": [result]})
+
+
+def test_tavily_result_with_several_fields_of_wrong_type_reports_every_field():
+ result = {"title": 7, "url": None, "content": ["snippet"]}
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"results": [result]})
+
+ assert [error["loc"] for error in exc_info.value.errors()] == [("title",), ("url",), ("snippet",)]
+
+
+def test_tavily_malformed_result_is_rejected_without_naming_its_position():
+ results = [{}] * 429 + ["not-an-object"]
+
+ with pytest.raises(ValidationError) as exc_info:
+ _transform({"results": results})
+
+ assert "429" not in str(exc_info.value)
+
+
@pytest.fixture(autouse=True)
def _vcr_outcome_gate(request, vcr):
install_live_call_probe(request, vcr)
diff --git a/tests/unit/llms/vertex_ai/fine_tuning/test_handler.py b/tests/unit/llms/vertex_ai/fine_tuning/test_handler.py
index ce9def6ab1d..5447d553f30 100644
--- a/tests/unit/llms/vertex_ai/fine_tuning/test_handler.py
+++ b/tests/unit/llms/vertex_ai/fine_tuning/test_handler.py
@@ -3,7 +3,11 @@ import json
from typing import Optional
from unittest.mock import AsyncMock, MagicMock, patch
+import google.auth
+import google.auth.credentials
+import httpx
import pytest
+import respx
import litellm
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@@ -572,3 +576,105 @@ def _restore_litellm_state(state) -> None:
for attr, value in state.items():
if hasattr(litellm, attr):
setattr(litellm, attr, value)
+
+
+_VERTEX_TUNING_JOB = {
+ "name": "projects/test-project/locations/us-central1/tuningJobs/42",
+ "tunedModelDisplayName": "gemini-1.0-pro-002-tuned",
+ "baseModel": "gemini-1.0-pro-002",
+ "supervisedTuningSpec": {"trainingDatasetUri": "gs://bucket/train.jsonl"},
+ "state": "JOB_STATE_RUNNING",
+ "createTime": "2024-12-31T22:40:20.211140Z",
+ "updateTime": "2024-12-31T22:40:20.211140Z",
+}
+_BOTH_CALL_STYLES = pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"])
+
+
+class _StaticGoogleCredentials(google.auth.credentials.Credentials):
+ def refresh(self, request):
+ self.token = "test-token"
+
+
+@pytest.fixture
+def vertex_tuning_endpoint(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> respx.Route:
+ monkeypatch.setattr(google.auth, "default", lambda scopes: (_StaticGoogleCredentials(), "test-project"))
+ monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+ return respx_mock.post(host="us-central1-aiplatform.googleapis.com")
+
+
+def _vertex_json_response(payload: object) -> httpx.Response:
+ return httpx.Response(200, content=json.dumps(payload).encode())
+
+
+async def _created_vertex_tuning_job(run_async: bool, endpoint: respx.Route, payload: object):
+ endpoint.mock(return_value=_vertex_json_response(payload))
+ pending = VertexFineTuningAPI().create_fine_tuning_job(
+ _is_async=run_async,
+ create_fine_tuning_job_data=FineTuningJobCreate(
+ training_file="gs://bucket/train.jsonl", model="gemini-1.0-pro-002"
+ ),
+ vertex_project="test-project",
+ vertex_location="us-central1",
+ vertex_credentials=None,
+ api_base=None,
+ timeout=600.0,
+ )
+ return await pending if run_async else pending
+
+
+@_BOTH_CALL_STYLES
+@pytest.mark.asyncio()
+async def test_vertex_fine_tuning_job_creation_returns_the_tuning_job_as_an_openai_job(
+ run_async: bool, vertex_tuning_endpoint: respx.Route
+):
+ job = await _created_vertex_tuning_job(run_async, vertex_tuning_endpoint, _VERTEX_TUNING_JOB)
+
+ assert (job.id, job.created_at, job.fine_tuned_model, job.model, job.status, job.training_file) == (
+ "projects/test-project/locations/us-central1/tuningJobs/42",
+ 1735684820,
+ "gemini-1.0-pro-002-tuned",
+ "gemini-1.0-pro-002",
+ "running",
+ "gs://bucket/train.jsonl",
+ )
+
+
+@_BOTH_CALL_STYLES
+@pytest.mark.parametrize("body", [[_VERTEX_TUNING_JOB], [["state", "JOB_STATE_RUNNING"]], "state", 403, 2.5, None])
+@pytest.mark.asyncio()
+async def test_vertex_fine_tuning_job_creation_reports_a_tuning_job_body_that_is_not_an_object_as_a_type_error(
+ run_async: bool, vertex_tuning_endpoint: respx.Route, body: object
+):
+ with pytest.raises(TypeError, match=r"argument after \*\* must be a mapping"):
+ await _created_vertex_tuning_job(run_async, vertex_tuning_endpoint, body)
+
+
+@pytest.mark.parametrize(
+ "body",
+ [
+ _VERTEX_TUNING_JOB,
+ {"candidates": [{"content": {"parts": [{"text": "hello"}]}}], "usageMetadata": {"totalTokenCount": 3}},
+ {},
+ [{"a": 1}, 2, "x", None],
+ "text",
+ 403,
+ 2.5,
+ True,
+ None,
+ ],
+)
+@pytest.mark.asyncio()
+async def test_vertex_passthrough_post_returns_the_upstream_json_body_unchanged(
+ vertex_tuning_endpoint: respx.Route, body: object
+):
+ vertex_tuning_endpoint.mock(return_value=_vertex_json_response(body))
+
+ returned = await VertexFineTuningAPI().pass_through_vertex_ai_POST_request(
+ request_data={"baseModel": "gemini-1.0-pro-002"},
+ vertex_project="test-project",
+ vertex_location="us-central1",
+ vertex_credentials=None,
+ request_route="/tuningJobs",
+ )
+
+ assert (type(returned), returned) == (type(body), body)
diff --git a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py
index d3e94e5aa29..756e6441f1f 100644
--- a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py
+++ b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py
@@ -1,13 +1,18 @@
import base64
import json
import os
+from datetime import datetime
from io import BytesIO
from typing import Dict
from unittest.mock import MagicMock, patch
import httpx
import pytest
+from pydantic import ValidationError
+import litellm
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.vertex_ai.image_edit.vertex_gemini_transformation import (
VertexAIGeminiImageEditConfig,
)
@@ -410,3 +415,184 @@ class TestVertexAIImagenImageEditTransformation:
self.config._read_all_bytes(bytearray(b"test_bytearray"))
== b"test_bytearray"
)
+
+
+_IMAGEN_EDIT_IMAGE_KEY = "bytesBase64Encoded"
+
+
+def _imagen_edit_predict_response(payload: object) -> httpx.Response:
+ return httpx.Response(200, content=json.dumps(payload).encode(), headers={"content-type": "application/json"})
+
+
+def _imagen_edit_response_images(payload: object) -> list[tuple[object, object]]:
+ response = VertexAIImagenImageEditConfig().transform_image_edit_response(
+ model="imagen-3.0-capability-001",
+ raw_response=_imagen_edit_predict_response(payload),
+ logging_obj=Logging(
+ model="imagen-3.0-capability-001",
+ messages=[],
+ stream=False,
+ call_type="image_edit",
+ start_time=datetime(2026, 1, 1),
+ litellm_call_id="imagen-edit-call",
+ function_id="imagen-edit-function",
+ ),
+ )
+ return [(image.b64_json, image.url) for image in response.data]
+
+
+def _imagen_edit_rejection_text(payload: object) -> str:
+ with pytest.raises(ValidationError) as rejection:
+ _imagen_edit_response_images(payload)
+ return str(rejection.value)
+
+
+def _edited_imagen_images(payload: object) -> list[tuple[object, object]]:
+ transport = httpx.MockTransport(lambda request: _imagen_edit_predict_response(payload))
+ response = litellm.image_edit(
+ model="vertex_ai/imagen-3.0-capability-001",
+ image=BytesIO(b"image-bytes"),
+ prompt="add a hat",
+ api_base="https://vertex.invalid",
+ vertex_project="test-project",
+ vertex_location="us-central1",
+ client=HTTPHandler(client=httpx.Client(transport=transport)),
+ )
+ return [(image.b64_json, image.url) for image in response.data]
+
+
+def test_imagen_image_edit_response_lists_prediction_images_in_order_and_skips_entries_without_image_bytes():
+ images = _imagen_edit_response_images(
+ {
+ "predictions": [
+ {_IMAGEN_EDIT_IMAGE_KEY: "first", "mimeType": "image/png"},
+ "plain text",
+ [],
+ {},
+ ["other"],
+ {"raiFilteredReason": "blocked"},
+ {_IMAGEN_EDIT_IMAGE_KEY: "second"},
+ ],
+ "deployedModelId": "1",
+ }
+ )
+
+ assert images == [("first", None), ("second", None)]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {},
+ {"predictions": []},
+ {"predictions": ""},
+ {"predictions": "plain text"},
+ {"predictions": _IMAGEN_EDIT_IMAGE_KEY},
+ {"predictions": {}},
+ {"predictions": {"mimeType": "image/png"}},
+ {"predictions": [{}]},
+ {"predictions": [{"mimeType": "image/png"}]},
+ {"predictions": ["plain", ""]},
+ {"predictions": [["a", "b"], []]},
+ {"predictions": [[[_IMAGEN_EDIT_IMAGE_KEY]]]},
+ {"Predictions": [{_IMAGEN_EDIT_IMAGE_KEY: "a"}]},
+ {"predictions": [{"bytes_base64_encoded": "a"}]},
+ ],
+)
+def test_imagen_image_edit_response_without_prediction_image_bytes_has_no_images(payload: object):
+ assert _imagen_edit_response_images(payload) == []
+
+
+@pytest.mark.parametrize("image_bytes", [None, "", "aGVsbG8=", "1.5", "true", "é中"])
+def test_imagen_image_edit_response_keeps_the_image_bytes_value_exactly_as_sent(image_bytes: str | None):
+ images = _imagen_edit_response_images({"predictions": [{_IMAGEN_EDIT_IMAGE_KEY: image_bytes}]})
+
+ assert repr(images) == repr([(image_bytes, None)])
+
+
+def test_imagen_image_edit_response_finds_the_image_after_hundreds_of_imageless_predictions():
+ images = _imagen_edit_response_images({"predictions": [*[{}] * 403, {_IMAGEN_EDIT_IMAGE_KEY: "late"}]})
+
+ assert images == [("late", None)]
+
+
+@pytest.mark.parametrize("image_bytes", [5, 1.5, True, ["a"], {"a": 1}])
+def test_imagen_image_edit_response_rejects_non_string_image_bytes_naming_b64_json(image_bytes: object):
+ with pytest.raises(ValidationError, match="b64_json"):
+ _imagen_edit_response_images({"predictions": [{_IMAGEN_EDIT_IMAGE_KEY: image_bytes}]})
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {"predictions": None},
+ {"predictions": 7},
+ {"predictions": 1.5},
+ {"predictions": True},
+ {"predictions": {_IMAGEN_EDIT_IMAGE_KEY: "vertex-secret"}},
+ {"predictions": [None]},
+ {"predictions": [7]},
+ {"predictions": [1.5]},
+ {"predictions": [False]},
+ {"predictions": ["vertex-secret " + _IMAGEN_EDIT_IMAGE_KEY]},
+ {"predictions": [[_IMAGEN_EDIT_IMAGE_KEY, "vertex-secret"]]},
+ {"predictions": [{_IMAGEN_EDIT_IMAGE_KEY: "vertex-secret"}, None]},
+ ],
+)
+def test_imagen_image_edit_response_rejects_malformed_predictions_without_echoing_them(payload: object):
+ rejection_text = _imagen_edit_rejection_text(payload)
+
+ assert "vertex-secret" not in rejection_text and "input_value" not in rejection_text, rejection_text
+
+
+@pytest.mark.parametrize("position", [403, 429])
+@pytest.mark.parametrize("malformed", [None, 7, _IMAGEN_EDIT_IMAGE_KEY, [_IMAGEN_EDIT_IMAGE_KEY]])
+def test_imagen_image_edit_response_rejection_text_is_the_same_wherever_the_malformed_prediction_sits(
+ position: int, malformed: object
+):
+ rejection_at_start = _imagen_edit_rejection_text({"predictions": [malformed]})
+
+ rejection_at_position = _imagen_edit_rejection_text({"predictions": [*[{}] * position, malformed]})
+
+ assert rejection_at_position == rejection_at_start and str(position) not in rejection_at_position
+
+
+@pytest.mark.parametrize("body", [[], [{_IMAGEN_EDIT_IMAGE_KEY: "a"}], "predictions", 7, 2.5, True, None])
+def test_imagen_image_edit_response_rejects_a_predict_body_that_is_not_an_object_with_attribute_error(body: object):
+ with pytest.raises(AttributeError, match="object has no attribute 'get'"):
+ _imagen_edit_response_images(body)
+
+
+def test_imagen_image_edit_returns_the_prediction_images_of_the_predict_endpoint():
+ images = _edited_imagen_images(
+ {
+ "predictions": [
+ {_IMAGEN_EDIT_IMAGE_KEY: "first"},
+ {"mimeType": "image/png"},
+ {_IMAGEN_EDIT_IMAGE_KEY: "second"},
+ ]
+ }
+ )
+
+ assert images == [("first", None), ("second", None)]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {"predictions": None},
+ {"predictions": 7},
+ {"predictions": [None]},
+ {"predictions": ["vertex-secret " + _IMAGEN_EDIT_IMAGE_KEY]},
+ {"predictions": [*[{}] * 403, [_IMAGEN_EDIT_IMAGE_KEY, "vertex-secret"]]},
+ ],
+)
+def test_imagen_image_edit_maps_malformed_predictions_to_a_connection_error_without_echoing_them(payload: object):
+ with pytest.raises(litellm.APIConnectionError) as mapped:
+ _edited_imagen_images(payload)
+
+ assert (type(mapped.value), mapped.value.status_code, "vertex-secret" in str(mapped.value)) == (
+ litellm.APIConnectionError,
+ 500,
+ False,
+ )
diff --git a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py
index 43f6b47e9bd..5ac16283094 100644
--- a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py
+++ b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py
@@ -8,8 +8,9 @@ import pytest
from google.oauth2.credentials import Credentials
from pydantic import ValidationError
+import litellm
from litellm.litellm_core_utils.litellm_logging import Logging
-
+from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.vertex_ai.image_generation import (
get_vertex_ai_image_generation_config,
)
@@ -715,6 +716,293 @@ def test_gemini_image_generation_response_rejects_non_object_usage_metadata_with
assert "input_value" not in str(exc_info.value)
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {},
+ {"candidates": []},
+ {"candidates": ""},
+ {"candidates": {}},
+ {"candidates": [{}]},
+ {"candidates": [{"finishReason": "SAFETY"}]},
+ {"candidates": [{"content": {}}]},
+ {"candidates": [{"content": {"parts": ""}}]},
+ {"candidates": [{"content": {"parts": "plain text"}}]},
+ {"candidates": [{"content": {"parts": {"text": "only text"}}}]},
+ {"candidates": [{"content": {"parts": [{"text": "only text"}, "plain text", ["text"], [], ""]}}]},
+ {"candidates": [{"content": {"parts": [{"inlineData": {"mimeType": "image/png"}}]}}]},
+ {"candidates": [{"content": {"parts": [{"inlineData": "no image"}, {"inlineData": ["mimeType"]}]}}]},
+ {"candidates": [{"content": {"parts": [{"inlineData": ""}, {"inlineData": []}, {"inlineData": {}}]}}]},
+ {"candidates": [{"content": {"parts": [{"inline_data": {"data": "snake-case-is-not-read"}}]}}]},
+ ],
+)
+def test_gemini_image_generation_response_without_inline_image_data_has_no_images(payload: object):
+ assert _transform_gemini_response(payload).data == []
+
+
+def test_gemini_image_generation_response_keeps_image_order_and_thought_signatures():
+ response = _transform_gemini_response(
+ {
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {"text": "caption"},
+ {"inlineData": {"mimeType": "image/png", "data": "first"}, "thoughtSignature": "sig-1"},
+ "plain text",
+ ]
+ }
+ },
+ {
+ "content": {
+ "parts": [
+ {"inlineData": {"data": "second"}},
+ {"inlineData": {"data": None}, "thoughtSignature": ""},
+ {"inlineData": {"data": "fourth"}, "thoughtSignature": {"nested": ["sig"]}},
+ ]
+ }
+ },
+ ]
+ }
+ )
+
+ assert [(image.b64_json, image.provider_specific_fields) for image in response.data or []] == [
+ ("first", {"thought_signature": "sig-1"}),
+ ("second", None),
+ (None, None),
+ ("fourth", {"thought_signature": {"nested": ["sig"]}}),
+ ]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ ["vertex-secret"],
+ "vertex-secret",
+ {"candidates": None},
+ {"candidates": 7},
+ {"candidates": "vertex-secret"},
+ {"candidates": ["vertex-secret"]},
+ {"candidates": [{"content": None}]},
+ {"candidates": [{"content": "vertex-secret"}]},
+ {"candidates": [{"content": ["vertex-secret"]}]},
+ {"candidates": [{"content": {"parts": None}}]},
+ {"candidates": [{"content": {"parts": [None]}}]},
+ {"candidates": [{"content": {"parts": [7]}}]},
+ {"candidates": [{"content": {"parts": ["inlineData vertex-secret"]}}]},
+ {"candidates": [{"content": {"parts": [["inlineData", "vertex-secret"]]}}]},
+ {"candidates": [{"content": {"parts": {"inlineData": "vertex-secret"}}}]},
+ {"candidates": [{"content": {"parts": [{"inlineData": None}]}}]},
+ {"candidates": [{"content": {"parts": [{"inlineData": 7}]}}]},
+ {"candidates": [{"content": {"parts": [{"inlineData": "data vertex-secret"}]}}]},
+ {"candidates": [{"content": {"parts": [{"inlineData": ["data", "vertex-secret"]}]}}]},
+ ],
+)
+def test_gemini_image_generation_response_rejects_malformed_candidates_without_echoing_them(payload: object):
+ with pytest.raises(ValidationError) as exc_info:
+ _transform_gemini_response(payload)
+
+ assert "vertex-secret" not in str(exc_info.value)
+
+
+def _gemini_rejection_text(payload: object) -> str:
+ with pytest.raises(ValidationError) as exc_info:
+ _transform_gemini_response(payload)
+
+ return str(exc_info.value)
+
+
+@pytest.mark.parametrize("position", [403, 429])
+@pytest.mark.parametrize("malformed_entry", [None, 7, "inlineData"])
+def test_gemini_image_generation_response_rejection_text_is_the_same_wherever_the_malformed_entry_sits(
+ position: int, malformed_entry: object
+):
+ late_parts = [*[{"text": "caption"}] * position, malformed_entry]
+ late_candidates = [*[{"finishReason": "STOP"}] * position, malformed_entry]
+
+ assert _gemini_rejection_text({"candidates": [{"content": {"parts": late_parts}}]}) == _gemini_rejection_text(
+ {"candidates": [{"content": {"parts": [malformed_entry]}}]}
+ )
+ assert _gemini_rejection_text({"candidates": late_candidates}) == _gemini_rejection_text(
+ {"candidates": [malformed_entry]}
+ )
+
+
+_IMAGEN_IMAGE_KEY = "bytesBase64Encoded"
+
+
+def _imagen_predict_response(payload: object) -> httpx.Response:
+ return httpx.Response(200, content=json.dumps(payload).encode(), headers={"content-type": "application/json"})
+
+
+def _transform_imagen_response(payload: object, model_response: ImageResponse) -> ImageResponse:
+ return VertexAIImagenImageGenerationConfig().transform_image_generation_response(
+ model="imagegeneration@006",
+ raw_response=_imagen_predict_response(payload),
+ model_response=model_response,
+ logging_obj=Logging(
+ model="imagegeneration@006",
+ messages=[],
+ stream=False,
+ call_type="image_generation",
+ start_time=datetime.datetime(2026, 1, 1),
+ litellm_call_id="vertex-imagen-generation-test",
+ function_id="vertex-imagen-generation-test",
+ ),
+ request_data={},
+ optional_params={},
+ litellm_params={},
+ encoding=None,
+ )
+
+
+def _imagen_rejection_text(payload: object) -> str:
+ with pytest.raises(ValidationError) as exc_info:
+ _transform_imagen_response(payload, ImageResponse())
+
+ return str(exc_info.value)
+
+
+def _generate_imagen_image(payload: object) -> ImageResponse:
+ transport = httpx.MockTransport(lambda request: _imagen_predict_response(payload))
+ return litellm.image_generation(
+ model="vertex_ai/imagegeneration@006",
+ prompt="a cat",
+ api_base="https://vertex.invalid/v1/predict",
+ client=HTTPHandler(client=httpx.Client(transport=transport)),
+ )
+
+
+def test_imagen_image_generation_response_appends_prediction_images_in_order_to_the_given_response():
+ model_response = ImageResponse(data=[{"b64_json": "seeded", "url": None}])
+ seeded_image = model_response.data[0]
+
+ result = _transform_imagen_response(
+ {
+ "predictions": [
+ {_IMAGEN_IMAGE_KEY: "first", "mimeType": "image/png"},
+ {"raiFilteredReason": "filtered"},
+ "plain text",
+ [],
+ {_IMAGEN_IMAGE_KEY: "second"},
+ ],
+ "deployedModelId": "123",
+ },
+ model_response,
+ )
+
+ assert result is model_response
+ assert result.data[0] is seeded_image
+ assert [(image.b64_json, image.url) for image in result.data] == [
+ ("seeded", None),
+ ("first", None),
+ ("second", None),
+ ]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ {},
+ {"deployedModelId": "123"},
+ {"predictions": []},
+ {"predictions": ""},
+ {"predictions": "plain text"},
+ {"predictions": _IMAGEN_IMAGE_KEY},
+ {"predictions": {}},
+ {"predictions": {"mimeType": "image/png", "prompt": "a cat"}},
+ {"predictions": [{}]},
+ {"predictions": [{"mimeType": "image/png"}, {"raiFilteredReason": "filtered"}]},
+ {"predictions": ["plain text", ""]},
+ {"predictions": [["mimeType", "image/png"], []]},
+ {"predictions": [[[_IMAGEN_IMAGE_KEY]]]},
+ {"predictions": [{"bytes_base64_encoded": "snake-case-is-not-read"}]},
+ ],
+)
+def test_imagen_image_generation_response_without_prediction_image_bytes_has_no_images(payload: object):
+ assert _transform_imagen_response(payload, ImageResponse()).data == []
+
+
+@pytest.mark.parametrize("b64_value", [None, "", "aGVsbG8=", "1.5", "true"])
+def test_imagen_image_generation_response_keeps_the_image_bytes_value_exactly_as_sent(b64_value: object):
+ result = _transform_imagen_response({"predictions": [{_IMAGEN_IMAGE_KEY: b64_value}]}, ImageResponse())
+
+ assert [repr(image.b64_json) for image in result.data] == [repr(b64_value)]
+
+
+@pytest.mark.parametrize("b64_value", [5, 1.5, True, ["aGVsbG8="], {"data": "aGVsbG8="}])
+def test_imagen_image_generation_response_rejects_non_string_image_bytes_naming_b64_json(b64_value: object):
+ with pytest.raises(ValidationError, match="b64_json"):
+ _transform_imagen_response({"predictions": [{_IMAGEN_IMAGE_KEY: b64_value}]}, ImageResponse())
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ ["vertex-secret"],
+ "vertex-secret",
+ 7,
+ None,
+ {"predictions": None},
+ {"predictions": 7},
+ {"predictions": True},
+ {"predictions": [None]},
+ {"predictions": [7]},
+ {"predictions": [False]},
+ {"predictions": [f"{_IMAGEN_IMAGE_KEY} vertex-secret"]},
+ {"predictions": [[_IMAGEN_IMAGE_KEY, "vertex-secret"]]},
+ {"predictions": {_IMAGEN_IMAGE_KEY: "vertex-secret"}},
+ {"predictions": [{_IMAGEN_IMAGE_KEY: "kept"}, None]},
+ ],
+)
+def test_imagen_image_generation_response_rejects_malformed_predictions_without_echoing_them(payload: object):
+ rejection_text = _imagen_rejection_text(payload)
+
+ assert "vertex-secret" not in rejection_text
+ assert "input_value" not in rejection_text
+
+
+@pytest.mark.parametrize("position", [403, 429])
+@pytest.mark.parametrize("malformed_entry", [None, 7, _IMAGEN_IMAGE_KEY, [_IMAGEN_IMAGE_KEY]])
+def test_imagen_image_generation_response_rejection_text_is_the_same_wherever_the_malformed_prediction_sits(
+ position: int, malformed_entry: object
+):
+ late_predictions = [*[{_IMAGEN_IMAGE_KEY: "kept"}] * position, malformed_entry]
+
+ assert _imagen_rejection_text({"predictions": late_predictions}) == _imagen_rejection_text(
+ {"predictions": [malformed_entry]}
+ )
+
+
+def test_imagen_image_generation_returns_the_prediction_images_of_the_predict_endpoint():
+ result = _generate_imagen_image(
+ {"predictions": [{_IMAGEN_IMAGE_KEY: "first"}, {"mimeType": "image/png"}, {_IMAGEN_IMAGE_KEY: "second"}]}
+ )
+
+ assert [(image.b64_json, image.url) for image in result.data] == [("first", None), ("second", None)]
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ ["vertex-secret"],
+ "vertex-secret",
+ {"predictions": None},
+ {"predictions": [*[{_IMAGEN_IMAGE_KEY: "kept"}] * 403, f"{_IMAGEN_IMAGE_KEY} vertex-secret"]},
+ {"predictions": [*[{_IMAGEN_IMAGE_KEY: "kept"}] * 429, None]},
+ ],
+)
+def test_imagen_image_generation_maps_a_malformed_predict_body_to_a_connection_error_without_echoing_it(
+ payload: object,
+):
+ with pytest.raises(litellm.APIConnectionError) as exc_info:
+ _generate_imagen_image(payload)
+
+ assert type(exc_info.value) is litellm.APIConnectionError
+ assert exc_info.value.status_code == 500
+ assert "vertex-secret" not in str(exc_info.value)
+
+
def test_gemini_image_generation_validation_uses_supplied_vertex_token() -> None:
config: Final = VertexAIGeminiImageGenerationConfig()
credentials_json: Final = json.dumps({"type": "authorized_user", "client_id": "test-client"})
diff --git a/tests/unit/llms/vertex_ai/text_to_speech/test_text_to_speech_handler.py b/tests/unit/llms/vertex_ai/text_to_speech/test_text_to_speech_handler.py
new file mode 100644
index 00000000000..4a57316d5b3
--- /dev/null
+++ b/tests/unit/llms/vertex_ai/text_to_speech/test_text_to_speech_handler.py
@@ -0,0 +1,155 @@
+import base64
+import binascii
+import json
+from datetime import datetime
+from typing import Final
+
+import google.auth
+import google.auth.credentials
+import httpx
+import pytest
+import respx
+from pydantic import ValidationError
+
+import litellm
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.vertex_ai.text_to_speech.text_to_speech_handler import VertexTextToSpeechAPI
+
+SYNTHESIZE_URL: Final = "https://texttospeech.googleapis.com/v1/text:synthesize"
+AUDIO: Final = b"hello audio"
+ENCODED_AUDIO: Final = base64.b64encode(AUDIO).decode()
+BOTH_CALL_STYLES: Final = pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"])
+
+
+class _StaticGoogleCredentials(google.auth.credentials.Credentials):
+ def refresh(self, request: object) -> None:
+ self.token = "test-token"
+
+
+@pytest.fixture
+def synthesize_endpoint(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> respx.Route:
+ monkeypatch.setattr(google.auth, "default", lambda scopes: (_StaticGoogleCredentials(), "test-project"))
+ monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+ return respx_mock.post(SYNTHESIZE_URL)
+
+
+async def _synthesized_audio(run_async: bool, endpoint: respx.Route, payload: object) -> bytes:
+ endpoint.mock(return_value=httpx.Response(200, content=json.dumps(payload).encode()))
+ pending: Final = VertexTextToSpeechAPI().audio_speech(
+ logging_obj=Logging(
+ model="vertex_ai/",
+ messages=[],
+ stream=False,
+ call_type="speech",
+ start_time=datetime(2026, 1, 1),
+ litellm_call_id="speech-call",
+ function_id="speech-function",
+ ),
+ vertex_project="test-project",
+ vertex_location="us-central1",
+ vertex_credentials=None,
+ api_base=None,
+ timeout=10.0,
+ model="vertex_ai/",
+ input="hello world",
+ _is_async=run_async,
+ )
+ speech: Final = await pending if run_async else pending
+ return speech.content
+
+
+@BOTH_CALL_STYLES
+@pytest.mark.parametrize(
+ ("payload", "audio"),
+ [
+ ({"audioContent": ENCODED_AUDIO}, AUDIO),
+ ({"audioContent": ENCODED_AUDIO, "timepoints": [], "audioConfig": {"audioEncoding": "MP3"}}, AUDIO),
+ ({"audioContent": ENCODED_AUDIO[:8] + "\n" + ENCODED_AUDIO[8:] + "\n"}, AUDIO),
+ ({"audioContent": ""}, b""),
+ ],
+)
+async def test_vertex_speech_returns_the_decoded_audio_content(
+ run_async: bool, synthesize_endpoint: respx.Route, payload: object, audio: bytes
+) -> None:
+ assert await _synthesized_audio(run_async, synthesize_endpoint, payload) == audio
+
+
+@BOTH_CALL_STYLES
+async def test_vertex_speech_sends_the_default_voice_request_with_the_google_token(
+ run_async: bool, synthesize_endpoint: respx.Route
+) -> None:
+ await _synthesized_audio(run_async, synthesize_endpoint, {"audioContent": ENCODED_AUDIO})
+
+ sent: Final = synthesize_endpoint.calls.last.request
+ assert (json.loads(sent.content), sent.headers["authorization"], sent.headers["x-goog-user-project"]) == (
+ {
+ "input": {"text": "hello world"},
+ "voice": {"languageCode": "en-US", "name": "en-US-Studio-O"},
+ "audioConfig": {"audioEncoding": "LINEAR16", "speakingRate": "1"},
+ },
+ "Bearer test-token",
+ "test-project",
+ )
+
+
+@BOTH_CALL_STYLES
+@pytest.mark.parametrize(
+ "body",
+ [[], [{"audioContent": ENCODED_AUDIO}], "audioContent", 7, 403, 2.5, True, None, ["vertex-secret"]],
+)
+async def test_vertex_speech_rejects_a_synthesize_body_that_is_not_an_object_without_echoing_it(
+ run_async: bool, synthesize_endpoint: respx.Route, body: object
+) -> None:
+ with pytest.raises(ValidationError) as rejection:
+ await _synthesized_audio(run_async, synthesize_endpoint, body)
+
+ rejection_text: Final = str(rejection.value)
+ assert (
+ [error["type"] for error in rejection.value.errors()],
+ "vertex-secret" in rejection_text,
+ "403" in rejection_text,
+ ) == (["dict_type"], False, False), rejection_text
+
+
+@BOTH_CALL_STYLES
+@pytest.mark.parametrize("audio_content", [None, 5, 403, 1.5, True, [ENCODED_AUDIO], ["vertex-secret"], {"data": "x"}])
+async def test_vertex_speech_rejects_audio_content_that_is_not_a_string_without_echoing_it(
+ run_async: bool, synthesize_endpoint: respx.Route, audio_content: object
+) -> None:
+ with pytest.raises(ValidationError) as rejection:
+ await _synthesized_audio(run_async, synthesize_endpoint, {"audioContent": audio_content})
+
+ rejection_text: Final = str(rejection.value)
+ assert (
+ [error["type"] for error in rejection.value.errors()],
+ "vertex-secret" in rejection_text,
+ "403" in rejection_text,
+ ) == (["string_type"], False, False), rejection_text
+
+
+@BOTH_CALL_STYLES
+@pytest.mark.parametrize(
+ "payload",
+ [{}, {"audiocontent": ENCODED_AUDIO}, {"error": {"code": 403, "message": "denied", "status": "PERMISSION_DENIED"}}],
+)
+async def test_vertex_speech_reports_a_synthesize_body_without_audio_content_as_a_key_error(
+ run_async: bool, synthesize_endpoint: respx.Route, payload: object
+) -> None:
+ with pytest.raises(KeyError, match="audioContent"):
+ await _synthesized_audio(run_async, synthesize_endpoint, payload)
+
+
+@BOTH_CALL_STYLES
+async def test_vertex_speech_reports_badly_padded_audio_content_with_the_base64_error(
+ run_async: bool, synthesize_endpoint: respx.Route
+) -> None:
+ with pytest.raises(binascii.Error, match="Incorrect padding"):
+ await _synthesized_audio(run_async, synthesize_endpoint, {"audioContent": "abc"})
+
+
+@BOTH_CALL_STYLES
+async def test_vertex_speech_reports_non_ascii_audio_content_with_the_base64_error(
+ run_async: bool, synthesize_endpoint: respx.Route
+) -> None:
+ with pytest.raises(ValueError, match="string argument should contain only ASCII characters"):
+ await _synthesized_audio(run_async, synthesize_endpoint, {"audioContent": "é中"})
diff --git a/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py
index 8ca0a4a4335..dc6f35fcfe5 100644
--- a/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py
+++ b/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py
@@ -11,6 +11,7 @@ import responses
from pydantic import ValidationError
import litellm
+from litellm.litellm_core_utils.exception_mapping_utils import exception_type
from litellm.llms.vertex_ai.text_to_speech.transformation import (
VertexAILyriaTextToSpeechConfig,
VertexAITextToSpeechConfig,
@@ -761,3 +762,170 @@ async def test_aspeech_vertex_ai_sends_ssml_input(
"voice": _UK_VOICE,
"audioConfig": _UK_AUDIO_CONFIG,
}
+
+
+_WAV_AUDIO_B64: Final = "UklGRiQAAABXQVZFZm10IA=="
+_WAV_AUDIO: Final = b"RIFF$\x00\x00\x00WAVEfmt "
+_UNLABELED_AUDIO_B64: Final = "bHlyaWEtMy1hdWRpbw=="
+_UNLABELED_AUDIO: Final = b"lyria-3-audio"
+
+
+@pytest.mark.parametrize(
+ ("model", "payload", "expected_audio", "expected_mime_type"),
+ [
+ pytest.param(
+ "lyria-002",
+ {
+ "predictions": [
+ {"audioContent": None, "bytesBase64Encoded": _WAV_AUDIO_B64, "mimeType": None},
+ "trailing entry that is never read",
+ ],
+ "deployedModelId": "123",
+ },
+ _WAV_AUDIO,
+ "audio/wav",
+ id="predict-null-fields-and-unread-trailing-prediction",
+ ),
+ pytest.param(
+ "lyria-002",
+ {"predictions": [{"audioContent": _WAV_AUDIO_B64, "mimeType": ""}]},
+ _WAV_AUDIO,
+ "audio/wav",
+ id="predict-empty-mime-type-is-sniffed",
+ ),
+ pytest.param(
+ "lyria-002",
+ {"predictions": [{"audioContent": _WAV_AUDIO_B64, "mimeType": "audio/x-custom"}]},
+ _WAV_AUDIO,
+ "audio/x-custom",
+ id="predict-declared-mime-type-wins-over-sniffing",
+ ),
+ pytest.param(
+ "lyria-002",
+ {"predictions": [{"audioContent": "", "bytesBase64Encoded": ""}]},
+ b"",
+ "audio/mpeg",
+ id="predict-empty-audio-strings-yield-empty-audio",
+ ),
+ pytest.param(
+ "lyria-3-pro-preview",
+ {
+ "steps": [
+ {"type": "thought", "content": "free text that is never read as content parts"},
+ {"type": "model_output", "content": None},
+ {"type": "model_output", "content": [{"type": "audio", "data": _WAV_AUDIO_B64, "mime_type": None}]},
+ ]
+ },
+ _WAV_AUDIO,
+ "audio/wav",
+ id="interactions-unread-step-content-and-null-fields",
+ ),
+ pytest.param(
+ "lyria-3-pro-preview",
+ {"steps": None, "outputs": [{"type": "audio", "data": _UNLABELED_AUDIO_B64, "mime_type": "audio/mpeg"}]},
+ _UNLABELED_AUDIO,
+ "audio/mpeg",
+ id="interactions-null-steps-fall-back-to-outputs",
+ ),
+ pytest.param(
+ "lyria-3-clip-preview",
+ {
+ "outputs": [
+ {"type": "audio", "data": 7, "mime_type": 5},
+ {"type": "audio", "data": _UNLABELED_AUDIO_B64, "mime_type": "audio/mpeg"},
+ {"type": "audio", "data": ""},
+ ]
+ },
+ _UNLABELED_AUDIO,
+ "audio/mpeg",
+ id="interactions-last-non-empty-audio-part-wins-over-an-earlier-malformed-one",
+ ),
+ ],
+)
+def test_lyria_transform_response_tolerates_optional_and_unread_provider_fields(
+ model: str, payload: dict[str, object], expected_audio: bytes, expected_mime_type: str
+):
+ response: Final = VertexAILyriaTextToSpeechConfig().transform_text_to_speech_response(
+ model=model,
+ raw_response=httpx.Response(200, json=payload),
+ logging_obj=MagicMock(),
+ )
+
+ assert response.content == expected_audio
+ assert response.response.headers["content-type"] == expected_mime_type
+
+
+@pytest.mark.parametrize(
+ ("model", "payload"),
+ [
+ ("lyria-002", {}),
+ ("lyria-002", {"predictions": None}),
+ ("lyria-002", {"predictions": []}),
+ ("lyria-002", {"predictions": [{}]}),
+ ("lyria-002", {"predictions": [{"audioContent": "", "bytesBase64Encoded": None}]}),
+ ("lyria-3-pro-preview", {}),
+ ("lyria-3-pro-preview", {"steps": [], "outputs": None}),
+ ("lyria-3-pro-preview", {"steps": [{"type": "model_output", "content": []}]}),
+ ("lyria-3-pro-preview", {"outputs": [{"type": "audio", "data": ""}, {"type": "text", "data": "ignored"}]}),
+ ],
+)
+def test_lyria_transform_response_without_audio_reports_it_missing(model: str, payload: dict[str, object]):
+ with pytest.raises(ValueError, match=f"No generated audio found in Vertex AI {model} response"):
+ VertexAILyriaTextToSpeechConfig().transform_text_to_speech_response(
+ model=model,
+ raw_response=httpx.Response(200, json=payload),
+ logging_obj=MagicMock(),
+ )
+
+
+@pytest.mark.parametrize(
+ ("model", "payload"),
+ [
+ ("lyria-002", ["not", "an", "object"]),
+ ("lyria-002", {"predictions": {"audioContent": _WAV_AUDIO_B64}}),
+ ("lyria-002", {"predictions": ["not-an-object"]}),
+ ("lyria-002", {"predictions": [{"audioContent": 7}]}),
+ ("lyria-002", {"predictions": [{"audioContent": _WAV_AUDIO_B64, "mimeType": ["audio/wav"]}]}),
+ ("lyria-3-pro-preview", ["not", "an", "object"]),
+ ("lyria-3-pro-preview", {"steps": ["not-an-object"]}),
+ ("lyria-3-pro-preview", {"steps": [{"type": "model_output", "content": ["not-an-object"]}]}),
+ ("lyria-3-pro-preview", {"outputs": [{"type": "audio", "data": [_WAV_AUDIO_B64]}]}),
+ ("lyria-3-pro-preview", {"outputs": [{"type": "audio", "data": _WAV_AUDIO_B64, "mime_type": 5}]}),
+ ],
+)
+def test_lyria_transform_response_rejects_malformed_payloads_without_echoing_them(model: str, payload: object):
+ with pytest.raises(ValidationError) as exc_info:
+ VertexAILyriaTextToSpeechConfig().transform_text_to_speech_response(
+ model=model,
+ raw_response=httpx.Response(200, json=payload),
+ logging_obj=MagicMock(),
+ )
+
+ assert "input_value" not in str(exc_info.value)
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ pytest.param({"steps": [{"type": "thought"}] * 403 + ["not-an-object"]}, id="step"),
+ pytest.param({"outputs": [{"type": "text"}] * 403 + [None]}, id="output"),
+ pytest.param(
+ {"steps": [{"type": "model_output", "content": [{"type": "text"}] * 403 + [7]}]},
+ id="content-part",
+ ),
+ ],
+)
+def test_lyria_malformed_item_position_does_not_change_the_mapped_error(payload: dict[str, object]):
+ with pytest.raises(ValidationError) as exc_info:
+ VertexAILyriaTextToSpeechConfig().transform_text_to_speech_response(
+ model="lyria-3-pro-preview",
+ raw_response=httpx.Response(200, json=payload),
+ logging_obj=MagicMock(),
+ )
+
+ with pytest.raises(litellm.APIConnectionError):
+ exception_type(
+ model="vertex_ai/lyria-3-pro-preview",
+ custom_llm_provider="vertex_ai",
+ original_exception=exc_info.value,
+ )
diff --git a/tests/unit/llms/vertex_ai/vertex_embeddings/__init__.py b/tests/unit/llms/vertex_ai/vertex_embeddings/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/vertex_ai/vertex_embeddings/test_embedding_handler.py b/tests/unit/llms/vertex_ai/vertex_embeddings/test_embedding_handler.py
new file mode 100644
index 00000000000..4ceae8d407a
--- /dev/null
+++ b/tests/unit/llms/vertex_ai/vertex_embeddings/test_embedding_handler.py
@@ -0,0 +1,289 @@
+import json
+from datetime import datetime
+from typing import Final
+
+import google.auth
+import google.auth.credentials
+import httpx
+import pytest
+import respx
+from pydantic import ValidationError
+
+import litellm
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.vertex_ai.vertex_embeddings.embedding_handler import VertexEmbedding
+from litellm.types.utils import EmbeddingResponse
+
+GECKO: Final = "textembedding-gecko"
+FINE_TUNED_ENDPOINT: Final = "1234567890"
+BGE: Final = "bge-small-en-v1.5"
+PUBLIC_CALL_PROJECT: Final = "vertex-embedding-handler-tests"
+BOTH_CALL_STYLES: Final = pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"])
+EVERY_MODEL_FAMILY: Final = pytest.mark.parametrize("model", [GECKO, FINE_TUNED_ENDPOINT, BGE])
+BODIES_THAT_ARE_NOT_OBJECTS: Final = [
+ [],
+ [{"predictions": []}],
+ ["predictions"],
+ "predictions",
+ "",
+ 7,
+ 403,
+ 2.5,
+ True,
+ None,
+ ["vertex-secret"],
+]
+
+
+class _StaticGoogleCredentials(google.auth.credentials.Credentials):
+ def refresh(self, request: object) -> None:
+ self.token = "test-token"
+
+
+@pytest.fixture
+def predict_endpoint(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> respx.Route:
+ monkeypatch.setattr(google.auth, "default", lambda scopes: (_StaticGoogleCredentials(), "test-project"))
+ monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+ return respx_mock.post(host="us-central1-aiplatform.googleapis.com")
+
+
+def _predict_response(payload: object) -> httpx.Response:
+ return httpx.Response(200, content=json.dumps(payload).encode())
+
+
+def _embedding_logging() -> Logging:
+ return Logging(
+ model=GECKO,
+ messages=[],
+ stream=False,
+ call_type="embedding",
+ start_time=datetime(2026, 1, 1),
+ litellm_call_id="embedding-call",
+ function_id="embedding-function",
+ )
+
+
+def _gecko_prediction(values: object, token_count: int) -> dict[str, object]:
+ return {"embeddings": {"values": values, "statistics": {"token_count": token_count, "truncated": False}}}
+
+
+def _embedding_rows(vectors: list[object]) -> list[dict[str, object]]:
+ return [{"object": "embedding", "index": index, "embedding": vector} for index, vector in enumerate(vectors)]
+
+
+async def _embedding_response(
+ run_async: bool, endpoint: respx.Route, model: str, payload: object, logging_obj: Logging
+) -> EmbeddingResponse:
+ endpoint.mock(return_value=_predict_response(payload))
+ pending: Final = VertexEmbedding().embedding(
+ model=model,
+ input=["hello", "world"],
+ print_verbose=None,
+ model_response=EmbeddingResponse(),
+ optional_params={},
+ logging_obj=logging_obj,
+ custom_llm_provider="vertex_ai",
+ timeout=10.0,
+ aembedding=run_async,
+ vertex_project="test-project",
+ vertex_location="us-central1",
+ vertex_credentials=None,
+ litellm_params={},
+ )
+ return await pending if run_async else pending
+
+
+@BOTH_CALL_STYLES
+@pytest.mark.parametrize(
+ ("model", "payload", "prompt_tokens"),
+ [
+ (
+ GECKO,
+ {
+ "predictions": [_gecko_prediction([0.1, 0.2, 0.3], 3), _gecko_prediction([1, 2.0, 3], 5)],
+ "metadata": {"billableCharacterCount": 10},
+ },
+ 8,
+ ),
+ (FINE_TUNED_ENDPOINT, {"predictions": [[[0.1, 0.2, 0.3]], [[1, 2.0, 3]]], "deployedModelId": "42"}, 0),
+ (BGE, {"predictions": [[0.1, 0.2, 0.3], [1, 2.0, 3]], "deployedModelId": "42"}, 0),
+ ],
+)
+async def test_vertex_embedding_returns_each_prediction_vector_exactly_as_sent(
+ run_async: bool, predict_endpoint: respx.Route, model: str, payload: object, prompt_tokens: int
+) -> None:
+ response: Final = await _embedding_response(run_async, predict_endpoint, model, payload, _embedding_logging())
+
+ assert (
+ json.dumps(response.data),
+ response.object,
+ response.model,
+ response.usage.prompt_tokens,
+ response.usage.completion_tokens,
+ response.usage.total_tokens,
+ ) == (
+ json.dumps(_embedding_rows([[0.1, 0.2, 0.3], [1, 2.0, 3]])),
+ "list",
+ model,
+ prompt_tokens,
+ 0,
+ prompt_tokens,
+ )
+
+
+@BOTH_CALL_STYLES
+@pytest.mark.parametrize("values", [[], "not-a-vector", None, [0.5, "x", None, True, {"a": 1}]])
+async def test_vertex_embedding_keeps_prediction_values_that_are_not_float_vectors(
+ run_async: bool, predict_endpoint: respx.Route, values: object
+) -> None:
+ response: Final = await _embedding_response(
+ run_async, predict_endpoint, GECKO, {"predictions": [_gecko_prediction(values, 3)]}, _embedding_logging()
+ )
+
+ assert json.dumps(response.data) == json.dumps(_embedding_rows([values]))
+
+
+@BOTH_CALL_STYLES
+async def test_vertex_embedding_returns_every_vector_of_a_large_batch_in_order(
+ run_async: bool, predict_endpoint: respx.Route
+) -> None:
+ vectors: Final[list[object]] = [[index, index + 0.5] for index in range(500)]
+
+ response: Final = await _embedding_response(
+ run_async, predict_endpoint, BGE, {"predictions": vectors}, _embedding_logging()
+ )
+
+ assert json.dumps(response.data) == json.dumps(_embedding_rows(vectors))
+
+
+@BOTH_CALL_STYLES
+async def test_vertex_embedding_logs_the_predict_body_as_json_text(
+ run_async: bool, predict_endpoint: respx.Route
+) -> None:
+ payload: Final = {"predictions": [[0.1, 0.2]], "deployedModelId": "42"}
+ logging_obj: Final = _embedding_logging()
+
+ await _embedding_response(run_async, predict_endpoint, BGE, payload, logging_obj)
+
+ assert logging_obj.model_call_details["original_response"] == json.dumps(payload)
+
+
+@BOTH_CALL_STYLES
+@EVERY_MODEL_FAMILY
+@pytest.mark.parametrize("body", BODIES_THAT_ARE_NOT_OBJECTS)
+async def test_vertex_embedding_rejects_a_predict_body_that_is_not_an_object_without_echoing_it(
+ run_async: bool, predict_endpoint: respx.Route, model: str, body: object
+) -> None:
+ with pytest.raises(ValidationError) as rejection:
+ await _embedding_response(run_async, predict_endpoint, model, body, _embedding_logging())
+
+ rejection_text: Final = str(rejection.value)
+ assert (
+ [error["type"] for error in rejection.value.errors()],
+ "vertex-secret" in rejection_text,
+ "403" in rejection_text,
+ ) == (["dict_type"], False, False), rejection_text
+
+
+@BOTH_CALL_STYLES
+async def test_vertex_embedding_logs_a_predict_body_that_is_not_an_object_before_rejecting_it(
+ run_async: bool, predict_endpoint: respx.Route
+) -> None:
+ logging_obj: Final = _embedding_logging()
+
+ with pytest.raises(ValidationError):
+ await _embedding_response(run_async, predict_endpoint, GECKO, ["not", "an", "object"], logging_obj)
+
+ assert logging_obj.model_call_details["original_response"] == ["not", "an", "object"]
+
+
+@BOTH_CALL_STYLES
+@EVERY_MODEL_FAMILY
+@pytest.mark.parametrize(
+ "payload",
+ [{}, {"error": {"code": 403, "message": "denied", "status": "PERMISSION_DENIED"}}],
+)
+async def test_vertex_embedding_reports_a_predict_object_without_predictions_as_a_key_error(
+ run_async: bool, predict_endpoint: respx.Route, model: str, payload: object
+) -> None:
+ with pytest.raises(KeyError, match="predictions"):
+ await _embedding_response(run_async, predict_endpoint, model, payload, _embedding_logging())
+
+
+def test_vertex_embedding_call_returns_the_prediction_vectors(predict_endpoint: respx.Route) -> None:
+ predict_endpoint.mock(
+ return_value=_predict_response({"predictions": [_gecko_prediction([0.1, 2], 3), _gecko_prediction([3.5], 4)]})
+ )
+
+ response: Final = litellm.embedding(
+ model=f"vertex_ai/{GECKO}",
+ input=["hello", "world"],
+ vertex_project=PUBLIC_CALL_PROJECT,
+ vertex_location="us-central1",
+ )
+
+ assert (json.dumps(response.data), response.usage.prompt_tokens) == (
+ json.dumps(_embedding_rows([[0.1, 2], [3.5]])),
+ 7,
+ )
+
+
+async def test_vertex_async_embedding_call_returns_the_prediction_vectors(predict_endpoint: respx.Route) -> None:
+ predict_endpoint.mock(
+ return_value=_predict_response({"predictions": [_gecko_prediction([0.1, 2], 3), _gecko_prediction([3.5], 4)]})
+ )
+
+ response: Final = await litellm.aembedding(
+ model=f"vertex_ai/{GECKO}",
+ input=["hello", "world"],
+ vertex_project=PUBLIC_CALL_PROJECT,
+ vertex_location="us-central1",
+ )
+
+ assert (json.dumps(response.data), response.usage.prompt_tokens) == (
+ json.dumps(_embedding_rows([[0.1, 2], [3.5]])),
+ 7,
+ )
+
+
+MALFORMED_PREDICT_BODIES: Final = [
+ ["vertex-secret"],
+ "predictions",
+ 403,
+ None,
+ {"predictions": [_gecko_prediction([0.5], 1) for _ in range(403)] + [None]},
+]
+
+
+@pytest.mark.parametrize("body", MALFORMED_PREDICT_BODIES)
+def test_vertex_embedding_call_maps_a_malformed_predict_body_to_a_connection_error_without_echoing_it(
+ predict_endpoint: respx.Route, body: object
+) -> None:
+ predict_endpoint.mock(return_value=_predict_response(body))
+
+ with pytest.raises(litellm.APIConnectionError) as failure:
+ litellm.embedding(
+ model=f"vertex_ai/{GECKO}",
+ input=["hello", "world"],
+ vertex_project=PUBLIC_CALL_PROJECT,
+ vertex_location="us-central1",
+ )
+
+ assert (failure.value.status_code, "vertex-secret" in str(failure.value)) == (500, False)
+
+
+@pytest.mark.parametrize("body", MALFORMED_PREDICT_BODIES)
+async def test_vertex_async_embedding_call_maps_a_malformed_predict_body_to_a_connection_error_without_echoing_it(
+ predict_endpoint: respx.Route, body: object
+) -> None:
+ predict_endpoint.mock(return_value=_predict_response(body))
+
+ with pytest.raises(litellm.APIConnectionError) as failure:
+ await litellm.aembedding(
+ model=f"vertex_ai/{GECKO}",
+ input=["hello", "world"],
+ vertex_project=PUBLIC_CALL_PROJECT,
+ vertex_location="us-central1",
+ )
+
+ assert (failure.value.status_code, "vertex-secret" in str(failure.value)) == (500, False)
diff --git a/tests/unit/llms/voyage/embedding/test_transformation_contextual.py b/tests/unit/llms/voyage/embedding/test_transformation_contextual.py
index 9fd630283bd..915fe1ea48d 100644
--- a/tests/unit/llms/voyage/embedding/test_transformation_contextual.py
+++ b/tests/unit/llms/voyage/embedding/test_transformation_contextual.py
@@ -1,12 +1,18 @@
import asyncio
+import datetime
import importlib
import json
import os
from unittest.mock import MagicMock, patch
+import httpx
import pytest
+from pydantic import ValidationError
import litellm
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.voyage.embedding.transformation_contextual import VoyageContextualEmbeddingConfig
+from litellm.types.utils import EmbeddingResponse
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@@ -473,3 +479,86 @@ _SCALAR_DEFAULTS = {
"api_key": getattr(litellm, "api_key", None),
"cohere_key": getattr(litellm, "cohere_key", None),
}
+
+
+EMBEDDING_ROW = {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}
+
+
+def _transform(payload: object) -> EmbeddingResponse:
+ logging_obj = Logging(
+ model="voyage/voyage-context-3",
+ messages=[],
+ stream=False,
+ call_type="embedding",
+ start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+ litellm_call_id="call-id",
+ function_id="function-id",
+ )
+ return VoyageContextualEmbeddingConfig().transform_embedding_response(
+ model="voyage-context-3",
+ raw_response=httpx.Response(200, json=payload),
+ model_response=EmbeddingResponse(),
+ logging_obj=logging_obj,
+ api_key="sk-test",
+ request_data={},
+ optional_params={},
+ litellm_params={},
+ )
+
+
+def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]:
+ assert response.usage is not None
+ return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens)
+
+
+def test_voyage_contextual_embedding_response_carries_provider_fields_and_usage():
+ response = _transform(
+ {
+ "object": "list",
+ "model": "voyage-context-3",
+ "data": [EMBEDDING_ROW],
+ "usage": {"prompt_tokens": 8, "total_tokens": 9},
+ }
+ )
+
+ assert response.model == "voyage-context-3"
+ assert response.object == "list"
+ assert response.data == [EMBEDDING_ROW]
+ assert _token_counts(response) == (9, 0, 9)
+
+
+@pytest.mark.parametrize(
+ ("payload", "expected"),
+ [
+ ({"data": [EMBEDDING_ROW]}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 4}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 6}}, (6, 0, 6)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": None}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": ""}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": []}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": "12"}}, (12, 0, 12)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 12.0}}, (12, 0, 12)),
+ (
+ {"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 3, "total_tokens": 5, "completion_tokens": 2}},
+ (5, 0, 5),
+ ),
+ ],
+)
+def test_voyage_contextual_embedding_usage_tolerates_missing_null_and_loosely_typed_counts(payload, expected):
+ response = _transform(payload)
+
+ assert response.data == [EMBEDDING_ROW]
+ assert _token_counts(response) == expected
+
+
+@pytest.mark.parametrize("usage", [None, [], "tokens", 3])
+def test_voyage_contextual_embedding_usage_that_is_not_an_object_is_rejected(usage):
+ with pytest.raises(ValidationError):
+ _transform({"data": [EMBEDDING_ROW], "usage": usage})
+
+
+@pytest.mark.parametrize("usage", [{"total_tokens": "many"}, {"total_tokens": 1.5}, {"total_tokens": [1]}])
+def test_voyage_contextual_embedding_token_count_that_is_not_an_integer_is_rejected(usage):
+ with pytest.raises(ValidationError):
+ _transform({"data": [EMBEDDING_ROW], "usage": usage})
diff --git a/tests/unit/llms/voyage/embedding/test_transformation_multimodal.py b/tests/unit/llms/voyage/embedding/test_transformation_multimodal.py
new file mode 100644
index 00000000000..6be8e313a84
--- /dev/null
+++ b/tests/unit/llms/voyage/embedding/test_transformation_multimodal.py
@@ -0,0 +1,91 @@
+import datetime
+
+import httpx
+import pytest
+from pydantic import ValidationError
+
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.voyage.embedding.transformation_multimodal import VoyageMultimodalEmbeddingConfig
+from litellm.types.utils import EmbeddingResponse
+
+EMBEDDING_ROW = {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}
+
+
+def _transform(payload: object) -> EmbeddingResponse:
+ logging_obj = Logging(
+ model="voyage/voyage-multimodal-3",
+ messages=[],
+ stream=False,
+ call_type="embedding",
+ start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+ litellm_call_id="call-id",
+ function_id="function-id",
+ )
+ return VoyageMultimodalEmbeddingConfig().transform_embedding_response(
+ model="voyage-multimodal-3",
+ raw_response=httpx.Response(200, json=payload),
+ model_response=EmbeddingResponse(),
+ logging_obj=logging_obj,
+ api_key="sk-test",
+ request_data={},
+ optional_params={},
+ litellm_params={},
+ )
+
+
+def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]:
+ assert response.usage is not None
+ return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens)
+
+
+def test_voyage_multimodal_embedding_response_carries_provider_fields_and_usage():
+ response = _transform(
+ {
+ "object": "list",
+ "model": "voyage-multimodal-3",
+ "data": [EMBEDDING_ROW],
+ "usage": {"prompt_tokens": 8, "total_tokens": 9},
+ }
+ )
+
+ assert response.model == "voyage-multimodal-3"
+ assert response.object == "list"
+ assert response.data == [EMBEDDING_ROW]
+ assert _token_counts(response) == (9, 0, 9)
+
+
+@pytest.mark.parametrize(
+ ("payload", "expected"),
+ [
+ ({"data": [EMBEDDING_ROW]}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 4}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 6}}, (6, 0, 6)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": None}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": ""}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": []}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": "12"}}, (12, 0, 12)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 12.0}}, (12, 0, 12)),
+ (
+ {"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 3, "total_tokens": 5, "completion_tokens": 2}},
+ (5, 0, 5),
+ ),
+ ],
+)
+def test_voyage_multimodal_embedding_usage_tolerates_missing_null_and_loosely_typed_counts(payload, expected):
+ response = _transform(payload)
+
+ assert response.data == [EMBEDDING_ROW]
+ assert _token_counts(response) == expected
+
+
+@pytest.mark.parametrize("usage", [None, [], "tokens", 3])
+def test_voyage_multimodal_embedding_usage_that_is_not_an_object_is_rejected(usage):
+ with pytest.raises(ValidationError):
+ _transform({"data": [EMBEDDING_ROW], "usage": usage})
+
+
+@pytest.mark.parametrize("usage", [{"total_tokens": "many"}, {"total_tokens": 1.5}, {"total_tokens": [1]}])
+def test_voyage_multimodal_embedding_token_count_that_is_not_an_integer_is_rejected(usage):
+ with pytest.raises(ValidationError):
+ _transform({"data": [EMBEDDING_ROW], "usage": usage})
diff --git a/tests/unit/llms/voyage/embedding/test_voyage_embedding_transformation.py b/tests/unit/llms/voyage/embedding/test_voyage_embedding_transformation.py
new file mode 100644
index 00000000000..41bc812ec4a
--- /dev/null
+++ b/tests/unit/llms/voyage/embedding/test_voyage_embedding_transformation.py
@@ -0,0 +1,91 @@
+import datetime
+
+import httpx
+import pytest
+from pydantic import ValidationError
+
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.voyage.embedding.transformation import VoyageEmbeddingConfig
+from litellm.types.utils import EmbeddingResponse
+
+EMBEDDING_ROW = {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}
+
+
+def _transform(payload: object) -> EmbeddingResponse:
+ logging_obj = Logging(
+ model="voyage/voyage-3",
+ messages=[],
+ stream=False,
+ call_type="embedding",
+ start_time=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
+ litellm_call_id="call-id",
+ function_id="function-id",
+ )
+ return VoyageEmbeddingConfig().transform_embedding_response(
+ model="voyage-3",
+ raw_response=httpx.Response(200, json=payload),
+ model_response=EmbeddingResponse(),
+ logging_obj=logging_obj,
+ api_key="sk-test",
+ request_data={},
+ optional_params={},
+ litellm_params={},
+ )
+
+
+def _token_counts(response: EmbeddingResponse) -> tuple[int, int, int]:
+ assert response.usage is not None
+ return (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens)
+
+
+def test_voyage_embedding_response_carries_provider_fields_and_usage():
+ response = _transform(
+ {
+ "object": "list",
+ "model": "voyage-3",
+ "data": [EMBEDDING_ROW],
+ "usage": {"prompt_tokens": 8, "total_tokens": 9},
+ }
+ )
+
+ assert response.model == "voyage-3"
+ assert response.object == "list"
+ assert response.data == [EMBEDDING_ROW]
+ assert _token_counts(response) == (9, 0, 9)
+
+
+@pytest.mark.parametrize(
+ ("payload", "expected"),
+ [
+ ({"data": [EMBEDDING_ROW]}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 4}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 6}}, (6, 0, 6)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": None}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": ""}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": []}}, (0, 0, 0)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": "12"}}, (12, 0, 12)),
+ ({"data": [EMBEDDING_ROW], "usage": {"total_tokens": 12.0}}, (12, 0, 12)),
+ (
+ {"data": [EMBEDDING_ROW], "usage": {"prompt_tokens": 3, "total_tokens": 5, "completion_tokens": 2}},
+ (5, 0, 5),
+ ),
+ ],
+)
+def test_voyage_embedding_usage_tolerates_missing_null_and_loosely_typed_counts(payload, expected):
+ response = _transform(payload)
+
+ assert response.data == [EMBEDDING_ROW]
+ assert _token_counts(response) == expected
+
+
+@pytest.mark.parametrize("usage", [None, [], "tokens", 3])
+def test_voyage_embedding_usage_that_is_not_an_object_is_rejected(usage):
+ with pytest.raises(ValidationError):
+ _transform({"data": [EMBEDDING_ROW], "usage": usage})
+
+
+@pytest.mark.parametrize("usage", [{"total_tokens": "many"}, {"total_tokens": 1.5}, {"total_tokens": [1]}])
+def test_voyage_embedding_token_count_that_is_not_an_integer_is_rejected(usage):
+ with pytest.raises(ValidationError):
+ _transform({"data": [EMBEDDING_ROW], "usage": usage})
diff --git a/tests/unit/llms/watsonx/completion/__init__.py b/tests/unit/llms/watsonx/completion/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/watsonx/completion/test_watsonx_completion_transformation.py b/tests/unit/llms/watsonx/completion/test_watsonx_completion_transformation.py
new file mode 100644
index 00000000000..4ec197dcd46
--- /dev/null
+++ b/tests/unit/llms/watsonx/completion/test_watsonx_completion_transformation.py
@@ -0,0 +1,60 @@
+import pytest
+
+from litellm.llms.watsonx.completion.transformation import IBMWatsonXAIConfig
+
+
+@pytest.mark.parametrize(
+ ("supplied", "stored"),
+ [
+ pytest.param({}, {"kwargs": {}}, id="nothing-supplied"),
+ pytest.param(
+ {"decoding_method": "", "max_new_tokens": 0, "include_stop_sequences": False, "stream": False},
+ {
+ "decoding_method": "",
+ "max_new_tokens": 0,
+ "include_stop_sequences": False,
+ "stream": False,
+ "kwargs": {},
+ },
+ id="falsy-values-are-kept",
+ ),
+ pytest.param(
+ {"top_k": 40, "top_p": None, "random_seed": None}, {"top_k": 40, "kwargs": {}}, id="none-is-skipped"
+ ),
+ pytest.param(
+ {"temperature": 0.2, "project_id": "proj-1", "space_id": None},
+ {"temperature": 0.2, "kwargs": {"project_id": "proj-1", "space_id": None}},
+ id="unknown-keywords-are-grouped-under-kwargs",
+ ),
+ ],
+)
+def test_init_stores_each_supplied_value_except_none_as_class_level_config(
+ supplied: dict[str, object], stored: dict[str, object]
+) -> None:
+ class ScopedWatsonXConfig(IBMWatsonXAIConfig):
+ pass
+
+ ScopedWatsonXConfig(**supplied)
+
+ assert ScopedWatsonXConfig.get_config() == stored
+
+
+def test_init_stores_the_supplied_object_itself_rather_than_a_copy() -> None:
+ class ScopedWatsonXConfig(IBMWatsonXAIConfig):
+ pass
+
+ stop_sequences = ["}", ")"]
+
+ ScopedWatsonXConfig(stop_sequences=stop_sequences)
+
+ assert ScopedWatsonXConfig.get_config()["stop_sequences"] is stop_sequences
+
+
+def test_init_called_again_with_none_keeps_the_value_stored_earlier() -> None:
+ class ScopedWatsonXConfig(IBMWatsonXAIConfig):
+ pass
+
+ ScopedWatsonXConfig(top_k=40)
+ ScopedWatsonXConfig(top_k=None, top_p=0.5)
+
+ assert ScopedWatsonXConfig.get_config() == {"top_k": 40, "top_p": 0.5, "kwargs": {}}
diff --git a/tests/unit/llms/xai/test_xai_common_utils.py b/tests/unit/llms/xai/test_xai_common_utils.py
new file mode 100644
index 00000000000..42a61062229
--- /dev/null
+++ b/tests/unit/llms/xai/test_xai_common_utils.py
@@ -0,0 +1,124 @@
+import re
+
+import httpx
+import pytest
+import respx
+from pydantic import ValidationError
+
+import litellm
+from litellm.llms.xai.common_utils import XAIModelInfo
+
+MODELS_URL = "https://api.x.ai/v1/models"
+MANY_MODELS = [{"id": f"grok-{index}"} for index in range(500)]
+
+
+@pytest.fixture(autouse=True)
+def _default_xai_api_base(monkeypatch: pytest.MonkeyPatch):
+ monkeypatch.delenv("XAI_API_BASE", raising=False)
+
+
+def _get_models(payload: object) -> list[str]:
+ with respx.mock:
+ respx.get(MODELS_URL).mock(return_value=httpx.Response(200, json=payload))
+ return XAIModelInfo().get_models(api_key="xai-test-key")
+
+
+def test_get_models_prefixes_every_listed_model_id_in_provider_order():
+ names = _get_models(
+ {
+ "object": "list",
+ "data": [
+ {"id": "grok-4", "object": "model", "created": 1752019200, "owned_by": "xai"},
+ {"id": "grok-3-mini", "object": "model", "created": 1743724800, "owned_by": "xai"},
+ {"id": ""},
+ ],
+ }
+ )
+
+ assert names == ["xai/grok-4", "xai/grok-3-mini", "xai/"]
+ assert [type(name) for name in names] == [str, str, str]
+
+
+@pytest.mark.parametrize("data", [[], "", {}])
+def test_get_models_returns_no_names_for_an_empty_data_container(data: object):
+ assert _get_models({"data": data}) == []
+
+
+def test_get_models_sends_the_api_key_as_bearer_token():
+ with respx.mock:
+ route = respx.get(MODELS_URL).mock(return_value=httpx.Response(200, json={"data": [{"id": "grok-4"}]}))
+ XAIModelInfo().get_models(api_key="xai-test-key")
+
+ assert route.calls[0].request.headers["authorization"] == "Bearer xai-test-key"
+
+
+@pytest.mark.parametrize(
+ ("payload", "missing_key"),
+ [
+ ({"object": "list"}, "data"),
+ ({"data": [{"object": "model"}]}, "id"),
+ ({"data": [*MANY_MODELS[:403], {"object": "model"}, *MANY_MODELS[404:]]}, "id"),
+ ],
+)
+def test_get_models_raises_key_error_naming_the_field_the_provider_left_out(payload: object, missing_key: str):
+ with pytest.raises(KeyError) as exc_info:
+ _get_models(payload)
+
+ assert exc_info.value.args == (missing_key,)
+
+
+@pytest.mark.parametrize(
+ ("payload", "error_type"),
+ [
+ ([{"id": "grok-4"}], "dict_type"),
+ ("data", "dict_type"),
+ ({"data": None}, "iterable_type"),
+ ({"data": 5}, "iterable_type"),
+ ({"data": "grok-4"}, "dict_type"),
+ ({"data": {"grok-4": {"id": "grok-4"}}}, "dict_type"),
+ ({"data": ["grok-4"]}, "dict_type"),
+ ({"data": [None]}, "dict_type"),
+ ({"data": [{"id": None}]}, "string_type"),
+ ({"data": [{"id": 12}]}, "string_type"),
+ ({"data": [{"id": ["grok-4"]}]}, "string_type"),
+ ({"data": [*MANY_MODELS[:403], "grok-403", *MANY_MODELS[404:]]}, "dict_type"),
+ ({"data": [*MANY_MODELS[:429], {"id": 429}, *MANY_MODELS[430:]]}, "string_type"),
+ ],
+)
+def test_get_models_rejects_a_malformed_listing_without_echoing_position_or_input(payload: object, error_type: str):
+ with pytest.raises(ValidationError) as exc_info:
+ _get_models(payload)
+
+ assert [(error["type"], error["loc"]) for error in exc_info.value.errors()] == [(error_type, ())]
+ assert "403" not in str(exc_info.value)
+ assert "429" not in str(exc_info.value)
+ assert "grok" not in str(exc_info.value)
+
+
+def test_get_models_reports_the_upstream_status_and_body_on_an_http_error():
+ with respx.mock:
+ respx.get(MODELS_URL).mock(return_value=httpx.Response(401, text="bad key"))
+ with pytest.raises(
+ Exception, match=re.escape("Failed to fetch models from XAI. Status code: 401, Response: bad key")
+ ):
+ XAIModelInfo().get_models(api_key="xai-test-key")
+
+
+@pytest.mark.parametrize(
+ ("api_key", "payload", "expected"),
+ [
+ ("xai-listing-key", {"data": [{"id": "grok-4"}, {"id": "grok-3"}]}, ["xai/grok-4", "xai/grok-3"]),
+ ("xai-bad-id-key", {"data": [{"id": "grok-4"}, {"id": 7}]}, []),
+ ("xai-null-data-key", {"data": None}, []),
+ ("xai-missing-data-key", {"object": "list"}, []),
+ ("xai-list-body-key", [{"id": "grok-4"}], []),
+ ],
+)
+def test_get_valid_models_lists_xai_models_and_falls_back_to_empty_on_a_malformed_listing(
+ api_key: str, payload: object, expected: list[str]
+):
+ with respx.mock:
+ respx.get(MODELS_URL).mock(return_value=httpx.Response(200, json=payload))
+ models = litellm.get_valid_models(check_provider_endpoint=True, custom_llm_provider="xai", api_key=api_key)
+
+ assert models == expected
diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py
index a5d0d0a3ecc..d2f2fec073c 100644
--- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py
+++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py
@@ -13,8 +13,11 @@ from dataclasses import dataclass
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
+import httpx
import pytest
+import respx
+import litellm
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.agents import AgentCaller
@@ -2689,3 +2692,59 @@ def test_forwarding_headers_minted_bearer_replaces_a_forwarded_authorization_of_
)
assert merged == {"X-Custom": "kept", "Authorization": "Bearer minted-token"}
+
+
+_UPSTREAM_STATUS_ONLY_ERROR: Final = {
+ "jsonrpc": "2.0",
+ "id": "req-1",
+ "error": {"code": -32603, "message": "Not Found"},
+}
+
+
+@pytest.mark.parametrize(
+ ("upstream_body", "relayed_event"),
+ [
+ (
+ b'{"jsonrpc":"2.0","id":"up-1","error":{"code":-32001,"message":"Task not found"}}',
+ {"jsonrpc": "2.0", "id": "up-1", "error": {"code": -32001, "message": "Task not found"}},
+ ),
+ (b'{"error":"quota exceeded"}', {"error": "quota exceeded", "id": "req-1"}),
+ (b'{"jsonrpc":"2.0","id":"up-1","result":{"ok":true}}', _UPSTREAM_STATUS_ONLY_ERROR),
+ (b'["error"]', _UPSTREAM_STATUS_ONLY_ERROR),
+ (b'"error"', _UPSTREAM_STATUS_ONLY_ERROR),
+ (b"null", _UPSTREAM_STATUS_ONLY_ERROR),
+ (b"7", _UPSTREAM_STATUS_ONLY_ERROR),
+ (b"\xff\xfe", _UPSTREAM_STATUS_ONLY_ERROR),
+ (b"", _UPSTREAM_STATUS_ONLY_ERROR),
+ (b'{"error":null}', {"error": None, "id": "req-1"}),
+ ],
+)
+@pytest.mark.asyncio
+async def test_failed_upstream_stream_relays_its_body_only_when_that_is_a_json_object_with_an_error(
+ monkeypatch: pytest.MonkeyPatch,
+ respx_mock: respx.MockRouter,
+ upstream_body: bytes,
+ relayed_event: Mapping[str, object],
+):
+ from litellm.proxy.agent_endpoints.a2a_endpoints import _a2a_sse_event_source
+
+ monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+ upstream: Final = respx_mock.post("http://backend-agent:10001").mock(
+ return_value=httpx.Response(404, content=upstream_body)
+ )
+
+ events: Final = [
+ event
+ async for event in _a2a_sse_event_source(
+ "http://backend-agent:10001",
+ {"jsonrpc": "2.0", "id": "req-1", "method": "tasks/resubscribe"},
+ request_id="req-1",
+ )
+ ]
+
+ assert json.dumps(events, sort_keys=True) == json.dumps([relayed_event], sort_keys=True), upstream_body
+ assert json.loads(upstream.calls.last.request.content) == {
+ "jsonrpc": "2.0",
+ "id": "req-1",
+ "method": "tasks/resubscribe",
+ }
diff --git a/tests/unit/proxy/client/cli/test_auth_commands.py b/tests/unit/proxy/client/cli/test_auth_commands.py
index b4a6aaf706e..27edf9f3eb3 100644
--- a/tests/unit/proxy/client/cli/test_auth_commands.py
+++ b/tests/unit/proxy/client/cli/test_auth_commands.py
@@ -2075,3 +2075,94 @@ def test_login_shows_the_error_detail_only_when_the_proxy_answers_with_a_json_ob
"Authentication failed: Starting CLI login failed: HTTP 429 from https://test.example.com/sso/cli/start"
f"{expected_detail}\n"
)
+
+
+@pytest.mark.parametrize(
+ ("obj", "expected_exit_code", "expected_stdout", "expected_stderr"),
+ [
+ ({}, 0, "sk-stored\n", ""),
+ ({"base_url": "https://test.example.com///", "base_url_explicit": True}, 0, "sk-stored\n", ""),
+ (
+ {"base_url": "https://other.example.com", "base_url_explicit": True},
+ 1,
+ "",
+ "Not authenticated for this server. Run 'lite login'.\n",
+ ),
+ ({"base_url": "https://other.example.com", "base_url_explicit": False}, 0, "sk-stored\n", ""),
+ (
+ {"base_url": "https://test.example.com", "api_key": "sk-context", "api_key_from_token_file": True},
+ 0,
+ "sk-context\n",
+ "",
+ ),
+ (
+ {"base_url": "https://test.example.com", "api_key": "sk-context", "api_key_from_token_file": False},
+ 0,
+ "sk-stored\n",
+ "",
+ ),
+ (
+ {"base_url": "https://other.example.com", "api_key": "sk-context", "api_key_from_token_file": True},
+ 0,
+ "sk-stored\n",
+ "",
+ ),
+ ],
+)
+def test_print_token_reads_the_connection_from_the_context_object(
+ isolated_home, obj, expected_exit_code, expected_stdout, expected_stderr
+):
+ _write_token_file(isolated_home, key="sk-stored")
+
+ result = CliRunner().invoke(print_token, obj=obj)
+
+ assert result.exit_code == expected_exit_code
+ assert result.stdout == expected_stdout
+ assert result.stderr == expected_stderr
+
+
+def test_print_token_without_a_context_object_reports_no_login_when_nothing_is_stored(isolated_home):
+ result = CliRunner().invoke(print_token)
+
+ assert result.exit_code == 1
+ assert result.stdout == ""
+ assert result.stderr == "Not authenticated. Run 'lite login'.\n"
+
+
+def test_print_token_without_a_context_object_fails_reading_it_once_a_login_is_stored(isolated_home):
+ _write_token_file(isolated_home, key="sk-stored")
+
+ result = CliRunner().invoke(print_token)
+
+ assert result.exit_code == 1
+ assert type(result.exception) is AttributeError
+ assert result.stdout == ""
+ assert result.stderr == ""
+
+
+@pytest.mark.parametrize(
+ "obj",
+ [
+ {},
+ {"api_key": "sk-context", "base_url_explicit": True},
+ ],
+)
+@responses.activate
+def test_login_without_a_base_url_in_the_context_fails_before_any_request(obj):
+ result = CliRunner().invoke(login, obj=obj)
+
+ assert result.exit_code == 1
+ assert type(result.exception) is KeyError
+ assert result.exception.args == ("base_url",)
+ assert result.output == ""
+ assert len(responses.calls) == 0
+
+
+@responses.activate
+def test_login_without_a_context_object_fails_before_any_request():
+ result = CliRunner().invoke(login)
+
+ assert result.exit_code == 1
+ assert type(result.exception) is TypeError
+ assert result.output == ""
+ assert len(responses.calls) == 0
diff --git a/tests/unit/proxy/client/cli/test_global_options.py b/tests/unit/proxy/client/cli/test_global_options.py
index d46cc2ad120..96b200b3983 100644
--- a/tests/unit/proxy/client/cli/test_global_options.py
+++ b/tests/unit/proxy/client/cli/test_global_options.py
@@ -266,3 +266,106 @@ def test_lite_entrypoint_prints_nothing_on_stderr(monkeypatch, capsys, requests_
assert "LiteLLM Proxy Server Version: 1.2.3" in captured.out
assert f"LiteLLM Proxy CLI Version: {litellm_version}" in captured.out
assert captured.err == ""
+
+
+@pytest.mark.parametrize(
+ ("args", "config", "expected_context", "expected_authorization"),
+ [
+ (
+ ["--base-url", "https://flag-proxy.example.com/", "--api-key", "sk-flag"],
+ {},
+ {
+ "base_url": "https://flag-proxy.example.com",
+ "api_key": "sk-flag",
+ "api_key_from_token_file": False,
+ "base_url_explicit": True,
+ },
+ "Bearer sk-flag",
+ ),
+ (
+ [],
+ {},
+ {
+ "base_url": "http://localhost:4000",
+ "api_key": None,
+ "api_key_from_token_file": True,
+ "base_url_explicit": False,
+ },
+ None,
+ ),
+ (
+ ["--api-key", ""],
+ {"base_url": "https://config-proxy.example.com"},
+ {
+ "base_url": "https://config-proxy.example.com",
+ "api_key": "",
+ "api_key_from_token_file": False,
+ "base_url_explicit": True,
+ },
+ None,
+ ),
+ ],
+)
+def test_cli_stores_the_resolved_connection_on_the_supplied_context_object(
+ cli_runner, isolated_home, requests_mock, args, config, expected_context, expected_authorization
+):
+ _write_config_file(isolated_home, config)
+ requests_mock.get(f"{expected_context['base_url']}/health/readiness", json={"litellm_version": "1.2.3"})
+ context: Final[dict[str, object]] = {}
+
+ result = cli_runner.invoke(cli, [*args, "version"], obj=context)
+
+ assert result.exit_code == 0
+ assert json.dumps(context) == json.dumps(expected_context)
+ assert "LiteLLM Proxy Server Version: 1.2.3" in result.output
+ assert requests_mock.last_request.headers.get("Authorization") == expected_authorization
+
+
+@pytest.mark.parametrize(
+ ("context", "expected_lines", "expected_authorization"),
+ [
+ (
+ {"base_url": "http://localhost:4000", "api_key": "sk-context"},
+ ["LiteLLM Proxy Server URL: http://localhost:4000", "LiteLLM Proxy Server Version: 1.2.3"],
+ "Bearer sk-context",
+ ),
+ (
+ {"base_url": "http://localhost:4000"},
+ ["LiteLLM Proxy Server URL: http://localhost:4000", "LiteLLM Proxy Server Version: 1.2.3"],
+ None,
+ ),
+ (
+ {"base_url": "http://localhost:4000", "api_key": None, "api_key_from_token_file": True},
+ ["LiteLLM Proxy Server URL: http://localhost:4000", "LiteLLM Proxy Server Version: 1.2.3"],
+ None,
+ ),
+ ],
+)
+def test_version_command_reads_the_connection_from_the_context_object(
+ cli_runner, requests_mock, context, expected_lines, expected_authorization
+):
+ requests_mock.get("http://localhost:4000/health/readiness", json={"litellm_version": "1.2.3"})
+
+ result = cli_runner.invoke(cli.commands["version"], [], obj=context)
+
+ assert result.exit_code == 0
+ assert result.output.splitlines() == [f"LiteLLM Proxy CLI Version: {litellm_version}", *expected_lines]
+ assert requests_mock.last_request.headers.get("Authorization") == expected_authorization
+
+
+def test_version_command_without_a_base_url_reports_the_failed_lookup(cli_runner):
+ result = cli_runner.invoke(cli.commands["version"], [], obj={"api_key": "sk-context"})
+
+ lines: Final = result.output.splitlines()
+ assert result.exit_code == 0
+ assert len(lines) == 2
+ assert lines[0] == f"LiteLLM Proxy CLI Version: {litellm_version}"
+ assert lines[1].startswith("Could not retrieve server version: ")
+
+
+def test_version_command_without_a_context_object_fails_reading_it(cli_runner):
+ result = cli_runner.invoke(cli.commands["version"], [])
+
+ assert result.exit_code == 1
+ assert type(result.exception) is AttributeError
+ assert result.output == ""
diff --git a/tests/unit/proxy/client/cli/test_keys_commands.py b/tests/unit/proxy/client/cli/test_keys_commands.py
index 5cc0fb70881..7053e7d36fb 100644
--- a/tests/unit/proxy/client/cli/test_keys_commands.py
+++ b/tests/unit/proxy/client/cli/test_keys_commands.py
@@ -2,14 +2,16 @@ import json
import os
from unittest.mock import patch
-import requests
-
-
-
import pytest
+import requests
+import responses
from click.testing import CliRunner
from litellm.proxy.client.cli import cli
+from litellm.proxy.client.cli.commands.keys import keys
+
+DESTINATION = "https://destination.example.com"
+SOURCE = "https://source.example.com"
@pytest.fixture
@@ -538,3 +540,84 @@ def test_keys_import_with_all_key_properties(mock_keys_client, cli_runner):
budget_id="budget-456",
config={"max_tokens": 1000},
)
+
+
+def _sent_requests():
+ return [
+ (call.request.method, call.request.url, call.request.headers.get("Authorization")) for call in responses.calls
+ ]
+
+
+@pytest.mark.parametrize(("api_key", "authorization"), [("sk-destination", "Bearer sk-destination"), (None, None)])
+@pytest.mark.parametrize(
+ ("args", "method", "path", "sent_body", "payload"),
+ [
+ (
+ ["list", "--format", "json"],
+ "GET",
+ "/key/list?return_full_object=true&include_team_keys=false",
+ None,
+ {"keys": [{"token": "abc123", "key_alias": "alias1"}], "total_count": 1},
+ ),
+ (["generate", "--key-alias", "ci"], "POST", "/key/generate", {"key_alias": "ci"}, {"key": "sk-new"}),
+ (
+ ["delete", "--keys", "sk-old"],
+ "POST",
+ "/key/delete",
+ {"keys": ["sk-old"], "key_aliases": None},
+ {"deleted_keys": ["sk-old"]},
+ ),
+ ],
+)
+@responses.activate
+def test_keys_command_calls_the_server_stored_on_the_cli_context(
+ cli_runner, args, method, path, sent_body, payload, api_key, authorization
+):
+ responses.add(method, f"{DESTINATION}{path.partition('?')[0]}", json=payload)
+
+ result = cli_runner.invoke(keys, args, obj={"base_url": DESTINATION, "api_key": api_key})
+
+ assert result.exit_code == 0, result.output
+ assert json.loads(result.output) == payload
+ assert _sent_requests() == [(method, f"{DESTINATION}{path}", authorization)]
+ assert (responses.calls[0].request.body and json.loads(responses.calls[0].request.body)) == sent_body
+
+
+@pytest.mark.parametrize(("api_key", "authorization"), [("sk-destination", "Bearer sk-destination"), (None, None)])
+@responses.activate
+def test_keys_import_reads_the_source_and_writes_to_the_server_stored_on_the_cli_context(
+ cli_runner, api_key, authorization
+):
+ responses.get(f"{SOURCE}/key/list", json={"keys": [{"key_alias": "imported", "user_id": "u1"}]})
+ responses.post(f"{DESTINATION}/key/generate", json={"key": "sk-imported"})
+
+ result = cli_runner.invoke(
+ keys,
+ ["import", "--source-base-url", SOURCE, "--source-api-key", "sk-source"],
+ obj={"base_url": DESTINATION, "api_key": api_key},
+ )
+
+ assert result.exit_code == 0, result.output
+ assert "Successfully imported: 1" in result.output
+ assert _sent_requests() == [
+ ("GET", f"{SOURCE}/key/list?page=1&size=100&return_full_object=true", "Bearer sk-source"),
+ ("POST", f"{DESTINATION}/key/generate", authorization),
+ ]
+ assert json.loads(responses.calls[1].request.body) == {"key_alias": "imported", "user_id": "u1"}
+
+
+@pytest.mark.parametrize(
+ "args",
+ [
+ ["list"],
+ ["generate", "--key-alias", "ci"],
+ ["delete", "--keys", "sk-old"],
+ ["import", "--source-base-url", SOURCE, "--source-api-key", "sk-source"],
+ ],
+)
+@responses.activate
+def test_keys_command_without_cli_context_values_fails_before_any_request(cli_runner, args):
+ result = cli_runner.invoke(keys, args)
+
+ assert result.exit_code == 1
+ assert _sent_requests() == []
diff --git a/tests/unit/proxy/client/test_http_commands.py b/tests/unit/proxy/client/test_http_commands.py
index 04894248ff2..b4dd8070134 100644
--- a/tests/unit/proxy/client/test_http_commands.py
+++ b/tests/unit/proxy/client/test_http_commands.py
@@ -118,3 +118,122 @@ def test_request_invalid_header(runner):
)
assert result.exit_code == 2 # Click error code for invalid parameter
assert "Invalid header format" in result.output
+
+
+@pytest.mark.parametrize(
+ "raw_body",
+ [
+ '{"count": "12", "ratio": 1, "whole": 1.0, "enabled": "true", "flag": true, "missing": null}',
+ '{"data": [{"id": 1, "tags": ["a", 2, null, false]}, "two", null], "nested": {"deep": {"n": "007"}}}',
+ '[{"id": "1"}, 2, "3", null]',
+ '"plain string"',
+ "12",
+ "null",
+ "{}",
+ ],
+)
+@responses.activate
+def test_request_prints_the_server_json_unchanged(runner, raw_body):
+ responses.add(
+ responses.GET,
+ "http://localhost:4000/models",
+ body=raw_body,
+ content_type="application/json",
+ status=200,
+ )
+ result = runner.invoke(
+ http,
+ ["request", "GET", "/models"],
+ obj={"base_url": "http://localhost:4000", "api_key": "sk-test-key"},
+ )
+ assert result.exit_code == 0
+ assert json.dumps(json.loads(result.stdout)) == json.dumps(json.loads(raw_body))
+
+
+@pytest.mark.parametrize(
+ "raw_body",
+ [
+ '{"error": {"message": "model not found", "code": "404", "status": 404, "retry": false}}',
+ '["first", {"detail": null}]',
+ '"denied"',
+ ],
+)
+@responses.activate
+def test_request_http_error_prints_the_json_error_body_unchanged(runner, raw_body):
+ responses.add(
+ responses.GET,
+ "http://localhost:4000/models",
+ body=raw_body,
+ content_type="application/json",
+ status=404,
+ )
+ result = runner.invoke(
+ http,
+ ["request", "GET", "/models"],
+ obj={"base_url": "http://localhost:4000", "api_key": "sk-test-key"},
+ )
+ assert result.exit_code == 1
+ assert result.stderr.splitlines()[:1] == ["Error: HTTP 404"]
+ assert json.dumps(json.loads(result.stdout)) == json.dumps(json.loads(raw_body))
+
+
+@responses.activate
+def test_request_http_error_with_a_non_json_body_echoes_the_text(runner):
+ responses.add(
+ responses.GET,
+ "http://localhost:4000/models",
+ body="upstream exploded",
+ status=502,
+ )
+ result = runner.invoke(
+ http,
+ ["request", "GET", "/models"],
+ obj={"base_url": "http://localhost:4000", "api_key": "sk-test-key"},
+ )
+ assert result.exit_code == 1
+ assert result.stdout == ""
+ assert result.stderr.splitlines()[:2] == ["Error: HTTP 502", "upstream exploded"]
+
+
+@pytest.mark.parametrize(
+ ("api_key", "expected_authorization"),
+ [
+ ("sk-test-key", "Bearer sk-test-key"),
+ (None, None),
+ ("", None),
+ ],
+)
+@responses.activate
+def test_request_sends_the_context_api_key_to_the_context_base_url(runner, api_key, expected_authorization):
+ responses.add(responses.GET, "http://proxy.internal:4100/models", json={"models": []}, status=200)
+ result = runner.invoke(
+ http,
+ ["request", "GET", "/models"],
+ obj={"base_url": "http://proxy.internal:4100/", "api_key": api_key},
+ )
+ assert result.exit_code == 0
+ assert len(responses.calls) == 1
+ assert responses.calls[0].request.url == "http://proxy.internal:4100/models"
+ assert responses.calls[0].request.headers.get("Authorization") == expected_authorization
+
+
+@pytest.mark.parametrize(
+ ("obj", "missing_key"),
+ [
+ ({"api_key": "sk-test-key"}, "base_url"),
+ ({"base_url": "http://localhost:4000"}, "api_key"),
+ ({}, "base_url"),
+ ],
+)
+def test_request_with_an_incomplete_context_fails_on_the_first_missing_value(runner, obj, missing_key):
+ result = runner.invoke(http, ["request", "GET", "/models"], obj=obj)
+ assert result.exit_code == 1
+ assert type(result.exception) is KeyError
+ assert result.exception.args == (missing_key,)
+
+
+def test_request_without_a_context_object_fails_instead_of_assuming_a_default_server(runner):
+ result = runner.invoke(http, ["request", "GET", "/models"], obj=None)
+ assert result.exit_code == 1
+ assert type(result.exception) is TypeError
+ assert result.output == ""
diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py b/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py
index 67ce4046ce0..80bde77afdc 100644
--- a/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py
+++ b/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py
@@ -3,10 +3,12 @@ Test DynamoAI Guardrails integration
"""
import importlib
+import json
import os
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
+import respx
import litellm
from litellm.caching.caching import DualCache
@@ -118,6 +120,89 @@ async def test_dynamoai_allows_content_with_none_action():
assert result == request_data
+_ANALYZE_URL = "https://api.dynamo.ai/v1/moderation/analyze/"
+_BLOCK_VERDICT = {
+ "finalAction": "BLOCK",
+ "appliedPolicies": [
+ {
+ "policy": {"id": "policy-123", "name": "Toxicity Policy", "method": "TOXICITY"},
+ "outputs": {"action": "BLOCK", "message": "Content contains toxic language"},
+ }
+ ],
+}
+
+
+def _post_call_guardrail() -> DynamoAIGuardrails:
+ return DynamoAIGuardrails(
+ guardrail_name="test-dynamoai",
+ api_key="test-api-key",
+ api_base="https://api.dynamo.ai",
+ event_hook="post_call",
+ default_on=True,
+ )
+
+
+def _assistant_reply(content: str | None) -> litellm.ModelResponse:
+ return litellm.ModelResponse(
+ choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": content}}]
+ )
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("httpx_transport")
+async def test_post_call_hook_sends_the_assistant_text_to_dynamoai_and_allows_a_clean_reply(
+ respx_mock: respx.MockRouter,
+) -> None:
+ route = respx_mock.post(_ANALYZE_URL).respond(json={"finalAction": "NONE", "appliedPolicies": []})
+ request_data = {"model": "gpt-5.5", "messages": [{"role": "user", "content": "hi"}]}
+
+ result = await _post_call_guardrail().async_post_call_success_hook(
+ data=request_data, user_api_key_dict=UserAPIKeyAuth(), response=_assistant_reply("Hello, how are you?")
+ )
+
+ assert result is None
+ sent = route.calls.last.request
+ assert json.loads(sent.content) == {"messages": [{"role": "assistant", "content": "Hello, how are you?"}]}
+ assert sent.headers["Authorization"] == "Bearer test-api-key"
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("httpx_transport")
+async def test_post_call_hook_raises_when_dynamoai_blocks_the_assistant_text(respx_mock: respx.MockRouter) -> None:
+ respx_mock.post(_ANALYZE_URL).respond(json=_BLOCK_VERDICT)
+
+ with pytest.raises(ValueError, match="Guardrail failed") as blocked:
+ await _post_call_guardrail().async_post_call_success_hook(
+ data={"model": "gpt-5.5", "messages": [{"role": "user", "content": "hi"}]},
+ user_api_key_dict=UserAPIKeyAuth(),
+ response=_assistant_reply("This is harmful content"),
+ )
+
+ assert str(blocked.value) == (
+ "Guardrail failed: 1 violation(s) detected\n\n"
+ "- TOXICITY POLICY:\n"
+ " Action: BLOCK\n"
+ " Method: TOXICITY\n"
+ " Message: Content contains toxic language\n"
+ " Policy ID: policy-123"
+ )
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("httpx_transport")
+async def test_post_call_hook_skips_dynamoai_when_the_reply_has_no_text(respx_mock: respx.MockRouter) -> None:
+ route = respx_mock.post(_ANALYZE_URL).respond(json=_BLOCK_VERDICT)
+
+ result = await _post_call_guardrail().async_post_call_success_hook(
+ data={"model": "gpt-5.5", "messages": [{"role": "user", "content": "hi"}]},
+ user_api_key_dict=UserAPIKeyAuth(),
+ response=_assistant_reply(None),
+ )
+
+ assert result is None
+ assert route.called is False
+
+
@pytest.fixture(autouse=True)
def _vcr_outcome_gate(request, vcr):
install_live_call_probe(request, vcr)
diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py b/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py b/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py
new file mode 100644
index 00000000000..4a591c08bf3
--- /dev/null
+++ b/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py
@@ -0,0 +1,124 @@
+import json
+
+import pytest
+import respx
+
+from litellm.caching.caching import DualCache
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.guardrails.guardrail_hooks.ibm_guardrails.ibm_detector import IBMGuardrailDetector
+
+pytestmark = pytest.mark.usefixtures("httpx_transport")
+
+_BASE_URL = "https://ibm-detector.example.test"
+_DETECTOR_SERVER_URL = f"{_BASE_URL}/api/v1/text/contents"
+_ORCHESTRATOR_URL = f"{_BASE_URL}/api/v2/text/detection/content"
+_NESTED_PARAMS: dict[str, object] = {
+ "threshold": 0.5,
+ "max_hits": 3,
+ "strict": True,
+ "labels": ["hap", None],
+ "regex": {"email": ".+@.+"},
+}
+_DETECTOR_PARAMS_CASES = [
+ pytest.param(None, {}, id="omitted"),
+ pytest.param({}, {}, id="empty"),
+ pytest.param(_NESTED_PARAMS, _NESTED_PARAMS, id="nested-mixed-types"),
+]
+_HAP_DETECTION = {
+ "start": 0,
+ "end": 5,
+ "text": "hello",
+ "detection": "HAP",
+ "detection_type": "hap",
+ "score": 0.97,
+}
+
+
+def _guardrail(*, is_detector_server: bool, detector_params: dict[str, object] | None) -> IBMGuardrailDetector:
+ return IBMGuardrailDetector(
+ guardrail_name="ibm-guard",
+ auth_token="ibm-token",
+ base_url=_BASE_URL,
+ detector_id="hap-detector",
+ is_detector_server=is_detector_server,
+ detector_params=detector_params,
+ event_hook="pre_call",
+ default_on=True,
+ )
+
+
+def _request_data() -> dict[str, object]:
+ return {"messages": [{"role": "user", "content": "hello there"}]}
+
+
+def test_constructor_keeps_the_configured_detector_params_object() -> None:
+ guardrail = _guardrail(is_detector_server=True, detector_params=_NESTED_PARAMS)
+
+ assert guardrail.detector_params is _NESTED_PARAMS
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("detector_params", "sent_params"), _DETECTOR_PARAMS_CASES)
+async def test_detector_server_request_carries_configured_detector_params(
+ respx_mock: respx.MockRouter, detector_params: dict[str, object] | None, sent_params: dict[str, object]
+) -> None:
+ route = respx_mock.post(_DETECTOR_SERVER_URL).respond(json=[[]])
+ data = _request_data()
+
+ returned = await _guardrail(is_detector_server=True, detector_params=detector_params).async_pre_call_hook(
+ user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type="completion"
+ )
+
+ sent = route.calls.last.request
+ assert returned is data
+ assert repr(json.loads(sent.content)) == repr({"contents": ["hello there"], "detector_params": sent_params})
+ assert sent.headers["detector-id"] == "hap-detector"
+ assert sent.headers["Authorization"] == "Bearer ibm-token"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("detector_params", "sent_params"), _DETECTOR_PARAMS_CASES)
+async def test_orchestrator_request_keys_detector_params_by_detector_id(
+ respx_mock: respx.MockRouter, detector_params: dict[str, object] | None, sent_params: dict[str, object]
+) -> None:
+ route = respx_mock.post(_ORCHESTRATOR_URL).respond(json={"detections": []})
+ data = _request_data()
+
+ returned = await _guardrail(is_detector_server=False, detector_params=detector_params).async_pre_call_hook(
+ user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type="completion"
+ )
+
+ sent = route.calls.last.request
+ assert returned is data
+ assert repr(json.loads(sent.content)) == repr(
+ {"content": "hello there", "detectors": {"hap-detector": sent_params}}
+ )
+ assert "detector-id" not in sent.headers
+
+
+@pytest.mark.asyncio
+async def test_detector_server_detection_blocks_the_request_with_a_violation_summary(
+ respx_mock: respx.MockRouter,
+) -> None:
+ respx_mock.post(_DETECTOR_SERVER_URL).respond(json=[[_HAP_DETECTION]])
+
+ with pytest.raises(ValueError, match=r"IBM Guardrail Detector failed: 1 violation\(s\) detected") as blocked:
+ await _guardrail(is_detector_server=True, detector_params=_NESTED_PARAMS).async_pre_call_hook(
+ user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=_request_data(), call_type="completion"
+ )
+
+ assert " - HAP (score: 0.970)\n Text: 'hello'" in str(blocked.value)
+
+
+@pytest.mark.asyncio
+async def test_orchestrator_detection_blocks_the_request_and_names_the_detector(
+ respx_mock: respx.MockRouter,
+) -> None:
+ respx_mock.post(_ORCHESTRATOR_URL).respond(json={"detections": [_HAP_DETECTION]})
+
+ with pytest.raises(ValueError, match=r"IBM Guardrail Detector failed: 1 violation\(s\) detected") as blocked:
+ await _guardrail(is_detector_server=False, detector_params=_NESTED_PARAMS).async_pre_call_hook(
+ user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=_request_data(), call_type="completion"
+ )
+
+ assert "- HAP (detector: hap-detector, score: 0.970)\n Text: 'hello'" in str(blocked.value)
diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py
index 24fb60e5389..973d6bc9f7b 100644
--- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py
+++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py
@@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch
import httpx
import litellm
import pytest
+from pydantic import ValidationError
from litellm.integrations.custom_logger import CustomLogger
@@ -2864,6 +2865,51 @@ def test_build_complete_streaming_response(all_chunks):
assert result.usage.total_tokens == 266
+_USABLE_MESSAGE_START: Final = 'data: {"type": "message_start", "message": {"id": "msg_real", "model": "claude-real"}}'
+
+
+@pytest.mark.parametrize(
+ "unusable_line",
+ [
+ "data: [1, 2]",
+ 'data: "message_start"',
+ "data: 7",
+ "data: null",
+ "data: not json",
+ 'data: {"type": "message_start"}',
+ 'data: {"type": "message_start", "message": null}',
+ 'data: {"type": "message_start", "message": {}}',
+ 'data: {"type": "message_start", "message": ""}',
+ 'data: {"type": "message_start", "message": []}',
+ 'data: {"type": "message_start", "message": {"id": 7, "model": ["claude"]}}',
+ 'data: {"type": "message_start", "message": {"id": "", "model": null}}',
+ 'data: {"type": "ping", "message": "not an object"}',
+ ],
+)
+def test_model_and_response_id_come_from_the_first_message_start_that_carries_them(unusable_line: str):
+ all_chunks: Final = [f"event: message_start\n{unusable_line}\n\n".encode(), _USABLE_MESSAGE_START]
+
+ extracted: Final = (
+ AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks(all_chunks),
+ AnthropicPassthroughLoggingHandler._extract_response_id_from_anthropic_chunks(all_chunks),
+ )
+
+ assert extracted == ("claude-real", "msg_real"), unusable_line
+
+
+@pytest.mark.parametrize("message", ['"upstream text"', "7", "[1, 2]", "true"])
+def test_message_start_whose_message_is_not_an_object_fails_without_echoing_the_upstream_value(message: str):
+ all_chunks: Final = [f'data: {{"type": "message_start", "message": {message}}}', _USABLE_MESSAGE_START]
+
+ with pytest.raises(ValidationError) as exc_info:
+ AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks(all_chunks)
+
+ assert str(exc_info.value).splitlines()[:2] == [
+ "1 validation error for dict[str,any]",
+ " Input should be a valid dictionary [type=dict_type]",
+ ], message
+
+
@pytest.mark.parametrize(
"end_user_id",
[{"litellm_metadata": {"user": "test"}}, {"metadata": {"user_id": "test"}}],
diff --git a/tests/unit/proxy/policy_engine/test_policy_registry.py b/tests/unit/proxy/policy_engine/test_policy_registry.py
index db7775b6238..1d50230a7a9 100644
--- a/tests/unit/proxy/policy_engine/test_policy_registry.py
+++ b/tests/unit/proxy/policy_engine/test_policy_registry.py
@@ -152,3 +152,134 @@ async def test_update_policy_in_db_rejects_a_malformed_pipeline_before_writing(
f"Error updating policy in DB: 1 validation error for GuardrailPipeline\n{rejected_field}\n"
)
assert table.writes == []
+
+
+_VERSION_CONTENT = {
+ "inherit": "base-policy",
+ "description": "blocks pii",
+ "guardrails_add": ["pii-guard"],
+ "guardrails_remove": ["legacy-guard"],
+ "condition": {"model": "gpt-4"},
+ "pipeline": _STORED_PIPELINE,
+}
+
+_CHANGED_VERSION_CONTENT = {
+ "inherit": None,
+ "description": "",
+ "guardrails_add": ["pii-guard", "toxicity-guard"],
+ "guardrails_remove": [],
+ "condition": {"model": "gpt-4o"},
+ "pipeline": None,
+}
+
+
+class _VersionsTable:
+ def __init__(self, rows: Mapping[str, SimpleNamespace]) -> None:
+ self.rows = rows
+ self.lookups: list[Mapping[str, object]] = []
+
+ async def find_unique(self, where: Mapping[str, object]) -> SimpleNamespace | None:
+ self.lookups.append(where)
+ return self.rows.get(str(where["policy_id"]))
+
+
+def _version_row(policy_id: str, content: Mapping[str, object]) -> SimpleNamespace:
+ return SimpleNamespace(
+ policy_id=policy_id,
+ policy_name="pii-policy",
+ version_number=1,
+ version_status="draft",
+ parent_version_id=None,
+ is_latest=True,
+ published_at=None,
+ production_at=None,
+ created_at=_NOW,
+ updated_at=_NOW,
+ created_by=None,
+ updated_by=None,
+ **{**_VERSION_CONTENT, **content},
+ )
+
+
+def _versions_client(table: _VersionsTable) -> SimpleNamespace:
+ return SimpleNamespace(db=SimpleNamespace(litellm_policytable=table))
+
+
+@pytest.mark.parametrize("field", list(_VERSION_CONTENT))
+@pytest.mark.asyncio
+async def test_compare_versions_reports_exactly_the_one_field_that_differs(field: str):
+ value_a = _VERSION_CONTENT[field]
+ value_b = _CHANGED_VERSION_CONTENT[field]
+ table = _VersionsTable(
+ {
+ "version-a": _version_row("version-a", {}),
+ "version-b": _version_row("version-b", {field: value_b}),
+ }
+ )
+
+ result = await PolicyRegistry().compare_versions("version-a", "version-b", _versions_client(table))
+
+ assert result.field_diffs == {field: {"version_a": value_a, "version_b": value_b}}
+ assert [type(value) for value in result.field_diffs[field].values()] == [type(value_a), type(value_b)]
+ assert (result.version_a.policy_id, result.version_b.policy_id) == ("version-a", "version-b")
+ assert table.lookups == [{"policy_id": "version-a"}, {"policy_id": "version-b"}]
+
+
+@pytest.mark.asyncio
+async def test_compare_versions_lists_every_differing_field_in_content_order_with_each_side():
+ table = _VersionsTable(
+ {
+ "version-a": _version_row("version-a", _CHANGED_VERSION_CONTENT),
+ "version-b": _version_row("version-b", {}),
+ }
+ )
+
+ result = await PolicyRegistry().compare_versions("version-a", "version-b", _versions_client(table))
+
+ assert list(result.field_diffs) == list(_VERSION_CONTENT)
+ assert result.field_diffs == {
+ field: {"version_a": _CHANGED_VERSION_CONTENT[field], "version_b": _VERSION_CONTENT[field]}
+ for field in _VERSION_CONTENT
+ }
+ assert result.field_diffs["inherit"]["version_a"] is None
+ assert result.field_diffs["pipeline"]["version_a"] is None
+
+
+@pytest.mark.asyncio
+async def test_compare_versions_ignores_version_metadata_and_treats_unset_guardrail_lists_as_empty():
+ row_a = _version_row("version-a", {"guardrails_add": None, "guardrails_remove": None})
+ row_b = _version_row("version-b", {"guardrails_add": [], "guardrails_remove": []})
+ row_b.policy_name = "renamed-policy"
+ row_b.version_number = 7
+ row_b.version_status = "production"
+ row_b.parent_version_id = "version-a"
+ row_b.is_latest = False
+ row_b.created_by = "admin"
+ table = _VersionsTable({"version-a": row_a, "version-b": row_b})
+
+ result = await PolicyRegistry().compare_versions("version-a", "version-b", _versions_client(table))
+
+ assert result.field_diffs == {}
+ assert (result.version_a.version_number, result.version_b.version_number) == (1, 7)
+ assert (result.version_a.guardrails_add, result.version_b.guardrails_add) == ([], [])
+
+
+@pytest.mark.parametrize(
+ ("stored_versions", "missing_version"),
+ [
+ pytest.param(("version-b",), "version-a", id="first-missing"),
+ pytest.param(("version-a",), "version-b", id="second-missing"),
+ pytest.param((), "version-a", id="both-missing-names-the-first"),
+ ],
+)
+@pytest.mark.asyncio
+async def test_compare_versions_names_the_missing_version_after_looking_both_up(
+ stored_versions: tuple[str, ...], missing_version: str
+):
+ table = _VersionsTable({policy_id: _version_row(policy_id, {}) for policy_id in stored_versions})
+
+ with pytest.raises(Exception, match=f"^Error comparing versions: Policy {missing_version} not found$") as raised:
+ await PolicyRegistry().compare_versions("version-a", "version-b", _versions_client(table))
+
+ assert type(raised.value) is Exception
+ assert table.lookups == [{"policy_id": "version-a"}, {"policy_id": "version-b"}]
diff --git a/tests/unit/proxy/response_polling/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py
index 81ca0114a8d..979ab499f3b 100644
--- a/tests/unit/proxy/response_polling/test_response_polling_handler.py
+++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py
@@ -22,6 +22,7 @@ from unittest.mock import AsyncMock, Mock, patch
import pytest
from fastapi import Request
+from pydantic import ValidationError
from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler
@@ -1842,3 +1843,47 @@ class TestEdgeCases:
stored = json.loads(call_args.kwargs["value"])
assert stored["output"] == []
+
+
+@pytest.mark.parametrize(
+ ("request_data", "expected_metadata"),
+ [
+ (
+ {"model": "gpt-4o", "input": "Hello", "metadata": {"trace": "t-1", "tags": ["a", None], "attempt": 2}},
+ {"trace": "t-1", "tags": ["a", None], "attempt": 2},
+ ),
+ ({"model": "gpt-4o", "metadata": {}}, {}),
+ ({"model": "gpt-4o", "metadata": None}, None),
+ ({"model": "gpt-4o"}, {}),
+ ({"id": "resp_from_client", "status": "completed", "output": ["x"], "usage": {"total_tokens": 9}}, {}),
+ ],
+)
+async def test_create_initial_state_carries_only_the_request_metadata_into_the_queued_response(
+ request_data, expected_metadata
+):
+ handler = ResponsePollingHandler(redis_cache=None)
+
+ response = await handler.create_initial_state(polling_id="litellm_poll_abc", request_data=request_data)
+
+ assert repr(response.metadata) == repr(expected_metadata)
+ assert type(response.created_at) is int
+ assert response.model_dump(exclude={"created_at", "metadata"}, exclude_none=True) == {
+ "id": "litellm_poll_abc",
+ "object": "response",
+ "status": "queued",
+ "output": [],
+ }
+ assert response.model_fields_set == {"id", "object", "status", "created_at", "output", "metadata", "usage"}
+
+
+@pytest.mark.parametrize("metadata", ["trace-1", ["trace", "t-1"], 7, True, b"trace"])
+async def test_create_initial_state_rejects_request_metadata_that_is_not_an_object(metadata):
+ handler = ResponsePollingHandler(redis_cache=None)
+
+ with pytest.raises(ValidationError) as rejected:
+ await handler.create_initial_state(polling_id="litellm_poll_abc", request_data={"metadata": metadata})
+
+ assert rejected.value.title == "ResponsesAPIResponse"
+ assert [(error["loc"], error["type"], repr(error["input"])) for error in rejected.value.errors()] == [
+ (("metadata",), "dict_type", repr(metadata))
+ ]
diff --git a/tests/unit/proxy/test_caching_routes.py b/tests/unit/proxy/test_caching_routes.py
index 6f80c281565..8a59793495a 100644
--- a/tests/unit/proxy/test_caching_routes.py
+++ b/tests/unit/proxy/test_caching_routes.py
@@ -2,11 +2,14 @@ import json
import pytest
from fastapi.testclient import TestClient
+from pydantic import ValidationError
import litellm
from litellm.caching import RedisCache
+from litellm.caching.caching import Cache
+from litellm.proxy.caching_routes import _extract_cache_params
from litellm.proxy.proxy_server import app
client = TestClient(app)
@@ -334,3 +337,63 @@ def test_cache_redis_info_no_cache():
# Restore original cache
litellm.cache = original_cache
+
+
+def _use_redis_backed_cache(monkeypatch, backend_attributes):
+ cache = Cache(type="redis", host="localhost", port=6379, password="hello")
+ vars(cache.cache).update(backend_attributes)
+ monkeypatch.setattr(litellm, "cache", cache)
+
+
+@pytest.mark.parametrize(
+ ("field", "value"),
+ [
+ ("host", "redis.internal"),
+ ("port", 6379),
+ ("port", "6379"),
+ ("namespace", "team-a"),
+ ("redis_version", 7.2),
+ ("redis_version", 7),
+ ("redis_version", "7.2.4"),
+ ],
+)
+def test_extract_cache_params_keeps_each_backend_value_and_its_type(monkeypatch, field, value):
+ _use_redis_backed_cache(monkeypatch, {field: value})
+
+ params = _extract_cache_params()
+
+ assert sorted(params) == ["host", "namespace", "port", "redis_kwargs", "redis_version"]
+ assert repr(params[field]) == repr(value)
+
+
+def test_extract_cache_params_reads_connection_settings_and_masks_the_password(monkeypatch):
+ _use_redis_backed_cache(monkeypatch, {})
+
+ redis_kwargs = _extract_cache_params()["redis_kwargs"]
+
+ assert redis_kwargs["host"] == "localhost"
+ assert repr(redis_kwargs["port"]) == "6379"
+ assert "hello" not in json.dumps(redis_kwargs)
+
+
+def test_extract_cache_params_reports_every_malformed_backend_field(monkeypatch):
+ _use_redis_backed_cache(monkeypatch, {"host": 7, "namespace": 5, "redis_version": ["7"]})
+
+ with pytest.raises(ValidationError) as exc_info:
+ _extract_cache_params()
+
+ assert exc_info.value.title == "HealthCheckCacheParams"
+ assert sorted({error["loc"][0] for error in exc_info.value.errors()}) == ["host", "namespace", "redis_version"]
+
+
+def test_extract_cache_params_is_empty_when_no_cache_is_configured(monkeypatch):
+ monkeypatch.setattr(litellm, "cache", None)
+
+ assert _extract_cache_params() == {}
+
+
+@pytest.mark.parametrize("attribute_name", [1, None, (1, 2), b"host"])
+def test_extract_cache_params_is_empty_when_a_backend_attribute_name_is_not_a_string(monkeypatch, attribute_name):
+ _use_redis_backed_cache(monkeypatch, {attribute_name: "redis.internal"})
+
+ assert _extract_cache_params() == {}, attribute_name
diff --git a/tests/unit/proxy/test_plugin_routes.py b/tests/unit/proxy/test_plugin_routes.py
index c8d2939385d..7d3fbda14fd 100644
--- a/tests/unit/proxy/test_plugin_routes.py
+++ b/tests/unit/proxy/test_plugin_routes.py
@@ -15,7 +15,10 @@ import asyncio
from unittest.mock import MagicMock
import pytest
+import respx
+from starlette.requests import Request
+import litellm
from litellm.proxy._types import (
ConfigGeneralSettings,
LitellmUserRoles,
@@ -230,3 +233,90 @@ def test_configured_custom_key_header_is_stripped() -> None:
assert "x-my-tenant-key" in _request_strip_headers()
finally:
proxy_server.general_settings = original
+
+
+def _plugin_proxy_request(headers: dict[str, str]) -> Request:
+ async def receive() -> dict[str, object]:
+ return {"type": "http.request", "body": b'{"q": 1}', "more_body": False}
+
+ return Request(
+ {
+ "type": "http",
+ "method": "POST",
+ "path": "/plugin-proxy/p/items",
+ "query_string": b"page=2",
+ "headers": [(name.lower().encode(), value.encode()) for name, value in headers.items()],
+ },
+ receive,
+ )
+
+
+@pytest.mark.parametrize(
+ ("user_id", "expected_identity_headers"),
+ [
+ ("admin-1", {"x-litellm-user-id": "admin-1", "x-litellm-user-role": "proxy_admin"}),
+ (None, {"x-litellm-user-role": "proxy_admin"}),
+ ("", {"x-litellm-user-role": "proxy_admin"}),
+ ],
+)
+def test_plugin_proxy_forwards_caller_identity_and_replaces_the_caller_credential(
+ monkeypatch: pytest.MonkeyPatch,
+ respx_mock: respx.MockRouter,
+ user_id: str | None,
+ expected_identity_headers: dict[str, str],
+) -> None:
+ from litellm.proxy.plugin_routes import plugin_proxy
+
+ monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+ litellm.in_memory_llm_clients_cache.flush_cache()
+ register_plugins_from_config(
+ {"plugins": [{"name": "p", "url": "http://plugin.example", "plugin_key": "sk-plugin"}]}
+ )
+ upstream = respx_mock.post("http://plugin.example/items?page=2").respond(text="plugin says hi")
+
+ response = asyncio.run(
+ plugin_proxy(
+ "p",
+ "items",
+ _plugin_proxy_request({"Authorization": "Bearer sk-caller", "X-Trace-Id": "t-1"}),
+ UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN, user_id=user_id),
+ )
+ )
+ register_plugins_from_config({})
+ litellm.in_memory_llm_clients_cache.flush_cache()
+
+ forwarded = upstream.calls[0].request
+ assert response.status_code == 200
+ assert response.body == b"plugin says hi"
+ assert forwarded.content == b'{"q": 1}'
+ assert forwarded.headers["authorization"] == "Bearer sk-plugin"
+ assert forwarded.headers["x-trace-id"] == "t-1"
+ assert {
+ name: value for name, value in forwarded.headers.items() if name.startswith("x-litellm-user-")
+ } == expected_identity_headers
+
+
+def test_plugin_proxy_refuses_non_admin_callers_without_contacting_the_plugin(
+ monkeypatch: pytest.MonkeyPatch,
+ respx_mock: respx.MockRouter,
+) -> None:
+ from litellm.proxy.plugin_routes import plugin_proxy
+
+ monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+ litellm.in_memory_llm_clients_cache.flush_cache()
+ register_plugins_from_config({"plugins": [{"name": "p", "url": "http://plugin.example"}]})
+ upstream = respx_mock.post("http://plugin.example/items?page=2").respond(text="plugin says hi")
+
+ response = asyncio.run(
+ plugin_proxy(
+ "p",
+ "items",
+ _plugin_proxy_request({"Authorization": "Bearer sk-caller"}),
+ UserAPIKeyAuth(api_key="sk-user", user_role=LitellmUserRoles.INTERNAL_USER, user_id="user-1"),
+ )
+ )
+ register_plugins_from_config({})
+
+ assert response.status_code == 403
+ assert response.body == b"Plugin proxy access requires proxy_admin role."
+ assert not upstream.called
diff --git a/tests/unit/proxy/types_utils/test_utils.py b/tests/unit/proxy/types_utils/test_utils.py
new file mode 100644
index 00000000000..eb7a57ba37e
--- /dev/null
+++ b/tests/unit/proxy/types_utils/test_utils.py
@@ -0,0 +1,250 @@
+import io
+import re
+import sys
+import tempfile
+from collections.abc import Iterator, Mapping
+from pathlib import Path
+
+import boto3
+import pytest
+
+from litellm.proxy.types_utils.utils import get_instance_fn
+
+_PLUGIN_SOURCE = """
+class Handler:
+ label = "plugin-handler"
+
+
+handler = Handler()
+alias = handler
+nothing = None
+zero = 0
+disabled = False
+empty = []
+label = "plugin"
+location = __file__
+
+
+def build():
+ return "built"
+"""
+
+_DYNAMIC_PLUGIN_SOURCE = """
+def __getattr__(name):
+ if name.startswith("make_"):
+ return name.removeprefix("make_").upper()
+ if name == "needs_dependency":
+ raise ImportError("optional dependency missing")
+ raise LookupError(f"plugin cannot provide {name}")
+"""
+
+_WARNING_PLUGIN_TEMPLATE = """
+import warnings
+
+
+def __getattr__(name):
+ warnings.warn(f"{{name}} is deprecated", DeprecationWarning, stacklevel={stacklevel})
+ return name.upper()
+"""
+
+_PLAIN_ATTRIBUTES = [
+ pytest.param("nothing", "NoneType", "None", id="none"),
+ pytest.param("zero", "int", "0", id="zero"),
+ pytest.param("disabled", "bool", "False", id="false"),
+ pytest.param("empty", "list", "[]", id="empty-list"),
+ pytest.param("label", "str", "'plugin'", id="string"),
+]
+
+_INSTALLED_PLUGIN = "types_utils_installed_plugin"
+
+
+class _FakeS3:
+ def __init__(self, objects: Mapping[tuple[str, str], bytes]) -> None:
+ self.objects = objects
+ self.requests: list[tuple[str, str]] = []
+
+ def get_object(self, Bucket: str, Key: str) -> Mapping[str, io.BytesIO]:
+ self.requests.append((Bucket, Key))
+ return {"Body": io.BytesIO(self.objects[(Bucket, Key)])}
+
+
+@pytest.fixture
+def local_config(tmp_path: Path) -> str:
+ (tmp_path / "local_plugin.py").write_text(_PLUGIN_SOURCE)
+ (tmp_path / "dynamic_plugin.py").write_text(_DYNAMIC_PLUGIN_SOURCE)
+ return str(tmp_path / "config.yaml")
+
+
+@pytest.fixture
+def installed_plugin(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[str]:
+ site = tmp_path / "site"
+ site.mkdir()
+ (site / f"{_INSTALLED_PLUGIN}.py").write_text(_PLUGIN_SOURCE)
+ monkeypatch.syspath_prepend(str(site))
+ yield _INSTALLED_PLUGIN
+ sys.modules.pop(_INSTALLED_PLUGIN, None)
+
+
+@pytest.fixture
+def downloads(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
+ directory = tmp_path / "downloads"
+ directory.mkdir()
+ monkeypatch.setattr(tempfile, "tempdir", str(directory))
+ return directory
+
+
+@pytest.fixture
+def plugin_bucket(downloads: Path, monkeypatch: pytest.MonkeyPatch) -> _FakeS3:
+ bucket = _FakeS3(
+ {
+ ("plugins-bucket", "loggers/remote_plugin.py"): _PLUGIN_SOURCE.encode(),
+ ("plugins-bucket", "loggers/warning_plugin.py"): _WARNING_PLUGIN_TEMPLATE.format(stacklevel=4).encode(),
+ }
+ )
+
+ def fake_client(
+ service_name: str,
+ aws_access_key_id: str,
+ aws_secret_access_key: str,
+ aws_session_token: str | None,
+ ) -> _FakeS3:
+ assert service_name == "s3"
+ return bucket
+
+ monkeypatch.setenv("AWS_ACCESS_KEY_ID", "unit-test")
+ monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "unit-test")
+ monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
+ monkeypatch.setattr(boto3, "client", fake_client)
+ return bucket
+
+
+@pytest.mark.parametrize(("attribute", "expected_type", "expected_repr"), _PLAIN_ATTRIBUTES)
+def test_get_instance_fn_returns_the_attribute_of_the_module_beside_the_config_unchanged(
+ local_config: str, attribute: str, expected_type: str, expected_repr: str
+):
+ result = get_instance_fn(f"local_plugin.{attribute}", config_file_path=local_config)
+
+ assert (type(result).__name__, repr(result)) == (expected_type, expected_repr)
+
+
+def test_get_instance_fn_returns_instances_and_callables_defined_in_the_local_module(local_config: str):
+ handler = get_instance_fn("local_plugin.handler", config_file_path=local_config)
+ build = get_instance_fn("local_plugin.build", config_file_path=local_config)
+
+ assert (type(handler).__name__, handler.label) == ("Handler", "plugin-handler")
+ assert build() == "built"
+
+
+def test_get_instance_fn_surfaces_a_name_the_local_module_lacks_as_attribute_error(local_config: str):
+ with pytest.raises(AttributeError) as raised:
+ get_instance_fn("local_plugin.missing", config_file_path=local_config)
+
+ assert type(raised.value) is AttributeError
+ assert (raised.value.obj.__name__, raised.value.name) == ("local_plugin", "missing")
+ assert raised.value.__cause__ is None
+
+
+@pytest.mark.parametrize(
+ ("attribute", "expected"),
+ [
+ pytest.param("make_logger", "LOGGER", id="computed"),
+ pytest.param("make_", "", id="computed-empty"),
+ ],
+)
+def test_get_instance_fn_returns_what_a_module_level_getattr_computes(local_config: str, attribute: str, expected: str):
+ result = get_instance_fn(f"dynamic_plugin.{attribute}", config_file_path=local_config)
+
+ assert (type(result), result) == (str, expected)
+
+
+def test_get_instance_fn_lets_an_error_from_a_module_level_getattr_through_unchanged(local_config: str):
+ with pytest.raises(LookupError, match=r"^plugin cannot provide other$") as raised:
+ get_instance_fn("dynamic_plugin.other", config_file_path=local_config)
+
+ assert type(raised.value) is LookupError
+
+
+def test_get_instance_fn_rewords_an_import_error_raised_while_reading_the_attribute(local_config: str):
+ with pytest.raises(ImportError, match=r"^Could not import needs_dependency from dynamic_plugin$") as raised:
+ get_instance_fn("dynamic_plugin.needs_dependency", config_file_path=local_config)
+
+ assert type(raised.value) is ImportError
+ assert (type(raised.value.__cause__), str(raised.value.__cause__)) == (ImportError, "optional dependency missing")
+
+
+def test_get_instance_fn_returns_the_very_object_an_installed_module_holds(installed_plugin: str):
+ handler = get_instance_fn(f"{installed_plugin}.handler")
+ alias = get_instance_fn(f"{installed_plugin}.alias")
+
+ assert handler is sys.modules[installed_plugin].handler
+ assert alias is handler
+
+
+def test_get_instance_fn_surfaces_a_name_an_installed_module_lacks_as_attribute_error(installed_plugin: str):
+ with pytest.raises(AttributeError) as raised:
+ get_instance_fn(f"{installed_plugin}.missing")
+
+ assert type(raised.value) is AttributeError
+ assert (raised.value.obj, raised.value.name) == (sys.modules[installed_plugin], "missing")
+ assert raised.value.__cause__ is None
+
+
+@pytest.mark.parametrize(("attribute", "expected_type", "expected_repr"), _PLAIN_ATTRIBUTES)
+def test_get_instance_fn_loads_the_attribute_from_the_bucket_object_unchanged(
+ plugin_bucket: _FakeS3, attribute: str, expected_type: str, expected_repr: str
+):
+ result = get_instance_fn(
+ f"s3://plugins-bucket/loggers/remote_plugin.{attribute}", config_file_path="/etc/litellm/config.yaml"
+ )
+
+ assert (type(result).__name__, repr(result)) == (expected_type, expected_repr)
+ assert plugin_bucket.requests == [("plugins-bucket", "loggers/remote_plugin.py")]
+
+
+def test_get_instance_fn_runs_the_bucket_module_from_a_temporary_file_and_removes_it(
+ plugin_bucket: _FakeS3, downloads: Path
+):
+ location = get_instance_fn(
+ "s3://plugins-bucket/loggers/remote_plugin.location", config_file_path="/etc/litellm/config.yaml"
+ )
+
+ assert (Path(location).parent, Path(location).suffix) == (downloads, ".py")
+ assert not Path(location).exists()
+
+
+def test_get_instance_fn_reports_a_name_the_bucket_module_lacks_as_a_failed_load(plugin_bucket: _FakeS3):
+ prefix = "Failed to load custom logger from s3://plugins-bucket/loggers/remote_plugin.missing: "
+
+ with pytest.raises(ImportError, match=f"^{re.escape(prefix)}") as raised:
+ get_instance_fn(
+ "s3://plugins-bucket/loggers/remote_plugin.missing", config_file_path="/etc/litellm/config.yaml"
+ )
+
+ assert type(raised.value) is ImportError
+ assert type(raised.value.__cause__) is AttributeError
+ assert str(raised.value) == prefix + str(raised.value.__cause__)
+ assert (raised.value.__cause__.obj.__name__, raised.value.__cause__.name) == ("loggers/remote_plugin", "missing")
+ assert plugin_bucket.requests == [("plugins-bucket", "loggers/remote_plugin.py")]
+
+
+def test_get_instance_fn_attributes_a_stacklevel_three_warning_from_a_local_plugin_to_its_own_caller(tmp_path: Path):
+ (tmp_path / "warning_plugin.py").write_text(_WARNING_PLUGIN_TEMPLATE.format(stacklevel=3))
+
+ with pytest.warns(DeprecationWarning, match=r"^old_handler is deprecated$") as recorded:
+ result = get_instance_fn("warning_plugin.old_handler", config_file_path=str(tmp_path / "config.yaml"))
+
+ assert result == "OLD_HANDLER"
+ assert [Path(w.filename) for w in recorded if str(w.message) == "old_handler is deprecated"] == [Path(__file__)]
+
+
+def test_get_instance_fn_attributes_a_stacklevel_four_warning_from_a_bucket_plugin_to_its_own_caller(
+ plugin_bucket: _FakeS3,
+):
+ with pytest.warns(DeprecationWarning, match=r"^old_handler is deprecated$") as recorded:
+ result = get_instance_fn(
+ "s3://plugins-bucket/loggers/warning_plugin.old_handler", config_file_path="/etc/litellm/config.yaml"
+ )
+
+ assert result == "OLD_HANDLER"
+ assert [Path(w.filename) for w in recorded if str(w.message) == "old_handler is deprecated"] == [Path(__file__)]
+ assert plugin_bucket.requests == [("plugins-bucket", "loggers/warning_plugin.py")]
diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_utils.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_utils.py
new file mode 100644
index 00000000000..5b71fcb8037
--- /dev/null
+++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_utils.py
@@ -0,0 +1,94 @@
+import json
+
+import pytest
+
+import litellm
+from litellm.proxy.vector_store_endpoints.utils import get_litellm_managed_vector_store
+from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
+
+
+def _registered_store(litellm_params: object) -> dict[str, object]:
+ return {
+ "vector_store_id": "vs_1",
+ "custom_llm_provider": "bedrock",
+ "vector_store_name": "kb",
+ "vector_store_metadata": '{"source": "db"}',
+ "team_id": None,
+ "litellm_params": litellm_params,
+ }
+
+
+@pytest.mark.parametrize(
+ ("stored_params", "expected_params"),
+ [
+ (
+ '{"api_base": "https://example.com", "embedding": {"api_key": "sk-1", "dims": [1, null]}}',
+ {"api_base": "https://example.com", "embedding": {"api_key": "sk-1", "dims": [1, None]}},
+ ),
+ (
+ '{"top_k": 12, "score_threshold": 12.0, "api_version": "12", "stream": true, "verify": false,'
+ ' "region": null, "max_tokens": 123456789012345678901234567890}',
+ {
+ "top_k": 12,
+ "score_threshold": 12.0,
+ "api_version": "12",
+ "stream": True,
+ "verify": False,
+ "region": None,
+ "max_tokens": 123456789012345678901234567890,
+ },
+ ),
+ ('{"z": 1, "a": 2, "a": 3}', {"z": 1, "a": 3}),
+ ("{}", {}),
+ ],
+)
+async def test_registry_store_with_json_object_string_params_resolves_to_decoded_params(
+ monkeypatch: pytest.MonkeyPatch, stored_params: str, expected_params: dict[str, object]
+):
+ store = _registered_store(stored_params)
+ monkeypatch.setattr(litellm, "vector_store_registry", VectorStoreRegistry([store]))
+
+ resolved = await get_litellm_managed_vector_store("vs_1")
+
+ assert resolved == {**store, "litellm_params": expected_params}
+ assert json.dumps(resolved["litellm_params"]) == json.dumps(expected_params)
+ assert store["litellm_params"] == stored_params
+
+
+@pytest.mark.parametrize(
+ "stored_params",
+ [
+ "[1, 2]",
+ '[["api_base", "https://example.com"]]',
+ '"{\\"api_base\\": \\"https://example.com\\"}"',
+ "12",
+ "12.5",
+ "true",
+ "null",
+ "{not json",
+ "",
+ ],
+)
+async def test_registry_store_with_non_object_or_invalid_json_string_params_resolves_to_empty_params(
+ monkeypatch: pytest.MonkeyPatch, stored_params: str
+):
+ store = _registered_store(stored_params)
+ monkeypatch.setattr(litellm, "vector_store_registry", VectorStoreRegistry([store]))
+
+ resolved = await get_litellm_managed_vector_store("vs_1")
+
+ assert resolved == {**store, "litellm_params": {}}
+ assert store["litellm_params"] == stored_params
+
+
+@pytest.mark.parametrize("stored_params", [{"api_base": "https://example.com"}, {}, None])
+async def test_registry_store_with_non_string_params_is_returned_as_is(
+ monkeypatch: pytest.MonkeyPatch, stored_params: object
+):
+ store = _registered_store(stored_params)
+ monkeypatch.setattr(litellm, "vector_store_registry", VectorStoreRegistry([store]))
+
+ resolved = await get_litellm_managed_vector_store("vs_1")
+
+ assert resolved is store
+ assert resolved["litellm_params"] is stored_params
diff --git a/tests/unit/secret_managers/test_google_secret_manager.py b/tests/unit/secret_managers/test_google_secret_manager.py
index ed93aa2861d..2051be08f7c 100644
--- a/tests/unit/secret_managers/test_google_secret_manager.py
+++ b/tests/unit/secret_managers/test_google_secret_manager.py
@@ -1,12 +1,21 @@
import base64
+import binascii
+import json
from dataclasses import dataclass
from typing import Final
+import google.auth
+import httpx
import litellm
+import litellm.proxy.proxy_server
import pytest
import respx
+from google.auth.credentials import Credentials
+from pydantic import ValidationError
from litellm.secret_managers.google_secret_manager import GoogleSecretManager
+from litellm.secret_managers.main import get_secret
+from litellm.types.secret_managers.main import KeyManagementSettings
@dataclass(frozen=True, slots=True)
@@ -81,3 +90,237 @@ def test_google_secret_manager_returns_cached_values_without_http(
assert missing_value is None
assert cached_value == "lite-llm"
assert upstream.calls.call_count == 0
+
+
+PROJECT_ID: Final = "unit-test-project"
+SECRET_NAME: Final = "UNIT_TEST_GOOGLE_SECRET"
+ACCESS_URL: Final = (
+ f"https://secretmanager.googleapis.com/v1/projects/{PROJECT_ID}/secrets/{SECRET_NAME}/versions/latest:access"
+)
+SECRET_VALUE: Final = "s3cr3t-value"
+ENCODED_SECRET: Final = base64.b64encode(SECRET_VALUE.encode()).decode("ascii")
+
+
+class _StaticCredentials(Credentials):
+ def refresh(self, request: object) -> None:
+ self.token = "unit-test-token"
+
+
+def _default_credentials(scopes: object = None) -> tuple[Credentials, str]:
+ return _StaticCredentials(), PROJECT_ID
+
+
+def _access_response(body: object) -> httpx.Response:
+ return httpx.Response(200, content=json.dumps(body).encode())
+
+
+@pytest.fixture
+def manager(monkeypatch: pytest.MonkeyPatch) -> GoogleSecretManager:
+ monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True)
+ monkeypatch.setattr(litellm, "secret_manager_client", None)
+ monkeypatch.setattr(litellm, "_key_management_system", None)
+ monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings())
+ monkeypatch.setattr(litellm.vertex_chat_completion, "_credentials_project_mapping", {})
+ monkeypatch.setattr(google.auth, "default", _default_credentials)
+ monkeypatch.setenv("GOOGLE_SECRET_MANAGER_PROJECT_ID", PROJECT_ID)
+ monkeypatch.delenv("GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER", raising=False)
+ monkeypatch.delenv("GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL", raising=False)
+ monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False)
+ monkeypatch.delenv("GCS_MOCK", raising=False)
+ monkeypatch.delenv(SECRET_NAME, raising=False)
+ return GoogleSecretManager()
+
+
+@pytest.mark.parametrize(
+ ("body", "expected"),
+ [
+ ({"payload": {"data": ENCODED_SECRET}}, SECRET_VALUE),
+ (
+ {
+ "name": f"projects/123456/secrets/{SECRET_NAME}/versions/7",
+ "payload": {"data": ENCODED_SECRET, "dataCrc32c": "2932599711"},
+ },
+ SECRET_VALUE,
+ ),
+ ({"payload": {"data": ENCODED_SECRET, "labels": {"nested": [1, None]}}, "replication": [1, 2]}, SECRET_VALUE),
+ ({"payload": {"data": "czNj\ncjN0 LXZhbHVl"}}, SECRET_VALUE),
+ ({"payload": {"data": base64.b64encode("clé ☃".encode()).decode("ascii")}}, "clé ☃"),
+ ({"payload": {"data": ""}}, ""),
+ ],
+)
+@respx.mock
+def test_secret_is_decoded_from_the_access_response(manager: GoogleSecretManager, body: object, expected: str) -> None:
+ route: Final = respx.get(ACCESS_URL).mock(return_value=_access_response(body))
+
+ secret: Final = manager.get_secret_from_google_secret_manager(SECRET_NAME)
+
+ assert secret == expected
+ assert type(secret) is str
+ assert manager.cache.cache_dict[SECRET_NAME] is secret
+ assert route.calls.last.request.headers["Authorization"] == "Bearer unit-test-token"
+
+
+@pytest.mark.parametrize(
+ "body",
+ [
+ {},
+ {"name": f"projects/123456/secrets/{SECRET_NAME}/versions/7"},
+ {"payload": {}},
+ {"payload": {"dataCrc32c": "2932599711"}},
+ {"payload": {"data": None}},
+ ],
+)
+@respx.mock
+def test_response_without_secret_data_reports_not_found_and_caches_the_miss(
+ manager: GoogleSecretManager, body: object
+) -> None:
+ respx.get(ACCESS_URL).mock(return_value=_access_response(body))
+
+ with pytest.raises(ValueError, match=rf"^secret {SECRET_NAME} not found in Google Secret Manager$") as raised:
+ manager.get_secret_from_google_secret_manager(SECRET_NAME)
+
+ assert type(raised.value) is ValueError
+ assert manager.cache.cache_dict[SECRET_NAME] is None
+ assert manager.get_secret_from_google_secret_manager(SECRET_NAME) is None
+
+
+@pytest.mark.parametrize(
+ "body",
+ [
+ [],
+ [{"payload": {"data": ENCODED_SECRET}}],
+ [["payload", {"data": ENCODED_SECRET}]],
+ ENCODED_SECRET,
+ 7,
+ None,
+ True,
+ {"payload": None},
+ {"payload": []},
+ {"payload": [["data", ENCODED_SECRET]]},
+ {"payload": ENCODED_SECRET},
+ {"payload": ""},
+ {"payload": 0},
+ {"payload": False},
+ ],
+)
+@respx.mock
+def test_response_that_is_not_a_json_object_is_rejected_without_caching(
+ manager: GoogleSecretManager, body: object
+) -> None:
+ respx.get(ACCESS_URL).mock(return_value=_access_response(body))
+
+ with pytest.raises(ValidationError) as raised:
+ manager.get_secret_from_google_secret_manager(SECRET_NAME)
+
+ assert [error["type"] for error in raised.value.errors()] == ["dict_type"]
+ assert ENCODED_SECRET not in str(raised.value)
+ assert SECRET_NAME not in manager.cache.cache_dict
+
+
+@pytest.mark.parametrize("data", [5, 0, 1.5, True, False, [], [ENCODED_SECRET], {}, {"value": ENCODED_SECRET}])
+@respx.mock
+def test_secret_data_that_is_not_text_is_rejected_without_caching(manager: GoogleSecretManager, data: object) -> None:
+ respx.get(ACCESS_URL).mock(return_value=_access_response({"payload": {"data": data}}))
+
+ with pytest.raises(ValidationError) as raised:
+ manager.get_secret_from_google_secret_manager(SECRET_NAME)
+
+ assert [error["type"] for error in raised.value.errors()] == ["string_type"]
+ assert ENCODED_SECRET not in str(raised.value)
+ assert SECRET_NAME not in manager.cache.cache_dict
+
+
+@pytest.mark.parametrize(
+ "body",
+ [
+ [{"payload": {"data": ENCODED_SECRET}}] * 403 + [7],
+ {"payload": [{"data": ENCODED_SECRET}] * 403 + [None]},
+ {"payload": {"data": [ENCODED_SECRET] * 403 + [429]}},
+ {"payload": {"data": {str(position): "Request timed out" for position in range(403, 430)}}},
+ ],
+)
+@respx.mock
+def test_rejection_of_a_long_response_names_no_position_or_content_from_it(
+ manager: GoogleSecretManager, body: object
+) -> None:
+ respx.get(ACCESS_URL).mock(return_value=_access_response(body))
+
+ with pytest.raises(ValidationError) as raised:
+ manager.get_secret_from_google_secret_manager(SECRET_NAME)
+
+ message: Final = str(raised.value)
+ assert [error["loc"] for error in raised.value.errors()] == [()]
+ assert [fragment for fragment in ("403", "429", "Request timed out", ENCODED_SECRET) if fragment in message] == []
+ assert SECRET_NAME not in manager.cache.cache_dict
+
+
+@pytest.mark.parametrize(
+ ("data", "error_class"),
+ [
+ ("abc", binascii.Error),
+ (base64.b64encode(b"\xff\xfe").decode("ascii"), UnicodeDecodeError),
+ ("éééé", ValueError),
+ ],
+)
+@respx.mock
+def test_undecodable_secret_text_keeps_raising_the_decoder_error(
+ manager: GoogleSecretManager, data: str, error_class: type[Exception]
+) -> None:
+ respx.get(ACCESS_URL).mock(return_value=_access_response({"payload": {"data": data}}))
+
+ with pytest.raises(error_class) as raised:
+ manager.get_secret_from_google_secret_manager(SECRET_NAME)
+
+ assert type(raised.value) is error_class
+ assert SECRET_NAME not in manager.cache.cache_dict
+
+
+@pytest.mark.usefixtures("manager")
+@respx.mock
+def test_get_secret_reads_through_the_configured_manager() -> None:
+ respx.get(ACCESS_URL).mock(return_value=_access_response({"payload": {"data": ENCODED_SECRET}}))
+
+ assert get_secret(SECRET_NAME) == SECRET_VALUE
+
+
+@pytest.mark.usefixtures("manager")
+@pytest.mark.parametrize(
+ "body",
+ [
+ [],
+ {"payload": None},
+ {"payload": ENCODED_SECRET},
+ {"payload": {"data": 5}},
+ {"payload": {"data": [ENCODED_SECRET]}},
+ {"payload": {}},
+ [{"payload": {"data": ENCODED_SECRET}}] * 403 + [7],
+ {"payload": {"data": [ENCODED_SECRET] * 403 + [429]}},
+ {"payload": "Request timed out"},
+ ],
+)
+@respx.mock
+def test_get_secret_falls_back_to_the_environment_when_the_response_is_unusable(
+ monkeypatch: pytest.MonkeyPatch, body: object
+) -> None:
+ monkeypatch.setenv(SECRET_NAME, "from-environment")
+ respx.get(ACCESS_URL).mock(return_value=_access_response(body))
+
+ assert get_secret(SECRET_NAME) == "from-environment"
+
+
+@pytest.mark.usefixtures("manager")
+@pytest.mark.parametrize(
+ "body",
+ [
+ [{"payload": {"data": ENCODED_SECRET}}] * 403 + [7],
+ {"payload": None},
+ {"payload": {"data": [ENCODED_SECRET] * 403 + [429]}},
+ {"payload": {"data": 429}},
+ {"payload": {}},
+ ],
+)
+@respx.mock
+def test_get_secret_returns_none_when_the_response_is_unusable_and_the_environment_has_no_value(body: object) -> None:
+ respx.get(ACCESS_URL).mock(return_value=_access_response(body))
+
+ assert get_secret(SECRET_NAME) is None