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