diff --git a/docs/my-website/docs/providers/watsonx/rerank.md b/docs/my-website/docs/providers/watsonx/rerank.md new file mode 100644 index 00000000000..0900ce96781 --- /dev/null +++ b/docs/my-website/docs/providers/watsonx/rerank.md @@ -0,0 +1,52 @@ +# watsonx.ai Rerank + +## Overview + +| Property | Details | +|----------|--------------------------------------------------------------------------| +| Description | watsonx.ai rerank integration | +| Provider Route on LiteLLM | `watsonx/` | +| Supported Operations | `/ml/v1/text/rerank` | +| Link to Provider Doc | [IBM WatsonX.ai ↗](https://cloud.ibm.com/apidocs/watsonx-ai#text-rerank) | + +## Quick Start + +### **LiteLLM SDK** + +```python +import os +from litellm import rerank + +os.environ["WATSONX_APIKEY"] = "YOUR_WATSONX_APIKEY" +os.environ["WATSONX_API_BASE"] = "YOUR_WATSONX_API_BASE" +os.environ["WATSONX_PROJECT_ID"] = "YOUR_WATSONX_PROJECT_ID" + +query="Best programming language for beginners?" +documents=[ + "Python is great for beginners due to simple syntax.", + "JavaScript runs in browsers and is versatile.", + "Rust has a steep learning curve but is very safe.", +] + +response = rerank( + model="watsonx/cross-encoder/ms-marco-minilm-l-12-v2", + query=query, + documents=documents, + top_n=2, + return_documents=True, +) + +print(response) +``` + +### **LiteLLM Proxy** + +```yaml +model_list: + - model_name: cross-encoder/ms-marco-minilm-l-12-v2 + litellm_params: + model: watsonx/cross-encoder/ms-marco-minilm-l-12-v2 + api_key: os.environ/WATSONX_APIKEY + api_base: os.environ/WATSONX_API_BASE + project_id: os.environ/WATSONX_PROJECT_ID +``` diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index a2371232302..b73b7741dd0 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -540,7 +540,7 @@ router_settings: | DEFAULT_IMAGE_WIDTH | Default width for images. Default is 300 | DEFAULT_IN_MEMORY_TTL | Default time-to-live for in-memory cache in seconds. Default is 5 | DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL | Default time-to-live in seconds for management objects (User, Team, Key, Organization) in memory cache. Default is 60 seconds. -| DEFAULT_MAX_LRU_CACHE_SIZE | Default maximum size for LRU cache. Default is 16 +| DEFAULT_MAX_LRU_CACHE_SIZE | Default maximum size for LRU cache. Default is 64 | DEFAULT_MAX_RECURSE_DEPTH | Default maximum recursion depth. Default is 100 | DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER | Default maximum recursion depth for sensitive data masker. Default is 10 | DEFAULT_MAX_RETRIES | Default maximum retry attempts. Default is 2 diff --git a/docs/my-website/docs/rerank.md b/docs/my-website/docs/rerank.md index 90f685d2bbd..9c76883d7fd 100644 --- a/docs/my-website/docs/rerank.md +++ b/docs/my-website/docs/rerank.md @@ -8,15 +8,15 @@ LiteLLM Follows the [cohere api request / response for the rerank api](https://c ## Overview -| Feature | Supported | Notes | -|---------|-----------|-------| -| Cost Tracking | ✅ | Works with all supported models | -| Logging | ✅ | Works across all integrations | -| End-user Tracking | ✅ | | -| Fallbacks | ✅ | Works between supported models | -| Loadbalancing | ✅ | Works between supported models | -| Guardrails | ✅ | Applies to input query only (not documents) | -| Supported Providers | Cohere, Together AI, Azure AI, DeepInfra, Nvidia NIM, Infinity, Fireworks AI, Voyage AI | | +| Feature | Supported | Notes | +|---------|-----------------------------------------------------------------------------------------------------|-------| +| Cost Tracking | ✅ | Works with all supported models | +| Logging | ✅ | Works across all integrations | +| End-user Tracking | ✅ | | +| Fallbacks | ✅ | Works between supported models | +| Loadbalancing | ✅ | Works between supported models | +| Guardrails | ✅ | Applies to input query only (not documents) | +| Supported Providers | Cohere, Together AI, Azure AI, DeepInfra, Nvidia NIM, Infinity, Fireworks AI, Voyage AI, watsonx.ai | | ## **LiteLLM Python SDK Usage** ### Quick Start @@ -123,17 +123,18 @@ curl http://0.0.0.0:4000/rerank \ #### ⚡️See all supported models and providers at [models.litellm.ai](https://models.litellm.ai/) -| Provider | Link to Usage | -|-------------|--------------------| -| Cohere (v1 + v2 clients) | [Usage](#quick-start) | -| Together AI| [Usage](../docs/providers/togetherai) | -| Azure AI| [Usage](../docs/providers/azure_ai#rerank-endpoint) | -| Jina AI| [Usage](../docs/providers/jina_ai) | -| AWS Bedrock| [Usage](../docs/providers/bedrock#rerank-api) | -| HuggingFace| [Usage](../docs/providers/huggingface_rerank) | -| Infinity| [Usage](../docs/providers/infinity) | -| vLLM| [Usage](../docs/providers/vllm#rerank-endpoint) | -| DeepInfra| [Usage](../docs/providers/deepinfra#rerank-endpoint) | -| Vertex AI| [Usage](../docs/providers/vertex#rerank-api) | -| Fireworks AI| [Usage](../docs/providers/fireworks_ai#rerank-endpoint) | -| Voyage AI| [Usage](../docs/providers/voyage#rerank) | \ No newline at end of file +| Provider | Link to Usage | +|--------------------------|------------------------------------------------------| +| Cohere (v1 + v2 clients) | [Usage](#quick-start) | +| Together AI | [Usage](../docs/providers/togetherai) | +| Azure AI | [Usage](../docs/providers/azure_ai#rerank-endpoint) | +| Jina AI | [Usage](../docs/providers/jina_ai) | +| AWS Bedrock | [Usage](../docs/providers/bedrock#rerank-api) | +| HuggingFace | [Usage](../docs/providers/huggingface_rerank) | +| Infinity | [Usage](../docs/providers/infinity) | +| vLLM | [Usage](../docs/providers/vllm#rerank-endpoint) | +| DeepInfra | [Usage](../docs/providers/deepinfra#rerank-endpoint) | +| Vertex AI | [Usage](../docs/providers/vertex#rerank-api) | +| Fireworks AI | [Usage](../docs/providers/fireworks_ai#rerank-endpoint) | +| Voyage AI | [Usage](../docs/providers/voyage#rerank) | +| IBM watsonx.ai | [Usage](../docs/providers/watsonx/rerank) | \ No newline at end of file diff --git a/litellm/__init__.py b/litellm/__init__.py index ea7bb6a78d8..a994db85b11 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1355,6 +1355,7 @@ if TYPE_CHECKING: from .llms.vertex_ai.rerank.transformation import VertexAIRerankConfig as VertexAIRerankConfig from .llms.fireworks_ai.rerank.transformation import FireworksAIRerankConfig as FireworksAIRerankConfig from .llms.voyage.rerank.transformation import VoyageRerankConfig as VoyageRerankConfig + from .llms.watsonx.rerank.transformation import IBMWatsonXRerankConfig as IBMWatsonXRerankConfig from .llms.clarifai.chat.transformation import ClarifaiConfig as ClarifaiConfig from .llms.ai21.chat.transformation import AI21ChatConfig as AI21ChatConfig from .llms.meta_llama.chat.transformation import LlamaAPIConfig as LlamaAPIConfig diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index cf446d15156..943acc6320f 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -155,6 +155,7 @@ LLM_CONFIG_NAMES = ( "VertexAIRerankConfig", "FireworksAIRerankConfig", "VoyageRerankConfig", + "IBMWatsonXRerankConfig", "ClarifaiConfig", "AI21ChatConfig", "LlamaAPIConfig", @@ -672,6 +673,7 @@ _LLM_CONFIGS_IMPORT_MAP = { "FireworksAIRerankConfig", ), "VoyageRerankConfig": (".llms.voyage.rerank.transformation", "VoyageRerankConfig"), + "IBMWatsonXRerankConfig": (".llms.watsonx.rerank.transformation", "IBMWatsonXRerankConfig"), "ClarifaiConfig": (".llms.clarifai.chat.transformation", "ClarifaiConfig"), "AI21ChatConfig": (".llms.ai21.chat.transformation", "AI21ChatConfig"), "LlamaAPIConfig": (".llms.meta_llama.chat.transformation", "LlamaAPIConfig"), diff --git a/litellm/constants.py b/litellm/constants.py index 3c11eb701dc..17ad742e419 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -287,7 +287,9 @@ MIN_NON_ZERO_TEMPERATURE = float(os.getenv("MIN_NON_ZERO_TEMPERATURE", 0.0001)) REPEATED_STREAMING_CHUNK_LIMIT = int( os.getenv("REPEATED_STREAMING_CHUNK_LIMIT", 100) ) # catch if model starts looping the same chunk while streaming. Uses high default to prevent false positives. -DEFAULT_MAX_LRU_CACHE_SIZE = int(os.getenv("DEFAULT_MAX_LRU_CACHE_SIZE", 16)) +# Shared maxsize for functools.lru_cache usage across hot paths. +# Defaulted to 64 to avoid cache thrash in multi-model production workloads. +DEFAULT_MAX_LRU_CACHE_SIZE = int(os.getenv("DEFAULT_MAX_LRU_CACHE_SIZE", 64)) _REALTIME_BODY_CACHE_SIZE = 1000 # Keep realtime helper caches bounded; workloads rarely exceed 1k models/intents INITIAL_RETRY_DELAY = float(os.getenv("INITIAL_RETRY_DELAY", 0.5)) MAX_RETRY_DELAY = float(os.getenv("MAX_RETRY_DELAY", 8.0)) @@ -576,6 +578,11 @@ OPENAI_CHAT_COMPLETION_PARAMS = [ "thinking", "web_search_options", "service_tier", + "store", + "prompt_cache_key", + "prompt_cache_retention", + "safety_identifier", + "verbosity", ] OPENAI_TRANSCRIPTION_PARAMS = [ diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index dae0bb1c2c0..02df747792d 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -448,7 +448,9 @@ def cost_per_token( # noqa: PLR0915 elif custom_llm_provider == "anthropic": return anthropic_cost_per_token(model=model, usage=usage_block) elif custom_llm_provider == "bedrock": - return bedrock_cost_per_token(model=model, usage=usage_block) + return bedrock_cost_per_token( + model=model, usage=usage_block, service_tier=service_tier + ) elif custom_llm_provider == "openai": return openai_cost_per_token( model=model, usage=usage_block, service_tier=service_tier @@ -2146,4 +2148,3 @@ def handle_realtime_stream_cost_calculation( return total_cost - diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 25ad0a570cb..a6e502a32b3 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -546,7 +546,11 @@ def convert_to_model_response_object( # noqa: PLR0915 message = litellm.Message(content=json_mode_content_str) finish_reason = "stop" if message is None: - provider_specific_fields = {} + # Preserve provider_specific_fields if already present + # in the response (e.g. from proxy passthrough) + provider_specific_fields = dict( + choice["message"].get("provider_specific_fields", None) or {} + ) message_keys = Message.model_fields.keys() for field in choice["message"].keys(): if field not in message_keys: diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index c665e084261..0cceddd9acf 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -22,6 +22,15 @@ from litellm.types.llms.anthropic import ( from litellm.types.llms.openai import AllMessageValues +def is_anthropic_oauth_key(value: Optional[str]) -> bool: + """Check if a value contains an Anthropic OAuth token (sk-ant-oat*).""" + if value is None: + return False + # Handle both raw token and "Bearer " format + if value.startswith("Bearer "): + value = value[7:] + return value.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX) + def optionally_handle_anthropic_oauth( headers: dict, api_key: Optional[str] ) -> tuple[dict, Optional[str]]: diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 304c707fa0b..dfaddb3c2b1 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -384,6 +384,14 @@ class BaseAWSLLM: model_id = BaseAWSLLM._get_model_id_from_model_with_spec( model_id, spec="moonshot" ) + elif "nova-2/" in model_id: + model_id = BaseAWSLLM._get_model_id_from_model_with_spec( + model_id, spec="nova-2" + ) + elif "nova/" in model_id: + model_id = BaseAWSLLM._get_model_id_from_model_with_spec( + model_id, spec="nova" + ) return model_id @staticmethod diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 25af852e09c..60a93b169c8 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -272,7 +272,18 @@ class BedrockConverseLLM(BaseAWSLLM): if unencoded_model_id is not None: modelId = self.encode_model_id(model_id=unencoded_model_id) else: - modelId = self.encode_model_id(model_id=model) + # Strip nova spec prefixes before encoding model ID for API URL + _model_for_id = model + _stripped = _model_for_id + for rp in ["bedrock/converse/", "bedrock/", "converse/"]: + if _stripped.startswith(rp): + _stripped = _stripped[len(rp):] + break + for _nova_prefix in ["nova-2/", "nova/"]: + if _stripped.startswith(_nova_prefix): + _model_for_id = _model_for_id.replace(_nova_prefix, "", 1) + break + modelId = self.encode_model_id(model_id=_model_for_id) fake_stream = litellm.AmazonConverseConfig().should_fake_stream( fake_stream=fake_stream, diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index cf2f9c9f774..4197fdac6c4 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -86,7 +86,7 @@ BEDROCK_COMPUTER_USE_TOOLS = [ UNSUPPORTED_BEDROCK_CONVERSE_BETA_PATTERNS = [ "advanced-tool-use", # Bedrock Converse doesn't support advanced-tool-use beta headers "prompt-caching", # Prompt caching not supported in Converse API - "compact-2026-01-12", # The compact beta feature is not currently supported on the Converse and ConverseStream APIs + "compact-2026-01-12", # The compact beta feature is not currently supported on the Converse and ConverseStream APIs ] # Models that support Bedrock's native structured outputs API (outputConfig.textFormat) @@ -299,45 +299,56 @@ class AmazonConverseConfig(BaseConfig): llm_provider="bedrock", ) - def _is_nova_lite_2_model(self, model: str) -> bool: + def _is_nova_2_model(self, model: str) -> bool: """ - Check if the model is a Nova Lite 2 model that supports reasoningConfig. + Check if the model is a Nova 2 model that supports reasoningConfig. - Nova Lite 2 models use a different reasoning configuration structure compared to + Nova 2 models use a different reasoning configuration structure compared to Anthropic's thinking parameter and GPT-OSS's reasoning_effort parameter. Supported models: - amazon.nova-2-lite-v1:0 + - amazon.nova-2-pro-preview-20251202-v1:0 - us.amazon.nova-2-lite-v1:0 - eu.amazon.nova-2-lite-v1:0 - apac.amazon.nova-2-lite-v1:0 + - (and other regional variants) Args: model: The model identifier Returns: - True if the model is a Nova Lite 2 model, False otherwise + True if the model is a Nova 2 model, False otherwise Examples: >>> config = AmazonConverseConfig() - >>> config._is_nova_lite_2_model("amazon.nova-2-lite-v1:0") + >>> config._is_nova_2_model("amazon.nova-2-lite-v1:0") True - >>> config._is_nova_lite_2_model("us.amazon.nova-2-lite-v1:0") + >>> config._is_nova_2_model("us.amazon.nova-2-lite-v1:0") True - >>> config._is_nova_lite_2_model("amazon.nova-pro-1-5-v1:0") + >>> config._is_nova_2_model("us.amazon.nova-2-pro-preview-20251202-v1:0") + True + >>> config._is_nova_2_model("amazon.nova-pro-1-5-v1:0") False - >>> config._is_nova_lite_2_model("amazon.nova-pro-v1:0") + >>> config._is_nova_2_model("amazon.nova-pro-v1:0") False """ - # Remove regional prefix if present (us., eu., apac.) + # Remove provider routing prefix if present (bedrock/converse/, bedrock/, converse/) model_without_region = model - for prefix in ["us.", "eu.", "apac."]: - if model.startswith(prefix): - model_without_region = model[len(prefix) :] + for routing_prefix in ["bedrock/converse/", "bedrock/", "converse/"]: + if model_without_region.startswith(routing_prefix): + model_without_region = model_without_region[len(routing_prefix) :] break - # Check if the model is specifically Nova Lite 2 - return "nova-2-lite" in model_without_region + # Remove regional prefix if present (us., eu., apac.) + for prefix in ["us.", "eu.", "apac."]: + if model_without_region.startswith(prefix): + model_without_region = model_without_region[len(prefix) :] + break + + # Check if the model is a Nova 2 model (matches nova-2-lite, nova-2-pro, etc.) + # Also check for nova-2/ spec prefix for imported models + return model_without_region.startswith("amazon.nova-2-") or model_without_region.startswith("nova-2/") def _map_web_search_options( self, web_search_options: dict, model: str @@ -425,7 +436,7 @@ class AmazonConverseConfig(BaseConfig): Different model families handle reasoning effort differently: - GPT-OSS models: Keep reasoning_effort as-is (passed to additionalModelRequestFields) - - Nova Lite 2 models: Transform to reasoningConfig structure + - Nova 2 models: Transform to reasoningConfig structure - Other models (Anthropic, etc.): Convert to thinking parameter Args: @@ -454,8 +465,8 @@ class AmazonConverseConfig(BaseConfig): # GPT-OSS models: keep reasoning_effort as-is # It will be passed through to additionalModelRequestFields optional_params["reasoning_effort"] = reasoning_effort - elif self._is_nova_lite_2_model(model): - # Nova Lite 2 models: transform to reasoningConfig + elif self._is_nova_2_model(model): + # Nova 2 models: transform to reasoningConfig reasoning_config = self._transform_reasoning_effort_to_reasoning_config( reasoning_effort ) @@ -509,6 +520,9 @@ class AmazonConverseConfig(BaseConfig): supported_params.append("tool_choice") supported_params.append("thinking") supported_params.append("reasoning_effort") + # For nova imported models, also add web_search_options + if "nova" in model.lower(): + supported_params.append("web_search_options") return supported_params ## Filter out 'cross-region' from model name @@ -543,8 +557,8 @@ class AmazonConverseConfig(BaseConfig): if "gpt-oss" in model: supported_params.append("reasoning_effort") - elif self._is_nova_lite_2_model(model): - # Nova Lite 2 models support reasoning_effort (transformed to reasoningConfig) + elif self._is_nova_2_model(model): + # Nova 2 models support reasoning_effort (transformed to reasoningConfig) # These models use a different reasoning structure than Anthropic's thinking parameter supported_params.append("reasoning_effort") elif ( @@ -929,8 +943,8 @@ class AmazonConverseConfig(BaseConfig): ) # Only update thinking tokens for non-GPT-OSS models and non-Nova-Lite-2 models - # Nova Lite 2 handles token budgeting differently through reasoningConfig - if "gpt-oss" not in model and not self._is_nova_lite_2_model(model): + # Nova 2 handles token budgeting differently through reasoningConfig + if "gpt-oss" not in model and not self._is_nova_2_model(model): self.update_optional_params_with_thinking_tokens( non_default_params=non_default_params, optional_params=optional_params ) @@ -1263,22 +1277,49 @@ class AmazonConverseConfig(BaseConfig): # "computer-use-2025-01-24" for Claude Sonnet 4.5, Haiku 4.5, Opus 4.1, Sonnet 4, Opus 4, and Sonnet 3.7 # "computer-use-2024-10-22" for older models model_lower = model.lower() - if "opus-4.6" in model_lower or "opus_4.6" in model_lower or "opus-4-6" in model_lower or "opus_4_6" in model_lower: + if ( + "opus-4.6" in model_lower + or "opus_4.6" in model_lower + or "opus-4-6" in model_lower + or "opus_4_6" in model_lower + ): computer_use_header = "computer-use-2025-11-24" - elif "opus-4.5" in model_lower or "opus_4.5" in model_lower or "opus-4-5" in model_lower or "opus_4_5" in model_lower: + elif ( + "opus-4.5" in model_lower + or "opus_4.5" in model_lower + or "opus-4-5" in model_lower + or "opus_4_5" in model_lower + ): computer_use_header = "computer-use-2025-11-24" - elif any(pattern in model_lower for pattern in [ - "sonnet-4.5", "sonnet_4.5", "sonnet-4-5", "sonnet_4_5", - "haiku-4.5", "haiku_4.5", "haiku-4-5", "haiku_4_5", - "opus-4.1", "opus_4.1", "opus-4-1", "opus_4_1", - "sonnet-4", "sonnet_4", - "opus-4", "opus_4", - "sonnet-3.7", "sonnet_3.7", "sonnet-3-7", "sonnet_3_7" - ]): + elif any( + pattern in model_lower + for pattern in [ + "sonnet-4.5", + "sonnet_4.5", + "sonnet-4-5", + "sonnet_4_5", + "haiku-4.5", + "haiku_4.5", + "haiku-4-5", + "haiku_4_5", + "opus-4.1", + "opus_4.1", + "opus-4-1", + "opus_4_1", + "sonnet-4", + "sonnet_4", + "opus-4", + "opus_4", + "sonnet-3.7", + "sonnet_3.7", + "sonnet-3-7", + "sonnet_3_7", + ] + ): computer_use_header = "computer-use-2025-01-24" else: computer_use_header = "computer-use-2024-10-22" - + anthropic_beta_list.append(computer_use_header) # Transform computer use tools to proper Bedrock format transformed_computer_tools = self._transform_computer_use_tools( @@ -1646,9 +1687,7 @@ class AmazonConverseConfig(BaseConfig): return message, returned_finish_reason - def _translate_message_content( - self, content_blocks: List[ContentBlock] - ) -> Tuple[ + def _translate_message_content(self, content_blocks: List[ContentBlock]) -> Tuple[ str, List[ChatCompletionToolCallChunk], Optional[List[BedrockConverseReasoningContentBlock]], @@ -1665,9 +1704,9 @@ class AmazonConverseConfig(BaseConfig): """ content_str = "" tools: List[ChatCompletionToolCallChunk] = [] - reasoningContentBlocks: Optional[ - List[BedrockConverseReasoningContentBlock] - ] = None + reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = ( + None + ) citationsContentBlocks: Optional[List[CitationsContentBlock]] = None for idx, content in enumerate(content_blocks): """ @@ -1794,9 +1833,9 @@ class AmazonConverseConfig(BaseConfig): chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"} content_str = "" tools: List[ChatCompletionToolCallChunk] = [] - reasoningContentBlocks: Optional[ - List[BedrockConverseReasoningContentBlock] - ] = None + reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = ( + None + ) citationsContentBlocks: Optional[List[CitationsContentBlock]] = None if message is not None: @@ -1815,17 +1854,17 @@ class AmazonConverseConfig(BaseConfig): provider_specific_fields["citationsContent"] = citationsContentBlocks if provider_specific_fields: - chat_completion_message[ - "provider_specific_fields" - ] = provider_specific_fields + chat_completion_message["provider_specific_fields"] = ( + provider_specific_fields + ) if reasoningContentBlocks is not None: - chat_completion_message[ - "reasoning_content" - ] = self._transform_reasoning_content(reasoningContentBlocks) - chat_completion_message[ - "thinking_blocks" - ] = self._transform_thinking_blocks(reasoningContentBlocks) + chat_completion_message["reasoning_content"] = ( + self._transform_reasoning_content(reasoningContentBlocks) + ) + chat_completion_message["thinking_blocks"] = ( + self._transform_thinking_blocks(reasoningContentBlocks) + ) chat_completion_message["content"] = content_str if ( json_mode is True diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 4c87f6fa994..8d40e349155 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -404,7 +404,7 @@ def extract_model_name_from_bedrock_arn(model: str) -> str: def strip_bedrock_routing_prefix(model: str) -> str: """Strip LiteLLM routing prefixes from model name.""" - for prefix in ["bedrock/", "converse/", "invoke/", "openai/"]: + for prefix in ["bedrock/", "converse/", "invoke/", "openai/", "nova-2/", "nova/"]: if model.startswith(prefix): model = model.split("/", 1)[1] return model @@ -427,7 +427,20 @@ def get_bedrock_base_model(model: str) -> str: - "us.meta.llama3-2-11b-instruct-v1:0" -> "meta.llama3-2-11b-instruct-v1" - "bedrock/converse/model" -> "model" - "anthropic.claude-3-5-sonnet-20241022-v2:0:51k" -> "anthropic.claude-3-5-sonnet-20241022-v2:0" + - "bedrock/nova-2/arn:aws:..." -> "amazon.nova-2-custom" + - "bedrock/nova/arn:aws:..." -> "amazon.nova-custom" """ + # Detect nova spec prefixes before stripping them + stripped = model + for rp in ["bedrock/converse/", "bedrock/", "converse/"]: + if stripped.startswith(rp): + stripped = stripped[len(rp):] + break + if stripped.startswith("nova-2/"): + return "amazon.nova-2-custom" + elif stripped.startswith("nova/"): + return "amazon.nova-custom" + model = strip_bedrock_routing_prefix(model) model = extract_model_name_from_bedrock_arn(model) model = strip_bedrock_throughput_suffix(model) @@ -594,6 +607,11 @@ class BedrockModelInfo(BaseLLMModelInfo): if prefix in model: return route_type + # Check for nova spec prefixes (nova/ and nova-2/) + _model_after_bedrock = model.replace("bedrock/", "", 1) + if _model_after_bedrock.startswith("nova-2/") or _model_after_bedrock.startswith("nova/"): + return "converse" + base_model = BedrockModelInfo.get_base_model(model) alt_model = BedrockModelInfo.get_non_litellm_routing_model_name(model=model) if ( diff --git a/litellm/llms/bedrock/cost_calculation.py b/litellm/llms/bedrock/cost_calculation.py index b20350d7325..ac99d4e36e7 100644 --- a/litellm/llms/bedrock/cost_calculation.py +++ b/litellm/llms/bedrock/cost_calculation.py @@ -3,7 +3,7 @@ Helper util for handling bedrock-specific cost calculation - e.g.: prompt caching """ -from typing import TYPE_CHECKING, Tuple +from typing import TYPE_CHECKING, Optional, Tuple from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token @@ -11,12 +11,17 @@ if TYPE_CHECKING: from litellm.types.utils import Usage -def cost_per_token(model: str, usage: "Usage") -> Tuple[float, float]: +def cost_per_token( + model: str, usage: "Usage", service_tier: Optional[str] = None +) -> Tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. Follows the same logic as Anthropic's cost per token calculation. """ return generic_cost_per_token( - model=model, usage=usage, custom_llm_provider="bedrock" - ) \ No newline at end of file + model=model, + usage=usage, + custom_llm_provider="bedrock", + service_tier=service_tier, + ) diff --git a/litellm/llms/watsonx/__init__.py b/litellm/llms/watsonx/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/watsonx/chat/__init__.py b/litellm/llms/watsonx/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/watsonx/completion/__init__.py b/litellm/llms/watsonx/completion/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/watsonx/embed/__init__.py b/litellm/llms/watsonx/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/watsonx/rerank/__init__.py b/litellm/llms/watsonx/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/watsonx/rerank/transformation.py b/litellm/llms/watsonx/rerank/transformation.py new file mode 100644 index 00000000000..7b4c2a07c3c --- /dev/null +++ b/litellm/llms/watsonx/rerank/transformation.py @@ -0,0 +1,204 @@ +""" +Transformation logic for IBM watsonx.ai's /ml/v1/text/rerank endpoint. + +Docs - https://cloud.ibm.com/apidocs/watsonx-ai#text-rerank +""" + +import uuid +from typing import Any, Dict, List, Optional, Union, cast + +import httpx + +from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj +from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.watsonx import ( + WatsonXAIEndpoint, +) +from litellm.types.rerank import ( + RerankResponse, + RerankResponseMeta, + RerankTokens, +) + +from ..common_utils import IBMWatsonXMixin, _generate_watsonx_token, _get_api_params + + +class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): + """ + IBM watsonx.ai Rerank API configuration + """ + + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: Optional[dict] = None, + ) -> str: + base_url = self._get_base_url(api_base=api_base) + endpoint = WatsonXAIEndpoint.RERANK.value + + url = base_url.rstrip("/") + endpoint + + params = optional_params or {} + + complete_url = self._add_api_version_to_url(url=url, api_version=(params.get("api_version", None))) + return complete_url + + def get_supported_cohere_rerank_params(self, model: str) -> list: + return [ + "query", + "documents", + "top_n", + "return_documents", + "max_tokens_per_doc", + ] + + def validate_environment( # type: ignore[override] + self, + headers: dict, + model: str, + api_key: Optional[str] = None, + optional_params: Optional[dict] = None, + ) -> Dict: + optional_params = optional_params or {} + + default_headers = { + "Content-Type": "application/json", + "Accept": "application/json", + } + + if "Authorization" in headers: + return {**default_headers, **headers} + token = cast( + Optional[str], + optional_params.pop("token", None) or get_secret_str("WATSONX_TOKEN"), + ) + zen_api_key = cast( + Optional[str], + optional_params.pop("zen_api_key", None) or get_secret_str("WATSONX_ZENAPIKEY"), + ) + if token: + headers["Authorization"] = f"Bearer {token}" + elif zen_api_key: + headers["Authorization"] = f"ZenApiKey {zen_api_key}" + else: + token = _generate_watsonx_token(api_key=api_key, token=token) + # build auth headers + headers["Authorization"] = f"Bearer {token}" + return {**default_headers, **headers} + + def map_cohere_rerank_params( + self, + non_default_params: Optional[dict], + model: str, + drop_params: bool, + query: str, + documents: List[Union[str, Dict[str, Any]]], + custom_llm_provider: Optional[str] = None, + top_n: Optional[int] = None, + rank_fields: Optional[List[str]] = None, + return_documents: Optional[bool] = True, + max_chunks_per_doc: Optional[int] = None, + max_tokens_per_doc: Optional[int] = None, + ) -> Dict: + """ + Map Cohere rerank params to IBM watsonx.ai rerank params + """ + optional_rerank_params = {} + if non_default_params is not None: + for k, v in non_default_params.items(): + if k == "query" and v is not None: + optional_rerank_params["query"] = v + elif k == "documents" and v is not None: + optional_rerank_params["inputs"] = [ + {"text": el} if isinstance(el, str) else el for el in v + ] + elif k == "top_n" and v is not None: + optional_rerank_params.setdefault("parameters", {}).setdefault("return_options", {})["top_n"] = v + elif k == "return_documents" and v is not None and isinstance(v, bool): + optional_rerank_params.setdefault("parameters", {}).setdefault("return_options", {})["inputs"] = v + elif k == "max_tokens_per_doc" and v is not None: + optional_rerank_params.setdefault("parameters", {})["truncate_input_tokens"] = v + + # IBM watsonx.ai require one of below parameters + elif k == "project_id" and v is not None: + optional_rerank_params["project_id"] = v + elif k == "space_id" and v is not None: + optional_rerank_params["space_id"] = v + + return dict(optional_rerank_params) + + def transform_rerank_request( + self, + model: str, + optional_rerank_params: Dict, + headers: dict, + ) -> dict: + """ + Transform request to IBM watsonx.ai rerank format + """ + watsonx_api_params = _get_api_params(params=optional_rerank_params, model=model) + watsonx_auth_payload = self._prepare_payload( + model=model, + api_params=watsonx_api_params, + ) + + return optional_rerank_params | watsonx_auth_payload + + def transform_rerank_response( + self, + model: str, + raw_response: httpx.Response, + model_response: RerankResponse, + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str] = None, + request_data: dict = {}, + optional_params: dict = {}, + litellm_params: dict = {}, + ) -> RerankResponse: + """ + Transform IBM watsonx.ai rerank response to LiteLLM RerankResponse format + """ + try: + raw_response_json = raw_response.json() + except Exception as e: + raise self.get_error_class( + error_message=f"Failed to parse response: {str(e)}", + status_code=raw_response.status_code, + headers=raw_response.headers, + ) + + _results: Optional[List[dict]] = raw_response_json.get("results") + if _results is None: + raise ValueError(f"No results found in the response={raw_response_json}") + + transformed_results = [] + + for result in _results: + transformed_result: Dict[str, Any] = { + "index": result["index"], + "relevance_score": result["score"], + } + + if "input" in result: + if isinstance(result["input"], str): + transformed_result["document"] = {"text": result["input"]} + else: + transformed_result["document"] = result["input"] + + transformed_results.append(transformed_result) + + response_id = raw_response_json.get("id") or raw_response_json.get("model_id") or str(uuid.uuid4()) + + # Extract usage information + _tokens = RerankTokens( + input_tokens=raw_response_json.get("input_token_count", 0), + ) + rerank_meta = RerankResponseMeta(tokens=_tokens) + + return RerankResponse( + id=response_id, + results=transformed_results, # type: ignore + meta=rerank_meta, + ) diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 71ae1348f39..32f209a763e 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -223,12 +223,14 @@ def get_known_models_from_wildcard( except ValueError: # safely fail return [] - if litellm_params is None: # need litellm params to extract litellm model name - return [] - - try: - provider = litellm_params.model.split("/", 1)[0] - except ValueError: + # Use provider from litellm_params when available, otherwise from wildcard prefix + # (e.g., "openai" from "openai/*" - needed for BYOK where wildcard isn't in router) + if litellm_params is not None: + try: + provider = litellm_params.model.split("/", 1)[0] + except ValueError: + provider = wildcard_provider_prefix + else: provider = wildcard_provider_prefix # get all known provider models @@ -282,7 +284,7 @@ def _get_wildcard_models( ## get litellm params from model if llm_router is not None: model_list = llm_router.get_model_list(model_name=model) - if model_list is not None: + if model_list: for router_model in model_list: wildcard_models = get_known_models_from_wildcard( wildcard_model=model, @@ -291,11 +293,22 @@ def _get_wildcard_models( ), ) all_wildcard_models.extend(wildcard_models) + else: + # Router has no deployment for this wildcard (e.g., BYOK team models) + # Fall back to expanding from known provider models + wildcard_models = get_known_models_from_wildcard( + wildcard_model=model, litellm_params=None + ) + if wildcard_models: + models_to_remove.add(model) + all_wildcard_models.extend(wildcard_models) else: # get all known provider models - wildcard_models = get_known_models_from_wildcard(wildcard_model=model) + wildcard_models = get_known_models_from_wildcard( + wildcard_model=model, litellm_params=None + ) - if wildcard_models is not None: + if wildcard_models: models_to_remove.add(model) all_wildcard_models.extend(wildcard_models) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4d77af513a8..3e8cc46521d 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -239,6 +239,8 @@ def clean_headers( """ Removes litellm api key from headers """ + from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + clean_headers = {} litellm_key_lower = ( litellm_key_header_name.lower() if litellm_key_header_name is not None else None @@ -246,8 +248,13 @@ def clean_headers( for header, value in headers.items(): header_lower = header.lower() + # Preserve Authorization header if it contains Anthropic OAuth token (sk-ant-oat*) + # This allows OAuth tokens to be forwarded to Anthropic-compatible providers + # via add_provider_specific_headers_to_request() + if header_lower == "authorization" and is_anthropic_oauth_key(value): + clean_headers[header] = value # Check if header should be excluded: either in special headers cache or matches custom litellm key - if header_lower not in _SPECIAL_HEADERS_CACHE and ( + elif header_lower not in _SPECIAL_HEADERS_CACHE and ( litellm_key_lower is None or header_lower != litellm_key_lower ): clean_headers[header] = value @@ -1717,6 +1724,8 @@ def add_provider_specific_headers_to_request( data: dict, headers: dict, ): + from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + anthropic_headers = {} # boolean to indicate if a header was added added_header = False @@ -1726,6 +1735,14 @@ def add_provider_specific_headers_to_request( anthropic_headers[header] = header_value added_header = True + # Check for Authorization header with Anthropic OAuth token (sk-ant-oat*) + # This needs to be handled via provider-specific headers to ensure it only + # goes to Anthropic-compatible providers, not all providers in the router + for header, value in headers.items(): + if header.lower() == "authorization" and is_anthropic_oauth_key(value): + anthropic_headers[header] = value + added_header = True + break if added_header is True: # Anthropic headers work across multiple providers # Store as comma-separated list so retrieval can match any of them diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 8910d37fbe7..f47fd6323f0 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -10,6 +10,7 @@ from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.together_ai.rerank.handler import TogetherAIRerank +from litellm.llms.watsonx.common_utils import IBMWatsonXMixin from litellm.rerank_api.rerank_utils import get_optional_rerank_params from litellm.secret_managers.main import get_secret, get_secret_str from litellm.types.rerank import RerankResponse @@ -29,7 +30,7 @@ async def arerank( model: str, query: str, documents: List[Union[str, Dict[str, Any]]], - custom_llm_provider: Optional[Literal["cohere", "together_ai", "deepinfra", "fireworks_ai", "voyage"]] = None, + custom_llm_provider: Optional[Literal["cohere", "together_ai", "deepinfra", "fireworks_ai", "voyage", "watsonx"]] = None, top_n: Optional[int] = None, rank_fields: Optional[List[str]] = None, return_documents: Optional[bool] = None, @@ -85,6 +86,7 @@ def rerank( # noqa: PLR0915 "deepinfra", "fireworks_ai", "voyage", + "watsonx", ] ] = None, top_n: Optional[int] = None, @@ -478,6 +480,31 @@ def rerank( # noqa: PLR0915 or get_secret_str("VOYAGE_API_BASE") ) + response = base_llm_http_handler.rerank( + model=model, + custom_llm_provider=_custom_llm_provider, + provider_config=rerank_provider_config, + optional_rerank_params=optional_rerank_params, + logging_obj=litellm_logging_obj, + timeout=optional_params.timeout, + api_key=api_key, + api_base=api_base, + _is_async=_is_async, + headers=headers or litellm.headers or {}, + client=client, + model_response=model_response, + ) + elif _custom_llm_provider == litellm.LlmProviders.WATSONX: + credentials = IBMWatsonXMixin.get_watsonx_credentials( + optional_params=dict(optional_params), api_key=dynamic_api_key, api_base=dynamic_api_base + ) + + api_key = credentials["api_key"] + api_base = credentials["api_base"] + + if credentials.get("token") is not None: + optional_rerank_params["token"] = credentials["token"] + response = base_llm_http_handler.rerank( model=model, custom_llm_provider=_custom_llm_provider, diff --git a/litellm/router.py b/litellm/router.py index 888c97ca0b1..9c821b11fbd 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7600,9 +7600,16 @@ class Router: Used by `.get_model_list` to get model list from model alias. """ returned_models: List[DeploymentTypedDict] = [] - for model_alias, model_value in self.model_group_alias.items(): - if model_name is not None and model_alias != model_name: - continue + + if model_name is not None: + # Fast path: direct dict lookup avoids scanning all aliases for non-alias model names. + if model_name not in self.model_group_alias: + return returned_models + alias_items = [(model_name, self.model_group_alias[model_name])] + else: + alias_items = list(self.model_group_alias.items()) + + for model_alias, model_value in alias_items: if isinstance(model_value, str): _router_model_name: str = model_value elif isinstance(model_value, dict): @@ -9099,4 +9106,3 @@ class Router: litellm._async_failure_callback = [] self.retry_policy = None self.flush_cache() - diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index bf3035fcc9f..70d4c6751db 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -335,13 +335,14 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): ): lowest_tpm = float("inf") potential_deployments = [] # if multiple deployments have the same low value + deployment_lookup = { + deployment.get("model_info", {}).get("id"): deployment + for deployment in healthy_deployments + } for item, item_tpm in all_deployments.items(): ## get the item from model list - _deployment = None item = item.split(":")[0] - for m in healthy_deployments: - if item == m["model_info"]["id"]: - _deployment = m + _deployment = deployment_lookup.get(item) if _deployment is None: continue # skip to next one elif item_tpm is None: diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 10acc343abd..3b0273f4c5d 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -58,7 +58,7 @@ def filter_team_based_models( request_team_id = metadata.get("user_api_key_team_id") or litellm_metadata.get( "user_api_key_team_id" ) - ids_to_remove = [] + ids_to_remove = set() if isinstance(healthy_deployments, dict): return healthy_deployments for deployment in healthy_deployments: @@ -67,7 +67,7 @@ def filter_team_based_models( if model_team_id is None: continue if model_team_id != request_team_id: - ids_to_remove.append(deployment.get("model_info", {}).get("id")) + ids_to_remove.add(_model_info.get("id")) return [ deployment @@ -125,4 +125,3 @@ def filter_web_search_deployments( if len(healthy_deployments) > 0 and len(final_deployments) == 0: verbose_logger.warning("No deployments support web search for request") return final_deployments - diff --git a/litellm/types/llms/watsonx.py b/litellm/types/llms/watsonx.py index 137090b032e..21e58500c6f 100644 --- a/litellm/types/llms/watsonx.py +++ b/litellm/types/llms/watsonx.py @@ -63,6 +63,7 @@ class WatsonXAIEndpoint(str, Enum): EMBEDDINGS = "/ml/v1/text/embeddings" PROMPTS = "/ml/v1/prompts" AVAILABLE_MODELS = "/ml/v1/foundation_model_specs" + RERANK = "/ml/v1/text/rerank" class WatsonXModelPattern(str, Enum): diff --git a/litellm/utils.py b/litellm/utils.py index 33f60f8d0a7..61402a1b30b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8151,6 +8151,8 @@ class ProviderConfigManager: return litellm.FireworksAIRerankConfig() elif litellm.LlmProviders.VOYAGE == provider: return litellm.VoyageRerankConfig() + elif litellm.LlmProviders.WATSONX == provider: + return litellm.IBMWatsonXRerankConfig() return litellm.CohereRerankConfig() @staticmethod diff --git a/tests/litellm/llms/bedrock/test_nova_imported_models.py b/tests/litellm/llms/bedrock/test_nova_imported_models.py new file mode 100644 index 00000000000..e3677aaf9e6 --- /dev/null +++ b/tests/litellm/llms/bedrock/test_nova_imported_models.py @@ -0,0 +1,92 @@ +""" +Tests for Nova imported/custom model support via spec prefixes (nova/, nova-2/). +""" + +import pytest + +from litellm.llms.bedrock.common_utils import ( + BedrockModelInfo, + get_bedrock_base_model, + strip_bedrock_routing_prefix, +) +from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + +NOVA_ARN = "arn:aws:bedrock:us-east-1:123456789012:custom-model-deployment/a1b2c3d4e5f6" +NOVA_MODEL = f"bedrock/nova/{NOVA_ARN}" +NOVA2_MODEL = f"bedrock/nova-2/{NOVA_ARN}" + + +class TestGetBedrockRoute: + def test_nova_prefix_routes_to_converse(self): + assert BedrockModelInfo.get_bedrock_route(NOVA_MODEL) == "converse" + + def test_nova2_prefix_routes_to_converse(self): + assert BedrockModelInfo.get_bedrock_route(NOVA2_MODEL) == "converse" + + def test_plain_arn_routes_to_invoke(self): + # Without spec prefix, ARN doesn't match converse models + result = BedrockModelInfo.get_bedrock_route(f"bedrock/{NOVA_ARN}") + assert result == "invoke" + + +class TestGetBedrockBaseModel: + def test_nova_prefix_returns_sentinel(self): + assert get_bedrock_base_model(f"nova/{NOVA_ARN}") == "amazon.nova-custom" + + def test_nova2_prefix_returns_sentinel(self): + assert get_bedrock_base_model(f"nova-2/{NOVA_ARN}") == "amazon.nova-2-custom" + + def test_bedrock_nova_prefix_returns_sentinel(self): + assert get_bedrock_base_model(NOVA_MODEL) == "amazon.nova-custom" + + def test_bedrock_nova2_prefix_returns_sentinel(self): + assert get_bedrock_base_model(NOVA2_MODEL) == "amazon.nova-2-custom" + + +class TestStripBedrockRoutingPrefix: + def test_strips_nova_prefix(self): + result = strip_bedrock_routing_prefix(f"nova/{NOVA_ARN}") + assert result == NOVA_ARN + + def test_strips_nova2_prefix(self): + result = strip_bedrock_routing_prefix(f"nova-2/{NOVA_ARN}") + assert result == NOVA_ARN + + +class TestIsNova2Model: + def setup_method(self): + self.config = AmazonConverseConfig() + + def test_standard_nova2_model(self): + assert self.config._is_nova_2_model("amazon.nova-2-lite-v1:0") is True + + def test_nova2_imported_model(self): + assert self.config._is_nova_2_model(NOVA2_MODEL) is True + + def test_nova_imported_model_is_not_nova2(self): + assert self.config._is_nova_2_model(NOVA_MODEL) is False + + def test_plain_nova_model(self): + assert self.config._is_nova_2_model("amazon.nova-pro-v1:0") is False + + +class TestGetSupportedOpenaiParams: + def setup_method(self): + self.config = AmazonConverseConfig() + + def test_nova_imported_has_tools_and_web_search(self): + params = self.config.get_supported_openai_params(NOVA_MODEL) + assert "tools" in params + assert "tool_choice" in params + assert "web_search_options" in params + + def test_nova2_imported_has_reasoning_effort(self): + params = self.config.get_supported_openai_params(NOVA2_MODEL) + assert "reasoning_effort" in params + assert "web_search_options" in params + + def test_nova2_imported_has_tools(self): + params = self.config.get_supported_openai_params(NOVA2_MODEL) + assert "tools" in params + assert "tool_choice" in params diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py index c151150f634..3b2087d25e9 100644 --- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py +++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py @@ -1037,6 +1037,193 @@ def test_convert_to_model_response_object_with_empty_dict_error(): assert result.choices[0].message.content == "Hello!" +def test_convert_to_model_response_object_preserves_provider_specific_fields_from_proxy(): + """ + Test that provider_specific_fields (e.g. Anthropic citations) are preserved + when the response already contains them (e.g. from a proxy passthrough). + + Regression test for https://github.com/BerriAI/litellm/issues/21153 + """ + citations = [ + [ + { + "type": "web_search_result_location", + "cited_text": "The Sony WH-1000XM5 remains one of the best...", + "url": "https://example.com/headphones-review", + "title": "Best Headphones 2025", + "supported_text": "Based on current reviews...", + } + ], + ] + web_search_results = [ + { + "url": "https://example.com/headphones-review", + "title": "Best Headphones 2025", + "snippet": "The Sony WH-1000XM5 remains one of the best...", + } + ] + + response_object = { + "id": "chatcmpl-proxy-123", + "object": "chat.completion", + "created": 1728933352, + "model": "anthropic/claude-opus-4-5-20251101", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Based on current reviews, the Sony WH-1000XM5 remains one of the best headphones.", + "tool_calls": [ + { + "id": "call_ws_123", + "type": "function", + "function": { + "name": "web_search", + "arguments": '{"query": "best headphones 2025"}', + }, + } + ], + "provider_specific_fields": { + "citations": citations, + "web_search_results": web_search_results, + }, + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 50, + "completion_tokens": 20, + "total_tokens": 70, + }, + } + + result = convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) + + assert isinstance(result, ModelResponse) + assert result.id == "chatcmpl-proxy-123" + + choice = result.choices[0] + assert choice.message.content == "Based on current reviews, the Sony WH-1000XM5 remains one of the best headphones." + assert choice.message.provider_specific_fields is not None + assert "citations" in choice.message.provider_specific_fields + assert choice.message.provider_specific_fields["citations"] == citations + assert "web_search_results" in choice.message.provider_specific_fields + assert choice.message.provider_specific_fields["web_search_results"] == web_search_results + + +def test_convert_to_model_response_object_provider_specific_fields_merges_extra_keys(): + """ + Test that provider_specific_fields from the response are merged with + any extra non-standard keys present in the message dict. + + Regression test for https://github.com/BerriAI/litellm/issues/21153 + """ + response_object = { + "id": "chatcmpl-merge-123", + "object": "chat.completion", + "created": 1728933352, + "model": "some-model", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello!", + "provider_specific_fields": { + "citations": [{"url": "https://example.com"}], + }, + "custom_extra_field": "extra_value", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + } + + result = convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) + + assert isinstance(result, ModelResponse) + psf = result.choices[0].message.provider_specific_fields + assert psf is not None + # Both the existing provider_specific_fields and the extra key should be present + assert "citations" in psf + assert psf["citations"] == [{"url": "https://example.com"}] + assert "custom_extra_field" in psf + assert psf["custom_extra_field"] == "extra_value" + + +def test_convert_to_model_response_object_no_provider_specific_fields_still_works(): + """ + Test that responses without provider_specific_fields continue to work as before. + + Ensures the fix for https://github.com/BerriAI/litellm/issues/21153 + doesn't break normal responses. + """ + response_object = { + "id": "chatcmpl-normal-123", + "object": "chat.completion", + "created": 1728933352, + "model": "gpt-4o", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello!", + "refusal": None, + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + } + + result = convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) + + assert isinstance(result, ModelResponse) + psf = result.choices[0].message.provider_specific_fields + # refusal is not a Message model field, so it should be in provider_specific_fields + assert psf is not None + assert "refusal" in psf + + def test_convert_to_model_response_object_with_error_code_only(): """ Test that errors with only a code (no message) are still treated as real errors. diff --git a/tests/llm_translation/test_optional_params.py b/tests/llm_translation/test_optional_params.py index 4699c31c378..6ecac7b36a2 100644 --- a/tests/llm_translation/test_optional_params.py +++ b/tests/llm_translation/test_optional_params.py @@ -1894,28 +1894,28 @@ def test_validate_openai_optional_params_stop_truncation(): result = validate_openai_optional_params(stop=stop_sequences) assert result == ["stop1", "stop2", "stop3", "stop4"] assert len(result) == 4 - + # Test with exactly 4 stop sequences - should not truncate stop_sequences_4 = ["stop1", "stop2", "stop3", "stop4"] result = validate_openai_optional_params(stop=stop_sequences_4) assert result == ["stop1", "stop2", "stop3", "stop4"] assert len(result) == 4 - + # Test with less than 4 stop sequences - should not truncate stop_sequences_2 = ["stop1", "stop2"] result = validate_openai_optional_params(stop=stop_sequences_2) assert result == ["stop1", "stop2"] assert len(result) == 2 - + # Test with single stop sequence as string - should return as is stop_string = "stop1" result = validate_openai_optional_params(stop=stop_string) assert result == "stop1" - + # Test with None - should return None result = validate_openai_optional_params(stop=None) assert result is None - + # Test with empty list - should return empty list result = validate_openai_optional_params(stop=[]) assert result == [] @@ -1928,7 +1928,7 @@ def test_validate_openai_optional_params_disable_stop_sequence_limit(): """ # Save original value original_value = litellm.disable_stop_sequence_limit - + try: # Test with disable_stop_sequence_limit = True - should NOT truncate litellm.disable_stop_sequence_limit = True @@ -1936,7 +1936,7 @@ def test_validate_openai_optional_params_disable_stop_sequence_limit(): result = validate_openai_optional_params(stop=stop_sequences) assert result == ["stop1", "stop2", "stop3", "stop4", "stop5", "stop6"] assert len(result) == 6 - + # Test with disable_stop_sequence_limit = False - should truncate to 4 litellm.disable_stop_sequence_limit = False stop_sequences = ["stop1", "stop2", "stop3", "stop4", "stop5", "stop6"] @@ -1965,19 +1965,83 @@ def test_validate_openai_optional_params_integration(): mock_response.usage.prompt_tokens = 10 mock_response.usage.completion_tokens = 5 mock_response.usage.total_tokens = 15 - - mock_client.return_value.chat.completions.create.return_value = mock_response - + + mock_client.return_value.chat.completions.create.return_value = ( + mock_response + ) + # Call completion with more than 4 stop sequences response = litellm.completion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello"}], stop=["stop1", "stop2", "stop3", "stop4", "stop5", "stop6"], - mock_response="Test response" # This will use mock + mock_response="Test response", # This will use mock ) - + # Verify the call was made (stop sequences should be truncated internally) assert response is not None except Exception as e: # Should not raise an exception pytest.fail(f"validate_openai_optional_params integration failed: {e}") + + +def test_drop_store_param_for_anthropic(): + """ + Test that the OpenAI-specific `store` parameter is correctly dropped + when calling Anthropic with drop_params=True. + + `store` is an OpenAI Chat Completion parameter (for storing completions + for distillation/evals) that Anthropic does not support. Without proper + handling, it leaks through to the Anthropic API and causes a + "store: Extra inputs are not permitted" error. + + Ref: https://github.com/BerriAI/litellm/issues/19700 + """ + optional_params = get_optional_params( + model="claude-sonnet-4-20250514", + custom_llm_provider="anthropic", + drop_params=True, + store=True, + ) + assert "store" not in optional_params + + +def test_additional_drop_params_store_for_anthropic(): + """ + Test that `additional_drop_params=["store"]` correctly strips the `store` + parameter for non-OpenAI providers like Anthropic. + + Ref: https://github.com/BerriAI/litellm/issues/19700 + """ + optional_params = get_optional_params( + model="claude-sonnet-4-20250514", + custom_llm_provider="anthropic", + additional_drop_params=["store"], + store=True, + ) + assert "store" not in optional_params + + +def test_store_in_openai_chat_completion_params(): + """ + Test that `store` is recognized as a standard OpenAI Chat Completion + parameter. This ensures it is correctly handled by helper functions + like `get_standard_openai_params()` and provider configs that rely on + `OPENAI_CHAT_COMPLETION_PARAMS`. + + Without `store` in this list, functions that filter by known OpenAI + params will silently drop it for OpenAI calls or incorrectly treat + it as a provider-specific param for non-OpenAI providers. + + Ref: https://github.com/BerriAI/litellm/issues/19700 + """ + from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS + + assert "store" in OPENAI_CHAT_COMPLETION_PARAMS + + # Verify get_standard_openai_params recognizes store + from litellm.utils import get_standard_openai_params + + result = get_standard_openai_params({"store": True, "temperature": 0.7}) + assert "store" in result + assert result["store"] is True diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index e34cad66ba6..1d9bb15b217 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -1928,6 +1928,24 @@ def test_get_known_models_from_wildcard( assert all(model in wildcard_models for model in expected_models) +def test_get_known_models_from_wildcard_without_litellm_params(): + """ + Test wildcard expansion without litellm_params (BYOK case - team has openai/* + but no deployment in router config). + """ + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + wildcard_models = get_known_models_from_wildcard( + wildcard_model="openai/*", litellm_params=None + ) + # Should return expanded OpenAI models (gpt-4o, gpt-4o-mini, etc.) + assert len(wildcard_models) > 0 + assert all(m.startswith("openai/") for m in wildcard_models) + # Check for common OpenAI models + model_ids = [m.split("/", 1)[1] for m in wildcard_models] + assert "gpt-4o" in model_ids or "gpt-3.5-turbo" in model_ids + + @pytest.mark.parametrize( "data, user_api_key_dict, expected_model", [ diff --git a/tests/router_unit_tests/test_get_model_list_alias_optimization.py b/tests/router_unit_tests/test_get_model_list_alias_optimization.py new file mode 100644 index 00000000000..31d992b6646 --- /dev/null +++ b/tests/router_unit_tests/test_get_model_list_alias_optimization.py @@ -0,0 +1,50 @@ +from litellm import Router + + +class NoItemsAliasDict(dict): + def items(self): + raise AssertionError("Unexpected full alias iteration via items()") + + +def test_get_model_list_from_model_alias_should_not_iterate_for_non_alias_lookup(): + router = Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo"}, + } + ], + model_group_alias={"alias-1": "gpt-4"}, + ) + router.model_group_alias = NoItemsAliasDict( + {f"alias-{idx}": "gpt-4" for idx in range(200)} + ) + + model_alias_list = router.get_model_list_from_model_alias( + model_name="gpt-3.5-turbo" + ) + assert model_alias_list == [] + + +def test_map_team_model_should_not_iterate_aliases_for_non_alias_team_model_name(): + router = Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo"}, + "model_info": { + "team_id": "team-1", + "team_public_model_name": "team-model", + }, + } + ], + model_group_alias={"alias-1": "gpt-4"}, + ) + router.model_group_alias = NoItemsAliasDict( + {f"alias-{idx}": "gpt-4" for idx in range(200)} + ) + + assert ( + router.map_team_model(team_model_name="team-model", team_id="team-1") + == "gpt-3.5-turbo" + ) diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index a321a24540f..ebffb56446e 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -283,3 +283,142 @@ class TestPassthroughOAuth: assert updated_headers["x-api-key"] == FAKE_REGULAR_KEY assert "authorization" not in updated_headers + + +class TestIsAnthropicOAuthKey: + """Tests for is_anthropic_oauth_key helper function.""" + + def test_oauth_token_raw(self): + """Raw OAuth token should be detected.""" + from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + + assert is_anthropic_oauth_key("sk-ant-oat01-abc123") is True + assert is_anthropic_oauth_key("sk-ant-oat02-xyz789") is True + + def test_oauth_token_bearer_format(self): + """Bearer-prefixed OAuth token should be detected.""" + from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + + assert is_anthropic_oauth_key("Bearer sk-ant-oat01-abc123") is True + assert is_anthropic_oauth_key("Bearer sk-ant-oat02-xyz789") is True + + def test_non_oauth_tokens(self): + """Non-OAuth values should return False.""" + from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + + assert is_anthropic_oauth_key(None) is False + assert is_anthropic_oauth_key("") is False + assert is_anthropic_oauth_key("sk-ant-api01-abc123") is False + assert is_anthropic_oauth_key("Bearer sk-ant-api01-abc123") is False + + def test_case_sensitivity(self): + """OAuth prefix matching should be case-sensitive.""" + from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + + assert is_anthropic_oauth_key("sk-ant-OAT01-abc123") is False + assert is_anthropic_oauth_key("SK-ANT-OAT01-abc123") is False + + def test_just_prefix(self): + """Just the prefix with no suffix should still match.""" + from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + + assert is_anthropic_oauth_key("sk-ant-oat") is True + + +class TestProxyOAuthHeaderForwarding: + """Tests for proxy-layer OAuth header preservation and forwarding.""" + + def test_clean_headers_preserves_oauth_authorization(self): + """clean_headers should preserve Authorization header with OAuth tokens.""" + from starlette.datastructures import Headers + + from litellm.proxy.litellm_pre_call_utils import clean_headers + + raw_headers = Headers( + raw=[ + (b"authorization", f"Bearer {FAKE_OAUTH_TOKEN}".encode()), + (b"content-type", b"application/json"), + ] + ) + cleaned = clean_headers(raw_headers) + + assert "authorization" in cleaned + assert cleaned["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + assert cleaned["content-type"] == "application/json" + + def test_clean_headers_strips_non_oauth_authorization(self): + """clean_headers should strip Authorization header with regular API keys.""" + from starlette.datastructures import Headers + + from litellm.proxy.litellm_pre_call_utils import clean_headers + + raw_headers = Headers( + raw=[ + (b"authorization", b"Bearer sk-regular-key-123"), + (b"content-type", b"application/json"), + ] + ) + cleaned = clean_headers(raw_headers) + + assert "authorization" not in cleaned + assert cleaned["content-type"] == "application/json" + + def test_add_provider_specific_headers_forwards_oauth(self): + """add_provider_specific_headers_to_request should forward OAuth Authorization + as a ProviderSpecificHeader scoped to Anthropic-compatible providers.""" + from litellm.proxy.litellm_pre_call_utils import ( + add_provider_specific_headers_to_request, + ) + + data: dict = {} + headers = { + "authorization": f"Bearer {FAKE_OAUTH_TOKEN}", + "content-type": "application/json", + } + + add_provider_specific_headers_to_request(data=data, headers=headers) + + assert "provider_specific_header" in data + psh = data["provider_specific_header"] + assert "anthropic" in psh["custom_llm_provider"] + assert "bedrock" in psh["custom_llm_provider"] + assert "vertex_ai" in psh["custom_llm_provider"] + assert psh["extra_headers"]["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + + def test_add_provider_specific_headers_ignores_non_oauth(self): + """add_provider_specific_headers_to_request should not create a + ProviderSpecificHeader for non-OAuth Authorization headers.""" + from litellm.proxy.litellm_pre_call_utils import ( + add_provider_specific_headers_to_request, + ) + + data: dict = {} + headers = { + "authorization": "Bearer sk-regular-key-123", + "content-type": "application/json", + } + + add_provider_specific_headers_to_request(data=data, headers=headers) + + assert "provider_specific_header" not in data + + def test_add_provider_specific_headers_combines_anthropic_and_oauth(self): + """When both anthropic-beta and OAuth Authorization are present, both + should be included in the ProviderSpecificHeader.""" + from litellm.proxy.litellm_pre_call_utils import ( + add_provider_specific_headers_to_request, + ) + + data: dict = {} + headers = { + "authorization": f"Bearer {FAKE_OAUTH_TOKEN}", + "anthropic-beta": "oauth-2025-04-20", + "content-type": "application/json", + } + + add_provider_specific_headers_to_request(data=data, headers=headers) + + assert "provider_specific_header" in data + psh = data["provider_specific_header"] + assert psh["extra_headers"]["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + assert psh["extra_headers"]["anthropic-beta"] == "oauth-2025-04-20" diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 018af654167..c773db21074 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -2701,37 +2701,37 @@ def test_empty_assistant_message_handling(): assert result[1]["content"][0]["text"] == "I'm doing well, thank you!" -def test_is_nova_lite_2_model(): - """Test the _is_nova_lite_2_model() method for detecting Nova 2 models.""" +def test_is_nova_2_model(): + """Test the _is_nova_2_model() method for detecting Nova 2 models.""" config = AmazonConverseConfig() # Test with amazon.nova-2-lite-v1:0 - assert config._is_nova_lite_2_model("amazon.nova-2-lite-v1:0") is True + assert config._is_nova_2_model("amazon.nova-2-lite-v1:0") is True # Test with regional variants - assert config._is_nova_lite_2_model("us.amazon.nova-2-lite-v1:0") is True - assert config._is_nova_lite_2_model("eu.amazon.nova-2-lite-v1:0") is True - assert config._is_nova_lite_2_model("apac.amazon.nova-2-lite-v1:0") is True + assert config._is_nova_2_model("us.amazon.nova-2-lite-v1:0") is True + assert config._is_nova_2_model("eu.amazon.nova-2-lite-v1:0") is True + assert config._is_nova_2_model("apac.amazon.nova-2-lite-v1:0") is True # Test with other Nova 2 variants (pro, micro) - assert config._is_nova_lite_2_model("amazon.nova-pro-1-5-v1:0") is False - assert config._is_nova_lite_2_model("amazon.nova-micro-1-5-v1:0") is False - assert config._is_nova_lite_2_model("us.amazon.nova-pro-1-5-v1:0") is False - assert config._is_nova_lite_2_model("eu.amazon.nova-micro-1-5-v1:0") is False + assert config._is_nova_2_model("amazon.nova-pro-1-5-v1:0") is False + assert config._is_nova_2_model("amazon.nova-micro-1-5-v1:0") is False + assert config._is_nova_2_model("us.amazon.nova-pro-1-5-v1:0") is False + assert config._is_nova_2_model("eu.amazon.nova-micro-1-5-v1:0") is False # Test with non-Nova-1.5 lite models (should return False) - assert config._is_nova_lite_2_model("amazon.nova-lite-v1:0") is False - assert config._is_nova_lite_2_model("amazon.nova-pro-v1:0") is False - assert config._is_nova_lite_2_model("amazon.nova-micro-v1:0") is False + assert config._is_nova_2_model("amazon.nova-lite-v1:0") is False + assert config._is_nova_2_model("amazon.nova-pro-v1:0") is False + assert config._is_nova_2_model("amazon.nova-micro-v1:0") is False # Test with Nova v1:0 models (should return False) - assert config._is_nova_lite_2_model("us.amazon.nova-lite-v1:0") is False - assert config._is_nova_lite_2_model("eu.amazon.nova-pro-v1:0") is False + assert config._is_nova_2_model("us.amazon.nova-lite-v1:0") is False + assert config._is_nova_2_model("eu.amazon.nova-pro-v1:0") is False # Test with completely different models (should return False) - assert config._is_nova_lite_2_model("anthropic.claude-3-5-sonnet-20240620-v1:0") is False - assert config._is_nova_lite_2_model("meta.llama3-70b-instruct-v1:0") is False - assert config._is_nova_lite_2_model("mistral.mistral-7b-instruct-v0:2") is False + assert config._is_nova_2_model("anthropic.claude-3-5-sonnet-20240620-v1:0") is False + assert config._is_nova_2_model("meta.llama3-70b-instruct-v1:0") is False + assert config._is_nova_2_model("mistral.mistral-7b-instruct-v0:2") is False def test_thinking_with_max_completion_tokens(): diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation_nova_2.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation_nova_2.py index 23243dac201..bac7aa08a04 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation_nova_2.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation_nova_2.py @@ -1,7 +1,10 @@ """ Unit tests for Amazon Nova 2 reasoning configuration transformation. -Tests the _transform_reasoning_effort_to_reasoning_config method in AmazonConverseConfig. +Tests request transformation, response parsing, multi-turn message translation, +and model detection for Nova 2 Lite and Nova 2 Pro via the Bedrock Converse API. + +Reference: https://docs.aws.amazon.com/nova/latest/nova2-userguide/using-converse-api.html """ import pytest @@ -12,6 +15,7 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path +import httpx import litellm from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig @@ -323,248 +327,52 @@ class TestNova15SupportedParameters: assert "response_format" in supported_params -class TestNova15ResponseParsing: - """Test suite for Nova 2 response parsing.""" +class TestNova2ResponseParsing: + """Test that reasoningContent blocks are parsed into reasoning_content strings.""" - def test_transform_reasoning_content_single_block(self): - """Test that reasoning content is extracted correctly from a single block.""" + def test_should_extract_single_reasoning_block(self): config = AmazonConverseConfig() - - reasoning_blocks = [ - {"reasoningText": {"text": "Let me think through this step by step..."}} - ] - - result = config._transform_reasoning_content(reasoning_blocks) - + result = config._transform_reasoning_content( + [{"reasoningText": {"text": "Let me think through this step by step..."}}] + ) assert result == "Let me think through this step by step..." - def test_transform_reasoning_content_multiple_blocks(self): - """Test that reasoning content is concatenated from multiple blocks.""" + def test_should_concatenate_multiple_reasoning_blocks(self): config = AmazonConverseConfig() - - reasoning_blocks = [ - {"reasoningText": {"text": "First, I need to analyze the problem. "}}, - {"reasoningText": {"text": "Then, I'll consider the solution."}}, - ] - - result = config._transform_reasoning_content(reasoning_blocks) - + result = config._transform_reasoning_content( + [ + {"reasoningText": {"text": "First, I need to analyze the problem. "}}, + {"reasoningText": {"text": "Then, I'll consider the solution."}}, + ] + ) assert ( result == "First, I need to analyze the problem. Then, I'll consider the solution." ) - def test_transform_reasoning_content_empty_blocks(self): - """Test that empty reasoning blocks return empty string.""" + def test_should_return_empty_string_for_empty_blocks(self): config = AmazonConverseConfig() - - reasoning_blocks = [] - - result = config._transform_reasoning_content(reasoning_blocks) - - assert result == "" - - def test_transform_thinking_blocks_with_text(self): - """Test that thinking blocks are populated correctly with text.""" - config = AmazonConverseConfig() - - reasoning_blocks = [{"reasoningText": {"text": "My reasoning process..."}}] - - result = config._transform_thinking_blocks(reasoning_blocks) - - assert len(result) == 1 - assert result[0]["type"] == "thinking" - assert result[0]["thinking"] == "My reasoning process..." - assert "signature" not in result[0] - - def test_transform_thinking_blocks_with_signature(self): - """Test that signature field is preserved when present.""" - config = AmazonConverseConfig() - - reasoning_blocks = [ - { - "reasoningText": { - "text": "My reasoning...", - "signature": "signature-hash-12345", - } - } - ] - - result = config._transform_thinking_blocks(reasoning_blocks) - - assert len(result) == 1 - assert result[0]["type"] == "thinking" - assert result[0]["thinking"] == "My reasoning..." - assert result[0]["signature"] == "signature-hash-12345" - - def test_transform_thinking_blocks_with_redacted_content(self): - """Test that redacted content blocks are handled correctly.""" - config = AmazonConverseConfig() - - reasoning_blocks = [ - {"reasoningText": {"text": "First part of reasoning..."}}, - {"redactedContent": {}}, - {"reasoningText": {"text": "Second part after redaction..."}}, - ] - - result = config._transform_thinking_blocks(reasoning_blocks) - - assert len(result) == 3 - assert result[0]["type"] == "thinking" - assert result[0]["thinking"] == "First part of reasoning..." - assert result[1]["type"] == "redacted_thinking" - assert result[2]["type"] == "thinking" - assert result[2]["thinking"] == "Second part after redaction..." - - def test_transform_thinking_blocks_multiple_blocks(self): - """Test that multiple thinking blocks are all transformed.""" - config = AmazonConverseConfig() - - reasoning_blocks = [ - {"reasoningText": {"text": "Step 1: Analyze the problem"}}, - { - "reasoningText": { - "text": "Step 2: Consider solutions", - "signature": "sig-abc", - } - }, - {"reasoningText": {"text": "Step 3: Choose best approach"}}, - ] - - result = config._transform_thinking_blocks(reasoning_blocks) - - assert len(result) == 3 - assert all(block["type"] == "thinking" for block in result) - assert result[0]["thinking"] == "Step 1: Analyze the problem" - assert result[1]["thinking"] == "Step 2: Consider solutions" - assert result[1]["signature"] == "sig-abc" - assert result[2]["thinking"] == "Step 3: Choose best approach" - - def test_transform_thinking_blocks_empty_list(self): - """Test that empty thinking blocks list returns empty list.""" - config = AmazonConverseConfig() - - reasoning_blocks = [] - - result = config._transform_thinking_blocks(reasoning_blocks) - - assert result == [] - - def test_response_parsing_integration(self): - """Test that response parsing works end-to-end with Nova 2 structure.""" - config = AmazonConverseConfig() - - # Simulate a Nova 2 response with reasoning content - reasoning_blocks = [ - { - "reasoningText": { - "text": "Let me analyze this carefully. ", - "signature": "test-signature", - } - }, - {"reasoningText": {"text": "Based on my analysis, the answer is clear."}}, - ] - - # Test reasoning content extraction - reasoning_content = config._transform_reasoning_content(reasoning_blocks) - assert ( - reasoning_content - == "Let me analyze this carefully. Based on my analysis, the answer is clear." - ) - - # Test thinking blocks transformation - thinking_blocks = config._transform_thinking_blocks(reasoning_blocks) - assert len(thinking_blocks) == 2 - assert thinking_blocks[0]["thinking"] == "Let me analyze this carefully. " - assert thinking_blocks[0]["signature"] == "test-signature" - assert ( - thinking_blocks[1]["thinking"] - == "Based on my analysis, the answer is clear." - ) + assert config._transform_reasoning_content([]) == "" -class TestNova15StreamingResponseParsing: - """Test suite for Nova 2 streaming response parsing.""" +class TestNova2StreamingResponseParsing: + """Test that streaming reasoningContent deltas produce reasoning_content on the delta.""" - def test_streaming_reasoning_content_start_event(self): - """Test that streaming start event with reasoningContent is handled correctly.""" + def test_should_extract_reasoning_content_from_delta(self): from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") - - # Simulate a start event with redacted reasoning content - chunk_data = { - "start": {"reasoningContent": {"redactedContent": {}}}, - "contentBlockIndex": 0, - } - - result = handler.converse_chunk_parser(chunk_data) - - # Verify thinking blocks are populated - assert result.choices[0].delta.thinking_blocks is not None - assert len(result.choices[0].delta.thinking_blocks) == 1 - assert result.choices[0].delta.thinking_blocks[0]["type"] == "redacted_thinking" - - def test_streaming_reasoning_content_delta_text(self): - """Test that streaming delta event with reasoning text is handled correctly.""" - from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder - - handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") - - # Simulate a delta event with reasoning text chunk_data = { "delta": {"reasoningContent": {"text": "Let me think about this..."}}, "contentBlockIndex": 0, } - result = handler.converse_chunk_parser(chunk_data) - - # Verify reasoning content is extracted assert result.choices[0].delta.reasoning_content == "Let me think about this..." - # Verify thinking blocks are populated - assert result.choices[0].delta.thinking_blocks is not None - assert len(result.choices[0].delta.thinking_blocks) == 1 - assert result.choices[0].delta.thinking_blocks[0]["type"] == "thinking" - assert ( - result.choices[0].delta.thinking_blocks[0]["thinking"] - == "Let me think about this..." - ) - - def test_streaming_reasoning_content_delta_signature(self): - """Test that streaming delta event with signature is handled correctly.""" + def test_should_accumulate_multiple_reasoning_deltas(self): from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") - - # Simulate a delta event with signature - chunk_data = { - "delta": {"reasoningContent": {"signature": "signature-hash-xyz"}}, - "contentBlockIndex": 0, - } - - result = handler.converse_chunk_parser(chunk_data) - - # Verify reasoning content is set to empty string for consistency - assert result.choices[0].delta.reasoning_content == "" - - # Verify thinking blocks are populated with signature - assert result.choices[0].delta.thinking_blocks is not None - assert len(result.choices[0].delta.thinking_blocks) == 1 - assert result.choices[0].delta.thinking_blocks[0]["type"] == "thinking" - assert ( - result.choices[0].delta.thinking_blocks[0]["signature"] - == "signature-hash-xyz" - ) - assert result.choices[0].delta.thinking_blocks[0]["thinking"] == "" - - def test_streaming_reasoning_content_multiple_deltas(self): - """Test that multiple reasoning content deltas are accumulated correctly.""" - from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder - - handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") - - # Simulate multiple delta events chunks = [ { "delta": {"reasoningContent": {"text": "First, "}}, @@ -579,30 +387,15 @@ class TestNova15StreamingResponseParsing: "contentBlockIndex": 0, }, ] - - results = [] - for chunk_data in chunks: - result = handler.converse_chunk_parser(chunk_data) - results.append(result) - - # Verify each delta has the correct reasoning content + results = [handler.converse_chunk_parser(c) for c in chunks] assert results[0].choices[0].delta.reasoning_content == "First, " assert results[1].choices[0].delta.reasoning_content == "I need to analyze " assert results[2].choices[0].delta.reasoning_content == "the problem." - # Verify thinking blocks are populated for each delta - for result in results: - assert result.choices[0].delta.thinking_blocks is not None - assert len(result.choices[0].delta.thinking_blocks) == 1 - assert result.choices[0].delta.thinking_blocks[0]["type"] == "thinking" - - def test_streaming_reasoning_then_text_content(self): - """Test that reasoning content followed by text content is handled correctly.""" + def test_should_stream_reasoning_then_text(self): from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") - - # Simulate reasoning content followed by text content chunks = [ { "delta": {"reasoningContent": {"text": "Let me think..."}}, @@ -611,184 +404,286 @@ class TestNova15StreamingResponseParsing: {"delta": {"text": "Based on my reasoning, "}, "contentBlockIndex": 1}, {"delta": {"text": "the answer is 42."}, "contentBlockIndex": 1}, ] - - results = [] - for chunk_data in chunks: - result = handler.converse_chunk_parser(chunk_data) - results.append(result) - - # Verify first chunk has reasoning content + results = [handler.converse_chunk_parser(c) for c in chunks] assert results[0].choices[0].delta.reasoning_content == "Let me think..." - assert results[0].choices[0].delta.thinking_blocks is not None - - # Verify subsequent chunks have text content assert results[1].choices[0].delta.content == "Based on my reasoning, " assert results[2].choices[0].delta.content == "the answer is 42." - def test_streaming_redacted_content_delta(self): - """Test that streaming delta with redacted content is handled correctly.""" + def test_should_populate_provider_specific_fields(self): from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") - - # Simulate a delta event with redacted content - chunk_data = { - "delta": {"reasoningContent": {"redactedContent": {}}}, - "contentBlockIndex": 0, - } - - result = handler.converse_chunk_parser(chunk_data) - - # Verify reasoning content is set to empty string for consistency - assert result.choices[0].delta.reasoning_content == "" - - # Verify thinking blocks contain redacted block - assert result.choices[0].delta.thinking_blocks is not None - assert len(result.choices[0].delta.thinking_blocks) == 1 - assert result.choices[0].delta.thinking_blocks[0]["type"] == "redacted_thinking" - - def test_streaming_provider_specific_fields(self): - """Test that provider_specific_fields are populated in streaming responses.""" - from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder - - handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") - - # Simulate a delta event with reasoning content chunk_data = { "delta": {"reasoningContent": {"text": "Reasoning text"}}, "contentBlockIndex": 0, } - result = handler.converse_chunk_parser(chunk_data) + psf = result.choices[0].delta.provider_specific_fields + assert psf is not None + assert psf["reasoningContent"]["text"] == "Reasoning text" - # Verify provider_specific_fields are populated - assert result.choices[0].delta.provider_specific_fields is not None - assert "reasoningContent" in result.choices[0].delta.provider_specific_fields - assert ( - result.choices[0].delta.provider_specific_fields["reasoningContent"]["text"] - == "Reasoning text" - ) - - def test_streaming_mixed_content_blocks(self): - """Test streaming with mixed content blocks (reasoning, text, tool calls).""" + def test_should_stream_reasoning_with_tool_calls(self): from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") - - # Simulate a complex streaming scenario chunks = [ - # Start with reasoning { - "delta": { - "reasoningContent": { - "text": "I need to call a tool to get information." - } - }, + "delta": {"reasoningContent": {"text": "I need to call a tool."}}, "contentBlockIndex": 0, }, - # Tool use start { "start": {"toolUse": {"toolUseId": "tool-123", "name": "get_weather"}}, "contentBlockIndex": 1, }, - # Tool use delta { "delta": {"toolUse": {"input": '{"location": "NYC"}'}}, "contentBlockIndex": 1, }, - # Text response {"delta": {"text": "The weather is sunny."}, "contentBlockIndex": 2}, ] - - results = [] - for chunk_data in chunks: - result = handler.converse_chunk_parser(chunk_data) - results.append(result) - - # Verify reasoning content in first chunk - assert ( - results[0].choices[0].delta.reasoning_content - == "I need to call a tool to get information." - ) - - # Verify tool call in second and third chunks - assert results[1].choices[0].delta.tool_calls is not None + results = [handler.converse_chunk_parser(c) for c in chunks] + assert results[0].choices[0].delta.reasoning_content == "I need to call a tool." assert ( results[1].choices[0].delta.tool_calls[0]["function"]["name"] == "get_weather" ) - assert results[2].choices[0].delta.tool_calls is not None - - # Verify text content in fourth chunk assert results[3].choices[0].delta.content == "The weather is sunny." - def test_extract_reasoning_content_str_with_text(self): - """Test extract_reasoning_content_str method with text.""" - from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder - handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") +# --------------------------------------------------------------------------- +# Model detection — _is_nova_2_model covers both Lite and Pro +# --------------------------------------------------------------------------- - reasoning_block = {"text": "This is reasoning text"} +NOVA_2_LITE = "amazon.nova-2-lite-v1:0" +NOVA_2_PRO = "us.amazon.nova-2-pro-preview-20251202-v1:0" - result = handler.extract_reasoning_content_str(reasoning_block) - assert result == "This is reasoning text" +class TestNova2ModelDetection: + """Verify _is_nova_2_model identifies all Nova 2 variants (lite, pro, regional, routed).""" - def test_extract_reasoning_content_str_without_text(self): - """Test extract_reasoning_content_str method without text (e.g., signature only).""" - from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + @pytest.mark.parametrize( + "model", + [ + "amazon.nova-2-lite-v1:0", + "amazon.nova-2-pro-preview-20251202-v1:0", + "us.amazon.nova-2-lite-v1:0", + "us.amazon.nova-2-pro-preview-20251202-v1:0", + "eu.amazon.nova-2-lite-v1:0", + "apac.amazon.nova-2-pro-preview-20251202-v1:0", + "bedrock/converse/amazon.nova-2-lite-v1:0", + "bedrock/converse/us.amazon.nova-2-pro-preview-20251202-v1:0", + "bedrock/amazon.nova-2-lite-v1:0", + "converse/us.amazon.nova-2-lite-v1:0", + "converse/amazon.nova-2-pro-preview-20251202-v1:0", + ], + ) + def test_should_recognize_nova_2_models(self, model): + assert AmazonConverseConfig()._is_nova_2_model(model) is True - handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") + @pytest.mark.parametrize( + "model", + [ + "amazon.nova-pro-v1:0", + "amazon.nova-lite-v1:0", + "amazon.nova-pro-1-5-v1:0", + "anthropic.claude-3-sonnet-20240229-v1:0", + "us.amazon.nova-pro-v1:0", + ], + ) + def test_should_not_match_non_nova_2_models(self, model): + assert AmazonConverseConfig()._is_nova_2_model(model) is False - reasoning_block = {"signature": "sig-123"} - result = handler.extract_reasoning_content_str(reasoning_block) +# --------------------------------------------------------------------------- +# End-to-end request body — reasoningConfig in additionalModelRequestFields +# --------------------------------------------------------------------------- - assert result is None - def test_translate_thinking_blocks_streaming_text(self): - """Test translate_thinking_blocks method with text.""" - from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder +class TestNova2EndToEndRequest: + """Verify transform_request places reasoningConfig correctly for both model variants.""" - handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") + def _build_request(self, model, effort, **extra): + config = AmazonConverseConfig() + optional_params = config.map_openai_params( + non_default_params={"reasoning_effort": effort, **extra}, + optional_params={}, + model=model, + drop_params=False, + ) + return config.transform_request( + model=model, + messages=[{"role": "user", "content": "What is 2+2?"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) - thinking_block = {"text": "Thinking content"} + @pytest.mark.parametrize("model", [NOVA_2_LITE, NOVA_2_PRO]) + def test_should_place_reasoning_config_in_additional_model_request_fields( + self, model + ): + body = self._build_request(model, "high") + additional = body.get("additionalModelRequestFields", {}) + assert additional["reasoningConfig"] == { + "type": "enabled", + "maxReasoningEffort": "high", + } + assert "reasoningConfig" not in body # not top-level + assert "thinking" not in body # not Anthropic-style - result = handler.translate_thinking_blocks(thinking_block) - - assert result is not None - assert len(result) == 1 - assert result[0]["type"] == "thinking" - assert result[0]["thinking"] == "Thinking content" - - def test_translate_thinking_blocks_streaming_signature(self): - """Test translate_thinking_blocks method with signature.""" - from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder - - handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") - - thinking_block = {"signature": "sig-abc"} - - result = handler.translate_thinking_blocks(thinking_block) - - assert result is not None - assert len(result) == 1 - assert result[0]["type"] == "thinking" - assert result[0]["signature"] == "sig-abc" + @pytest.mark.parametrize("model", [NOVA_2_LITE, NOVA_2_PRO]) + def test_should_coexist_with_inference_params(self, model): + body = self._build_request(model, "high", temperature=0.5, max_tokens=512) assert ( - result[0]["thinking"] == "" - ) # Empty string for consistency with Anthropic + body["additionalModelRequestFields"]["reasoningConfig"]["type"] == "enabled" + ) + inf = body.get("inferenceConfig", {}) + assert inf.get("temperature") == 0.5 + assert inf.get("maxTokens") == 512 - def test_translate_thinking_blocks_streaming_redacted(self): - """Test translate_thinking_blocks method with redacted content.""" - from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder - handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") +# --------------------------------------------------------------------------- +# End-to-end response — reasoningContent parsed to reasoning_content string +# --------------------------------------------------------------------------- - thinking_block = {"redactedContent": {}} - result = handler.translate_thinking_blocks(thinking_block) +class TestNova2EndToEndResponse: + """Verify transform_response produces reasoning_content from reasoningContent blocks.""" - assert result is not None - assert len(result) == 1 - assert result[0]["type"] == "redacted_thinking" + def _transform(self, content_blocks, model=NOVA_2_LITE): + config = AmazonConverseConfig() + body = { + "output": {"message": {"role": "assistant", "content": content_blocks}}, + "usage": {"inputTokens": 10, "outputTokens": 50, "totalTokens": 60}, + "stopReason": "end_turn", + "metrics": {"latencyMs": 100}, + } + resp = httpx.Response( + 200, json=body, request=httpx.Request("POST", "https://bedrock") + ) + return config.transform_response( + model=model, + raw_response=resp, + model_response=litellm.ModelResponse(), + logging_obj=None, + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + api_key=None, + json_mode=None, + ) + + def test_should_extract_reasoning_content_as_string(self): + result = self._transform( + [ + {"reasoningContent": {"reasoningText": {"text": "Step 1. "}}}, + {"reasoningContent": {"reasoningText": {"text": "Step 2."}}}, + {"text": "The answer is 4."}, + ] + ) + msg = result.choices[0].message + assert msg.content == "The answer is 4." + assert msg.reasoning_content == "Step 1. Step 2." + + def test_should_include_raw_blocks_in_provider_specific_fields(self): + result = self._transform( + [ + {"reasoningContent": {"reasoningText": {"text": "thinking..."}}}, + {"text": "done"}, + ] + ) + psf = result.choices[0].message.get("provider_specific_fields", {}) + assert "reasoningContentBlocks" in psf + + def test_should_omit_reasoning_content_when_absent(self): + result = self._transform([{"text": "Plain answer."}]) + assert not getattr(result.choices[0].message, "reasoning_content", None) + + +# --------------------------------------------------------------------------- +# Multi-turn — reasoning_content round-trips back to Bedrock format +# --------------------------------------------------------------------------- + + +class TestNova2MultiTurnMessageTranslation: + """Verify that assistant messages carrying reasoning from a previous turn are + correctly translated to Bedrock content blocks via _bedrock_converse_messages_pt.""" + + def _to_bedrock(self, messages, model=NOVA_2_LITE): + from litellm.litellm_core_utils.prompt_templates.factory import ( + _bedrock_converse_messages_pt, + ) + + return _bedrock_converse_messages_pt( + messages=messages, + model=model, + llm_provider="bedrock_converse", + ) + + def test_should_inline_unsigned_thinking_blocks_as_text(self): + """Without a signature, reasoning text becomes a plain text block.""" + bedrock_msgs = self._to_bedrock( + [ + {"role": "user", "content": "What is 2+2?"}, + { + "role": "assistant", + "content": "4.", + "thinking_blocks": [ + {"type": "thinking", "thinking": "Simple addition"}, + ], + }, + {"role": "user", "content": "Sure?"}, + ] + ) + assistant = next(m for m in bedrock_msgs if m["role"] == "assistant") + texts = [b["text"] for b in assistant["content"] if "text" in b] + assert "Simple addition" in texts + assert "4." in texts + + def test_should_keep_signed_thinking_blocks_as_reasoning_content(self): + """With a signature, reasoning is preserved as a reasoningContent block.""" + bedrock_msgs = self._to_bedrock( + [ + {"role": "user", "content": "What is 2+2?"}, + { + "role": "assistant", + "content": "4.", + "thinking_blocks": [ + {"type": "thinking", "thinking": "math", "signature": "sig-1"}, + ], + }, + {"role": "user", "content": "Sure?"}, + ] + ) + assistant = next(m for m in bedrock_msgs if m["role"] == "assistant") + rc_blocks = [b for b in assistant["content"] if "reasoningContent" in b] + assert len(rc_blocks) >= 1 + assert rc_blocks[0]["reasoningContent"]["reasoningText"]["text"] == "math" + assert rc_blocks[0]["reasoningContent"]["reasoningText"]["signature"] == "sig-1" + + def test_should_translate_inline_content_list_thinking_type(self): + """content=[{type:'thinking',...},{type:'text',...}] should also round-trip.""" + bedrock_msgs = self._to_bedrock( + [ + {"role": "user", "content": "Hi"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "hmm", "signature": "sig-2"}, + {"type": "text", "text": "Hello!"}, + ], + }, + {"role": "user", "content": "Bye"}, + ] + ) + assistant = next(m for m in bedrock_msgs if m["role"] == "assistant") + rc_blocks = [b for b in assistant["content"] if "reasoningContent" in b] + text_blocks = [ + b + for b in assistant["content"] + if "text" in b and "reasoningContent" not in b + ] + assert len(rc_blocks) >= 1 + assert any("Hello!" in b["text"] for b in text_blocks) diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py index dd0a3e36e46..2f9a0b63921 100644 --- a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py +++ b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py @@ -3,13 +3,9 @@ Integration tests for Vertex AI rerank functionality. These tests demonstrate end-to-end usage of the Vertex AI rerank feature. """ import importlib -import os -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import httpx -import pytest - -from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig class TestVertexAIRerankIntegration: @@ -20,16 +16,25 @@ class TestVertexAIRerankIntegration: importlib.reload(rerank_transformation_module) # Re-import after reload to get the fresh class - from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig as FreshConfig + from litellm.llms.vertex_ai.rerank.transformation import ( + VertexAIRerankConfig as FreshConfig, + ) self.config = FreshConfig() self.model = "semantic-ranker-default@latest" - @patch('litellm.llms.vertex_ai.rerank.transformation.VertexAIRerankConfig._ensure_access_token') - def test_end_to_end_rerank_flow(self, mock_ensure_access_token): - """Test complete rerank flow from request to response.""" - # Mock authentication - mock_ensure_access_token.return_value = ("test-access-token", "test-project-123") - + def test_end_to_end_rerank_flow(self): + """ + Test complete rerank flow from request to response. + + Uses instance-level mocking to avoid class-reference issues caused by + importlib.reload(litellm) in conftest.py. + """ + # Mock authentication at instance level + mock_ensure_access_token = MagicMock( + return_value=("test-access-token", "test-project-123") + ) + self.config._ensure_access_token = mock_ensure_access_token + # Test documents documents = [ "Gemini is a cutting edge large language model created by Google.", @@ -38,43 +43,40 @@ class TestVertexAIRerankIntegration: "Google's Gemini AI model represents a significant advancement in artificial intelligence technology." ] query = "What is Google Gemini?" - + # Step 1: Test request transformation - with patch.object(self.config, 'get_vertex_ai_credentials', return_value=None), \ - patch.object(self.config, 'get_vertex_ai_project', return_value="test-project-123"): - - # Validate environment - headers = self.config.validate_environment( - headers={}, - model=self.model, - api_key=None - ) - - # Transform request - request_data = self.config.transform_rerank_request( - model=self.model, - optional_rerank_params={ - "query": query, - "documents": documents, - "top_n": 2, - "return_documents": True - }, - headers=headers - ) - - # Verify request structure - assert request_data["model"] == self.model - assert request_data["query"] == query - assert request_data["topN"] == 2 - assert request_data["ignoreRecordDetailsInResponse"] == False - assert len(request_data["records"]) == 4 - - # Verify record structure - for i, record in enumerate(request_data["records"]): - assert record["id"] == str(i) # 0-based indexing - assert "title" in record - assert "content" in record - assert record["content"] == documents[i] + # Validate environment + headers = self.config.validate_environment( + headers={}, + model=self.model, + api_key=None + ) + + # Transform request + request_data = self.config.transform_rerank_request( + model=self.model, + optional_rerank_params={ + "query": query, + "documents": documents, + "top_n": 2, + "return_documents": True + }, + headers=headers + ) + + # Verify request structure + assert request_data["model"] == self.model + assert request_data["query"] == query + assert request_data["topN"] == 2 + assert request_data["ignoreRecordDetailsInResponse"] == False + assert len(request_data["records"]) == 4 + + # Verify record structure + for i, record in enumerate(request_data["records"]): + assert record["id"] == str(i) # 0-based indexing + assert "title" in record + assert "content" in record + assert record["content"] == documents[i] # Step 2: Test response transformation # Mock Vertex AI Discovery Engine response diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py index 2e631054143..d83f4612ee0 100644 --- a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py @@ -41,6 +41,18 @@ class TestVertexAIRerankTransform: for var, value in self._saved_env.items(): os.environ[var] = value +<<<<<<< litellm_oss_staging_02_17_2026 + def test_get_complete_url(self): + """ + Test URL generation for Vertex AI Discovery Engine rerank API. + + Uses instance-level mocking to avoid class-reference issues caused by + importlib.reload(litellm) in conftest.py. + """ + # Mock _ensure_access_token at instance level to return (token, project_id) + mock_ensure_access_token = MagicMock(return_value=("mock-token", None)) + self.config._ensure_access_token = mock_ensure_access_token +======= @patch('litellm.llms.vertex_ai.rerank.transformation.get_secret_str') @patch('litellm.llms.vertex_ai.rerank.transformation.VertexAIRerankConfig._ensure_access_token') def test_get_complete_url(self, mock_ensure_access_token, mock_get_secret_str): @@ -52,6 +64,7 @@ class TestVertexAIRerankTransform: def mock_get_secret(key): return os.environ.get(key) mock_get_secret_str.side_effect = mock_get_secret +>>>>>>> main # Test with project ID from environment with patch.dict(os.environ, {"VERTEXAI_PROJECT": "test-project-123"}): @@ -96,6 +109,33 @@ class TestVertexAIRerankTransform: finally: litellm.vertex_project = original_project +<<<<<<< litellm_oss_staging_02_17_2026 + def test_validate_environment(self): + """ + Test environment validation and header setup. + + Uses instance-level mocking to avoid class-reference issues caused by + importlib.reload(litellm) in conftest.py. + """ + # Mock the authentication at instance level + mock_ensure_access_token = MagicMock( + return_value=("test-access-token", "test-project-123") + ) + self.config._ensure_access_token = mock_ensure_access_token + + headers = self.config.validate_environment( + headers={}, + model=self.model, + api_key=None + ) + + expected_headers = { + "Authorization": "Bearer test-access-token", + "Content-Type": "application/json", + "X-Goog-User-Project": "test-project-123" + } + assert headers == expected_headers +======= @patch('litellm.llms.vertex_ai.rerank.transformation.get_secret_str') @patch('litellm.llms.vertex_ai.rerank.transformation.VertexAIRerankConfig._ensure_access_token') def test_validate_environment(self, mock_ensure_access_token, mock_get_secret_str): @@ -124,6 +164,7 @@ class TestVertexAIRerankTransform: "X-Goog-User-Project": "test-project-123" } assert headers == expected_headers +>>>>>>> main def test_transform_rerank_request_basic(self): """Test basic request transformation for Vertex AI Discovery Engine format.""" @@ -457,33 +498,40 @@ class TestVertexAIRerankTransform: assert params["top_n"] == 2 assert params["return_documents"] == True - @patch('litellm.llms.vertex_ai.rerank.transformation.VertexAIRerankConfig._ensure_access_token') - def test_validate_environment_with_optional_params(self, mock_ensure_access_token): - """Test that validate_environment accepts and uses optional_params for credentials.""" - # Mock the authentication - mock_ensure_access_token.return_value = ("test-access-token", "test-project-123") - + def test_validate_environment_with_optional_params(self): + """ + Test that validate_environment accepts and uses optional_params for credentials. + + Uses instance-level mocking to avoid class-reference issues caused by + importlib.reload(litellm) in conftest.py. + """ + # Mock the authentication at instance level + mock_ensure_access_token = MagicMock( + return_value=("test-access-token", "test-project-123") + ) + self.config._ensure_access_token = mock_ensure_access_token + optional_params = { "vertex_credentials": "path/to/credentials.json", "vertex_project": "custom-project-id", "query": "test query", "documents": ["doc1"] } - + headers = self.config.validate_environment( headers={}, model=self.model, api_key=None, optional_params=optional_params ) - + # Verify that _ensure_access_token was called with the credentials from optional_params mock_ensure_access_token.assert_called_once() call_args = mock_ensure_access_token.call_args # The first call argument should be credentials (which will be the value from optional_params) # We can't check the exact value easily due to how get_vertex_ai_credentials pops values, # but we can verify the headers were set correctly - + expected_headers = { "Authorization": "Bearer test-access-token", "Content-Type": "application/json", diff --git a/tests/test_litellm/llms/watsonx/rerank/__init__.py b/tests/test_litellm/llms/watsonx/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py b/tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py new file mode 100644 index 00000000000..4a2edd9810b --- /dev/null +++ b/tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py @@ -0,0 +1,224 @@ +""" +Tests for IBM watsonx.ai rerank transformation functionality. +""" +import json +import re +import uuid +from unittest.mock import MagicMock + +import httpx +import pytest + +from litellm.llms.watsonx.common_utils import ( + WatsonXAIError, +) +from litellm.llms.watsonx.rerank.transformation import IBMWatsonXRerankConfig +from litellm.types.rerank import RerankResponse + + +class TestIBMWatsonXRerankTransform: + def setup_method(self): + self.config = IBMWatsonXRerankConfig() + self.model = "watsonx/cross-encoder/ms-marco-minilm-l-12-v2" + + def test_get_complete_url(self): + """Test URL generation for IBM watsonx.ai rerank API.""" + + api_base = "https://us-south.ml.cloud.ibm.com" + model = "watsonx/cross-encoder/ms-marco-minilm-l-12-v2" + url = self.config.get_complete_url(api_base, model) + assert url == "https://us-south.ml.cloud.ibm.com/ml/v1/text/rerank?version=2024-03-13" + + def test_map_cohere_rerank_params_basic(self): + """Test basic parameter mapping for IBM watsonx.ai rerank.""" + params = self.config.map_cohere_rerank_params( + non_default_params={ + "query": "hello", + "documents": ["hello", "world"], + "top_n": 2, + "return_documents": True, + "max_tokens_per_doc": 100, + }, + model="test", + drop_params=False, + query="hello", + documents=["hello", "world"], + ) + assert params["query"] == "hello" + assert params["inputs"] == [{"text": "hello"}, {"text": "world"}] + assert params["parameters"]["return_options"]["top_n"] == 2 + assert params["parameters"]["return_options"]["inputs"] is True + assert params["parameters"]["truncate_input_tokens"] == 100 + + def test_transform_rerank_request(self): + """Test request transformation for IBM watsonx.ai format.""" + optional_params = { + "query": "What is the capital of France?", + "documents": [ + "Paris is the capital of France.", + "France is a country in Europe.", + ], + "top_n": 2, + "return_documents": True, + "project_id": uuid.uuid4(), + } + + request_body = self.config.transform_rerank_request( + model="cross-encoder/ms-marco-minilm-l-12-v2", optional_rerank_params=optional_params, headers={} + ) + + assert request_body["model_id"] == "cross-encoder/ms-marco-minilm-l-12-v2" + assert request_body["project_id"] is not None + assert request_body["query"] == "What is the capital of France?" + assert request_body["documents"] == optional_params["documents"] + assert request_body["top_n"] == 2 + assert request_body["return_documents"] is True + + def test_transform_rerank_response_success(self): + """Test successful response transformation.""" + # Mock IBM watsonx.ai response format + response_data = { + "model_id": self.model, + "results": [ + { + "index": 0, + "score": 6.53515625, + "input": {"text": "Python is great for beginners due to simple syntax."}, + }, + {"index": 1, "score": -7.1875, "input": {"text": "JavaScript runs in browsers and is versatile."}}, + ], + "input_token_count": 62, + } + + # Create mock httpx response + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = response_data + mock_response.status_code = 200 + mock_response.headers = {} + + # Create mock logging object + mock_logging = MagicMock() + + model_response = RerankResponse() + + result = self.config.transform_rerank_response( + model=self.model, + raw_response=mock_response, + model_response=model_response, + logging_obj=mock_logging, + ) + + # Verify response structure + # IBM watsonx.ai doesn't return "id", so it uses "model" as the id + assert result.id == "watsonx/cross-encoder/ms-marco-minilm-l-12-v2" + assert len(result.results) == 2 + assert result.results[0]["index"] == 0 + assert result.results[0]["relevance_score"] == 6.53515625 + assert result.results[0]["document"]["text"] == "Python is great for beginners due to simple syntax." + assert result.results[1]["index"] == 1 + assert result.results[1]["relevance_score"] == -7.1875 + assert result.results[1]["document"]["text"] == "JavaScript runs in browsers and is versatile." + + # # Verify metadata + assert result.meta["tokens"]["input_tokens"] == 62 + + def test_transform_rerank_response_without_documents(self): + """Test response transformation when return_documents is False.""" + response_data = { + "model_id": self.model, + "results": [ + { + "index": 0, + "score": 6.53515625, + }, + { + "index": 1, + "score": -7.1875, + }, + ], + "input_token_count": 62, + } + + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = response_data + mock_response.status_code = 200 + mock_response.headers = {} + + mock_logging = MagicMock() + model_response = RerankResponse() + + result = self.config.transform_rerank_response( + model=self.model, + raw_response=mock_response, + model_response=model_response, + logging_obj=mock_logging, + ) + + # Verify response structure + # IBM watsonx.ai doesn't return "id", so it uses "model" as the id + assert result.id == "watsonx/cross-encoder/ms-marco-minilm-l-12-v2" + assert len(result.results) == 2 + + assert result.results[0]["index"] == 0 + assert result.results[0]["relevance_score"] == 6.53515625 + assert "document" not in result.results[0] + + assert result.results[1]["index"] == 1 + assert result.results[1]["relevance_score"] == -7.1875 + assert "document" not in result.results[1] + + def test_transform_rerank_response_missing_results(self): + """Test that missing results raises ValueError.""" + response_data = { + "model": self.model, + "usage": {"total_tokens": 10}, + } + + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = response_data + mock_response.status_code = 200 + mock_response.headers = {} + + mock_logging = MagicMock() + model_response = RerankResponse() + + expected_error_msg = re.escape("No results found") + + with pytest.raises(ValueError, match=expected_error_msg): + self.config.transform_rerank_response( + model=self.model, + raw_response=mock_response, + model_response=model_response, + logging_obj=mock_logging, + ) + + def test_transform_rerank_response_invalid_json(self): + """Test error handling for invalid JSON response.""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "doc", 0) + mock_response.text = "Invalid JSON response" + mock_response.status_code = 500 + mock_response.headers = {} + + mock_logging = MagicMock() + model_response = RerankResponse() + + expected_error_msg = re.escape("Failed to parse response") + + with pytest.raises(Exception, match=expected_error_msg): + self.config.transform_rerank_response( + model=self.model, + raw_response=mock_response, + model_response=model_response, + logging_obj=mock_logging, + ) + + def test_get_supported_cohere_rerank_params(self): + """Test getting supported parameters for IBM watsonx.ai rerank.""" + supported_params = self.config.get_supported_cohere_rerank_params(self.model) + assert "query" in supported_params + assert "documents" in supported_params + assert "top_n" in supported_params + assert "return_documents" in supported_params + assert "max_tokens_per_doc" in supported_params + assert len(supported_params) == 5 diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index 789af480e72..193b014f03d 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -62,3 +62,27 @@ def test_get_complete_model_list_order(key_models, team_models, proxy_model_list infer_model_from_keys=False, llm_router=Router(model_list=model_list), ) == expected + + +def test_get_complete_model_list_byok_wildcard_expansion(): + """ + Test that wildcard models (e.g., openai/*) are expanded when the router has + no deployment for them - BYOK case where team has openai/* but proxy has + no openai config. + """ + from litellm.proxy.auth.model_checks import get_complete_model_list + from litellm import Router + + # Router with empty model_list - no openai/* deployment (BYOK scenario) + result = get_complete_model_list( + key_models=[], + team_models=["openai/*"], + proxy_model_list=[], + user_model=None, + infer_model_from_keys=False, + llm_router=Router(model_list=[]), + ) + # Should expand openai/* to actual OpenAI models + assert len(result) > 0 + assert all(m.startswith("openai/") for m in result) + assert "openai/*" not in result diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 74f5cf9bdd7..c2c20485b5e 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -1600,6 +1600,56 @@ def test_completion_cost_service_tier_priority(): ), "Costs from params and usage should be similar (both flex)" +def test_completion_cost_service_tier_for_bedrock(): + """Test that Bedrock cost calculation applies service_tier-specific pricing.""" + from litellm import completion_cost + + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model = "bedrock/us-east-1/test-bedrock-service-tier-cost-model" + litellm.register_model( + model_cost={ + model: { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "input_cost_per_token_priority": 0.01, + "output_cost_per_token_priority": 0.02, + "input_cost_per_token_flex": 0.0005, + "output_cost_per_token_flex": 0.001, + "litellm_provider": "bedrock", + "max_tokens": 8192, + } + } + ) + + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + response = ModelResponse(usage=usage, model=model) + + default_cost = completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock", + ) + + priority_cost = completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock", + optional_params={"service_tier": "priority"}, + ) + + response_with_flex_tier = ModelResponse(usage=usage, model=model) + setattr(response_with_flex_tier, "service_tier", "flex") + flex_cost = completion_cost( + completion_response=response_with_flex_tier, + model=model, + custom_llm_provider="bedrock", + ) + + assert priority_cost > default_cost > flex_cost > 0 + + def test_gemini_cache_tokens_details_no_negative_values(): """ Test for Issue #18750: Negative text_tokens with Gemini caching