mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
Merge pull request #21361 from BerriAI/litellm_oss_staging_02_17_2026
Litellm oss staging 02 17 2026
This commit is contained in:
Sameer Kankute
GitHub
commit
aa255c7e63
42 changed files with 1774 additions and 553 deletions
52
docs/my-website/docs/providers/watsonx/rerank.md
Normal file
52
docs/my-website/docs/providers/watsonx/rerank.md
Normal file
|
|
@ -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
|
||||
```
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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) |
|
||||
| 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) |
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
|
|
@ -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 = [
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 <token>" 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]]:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
)
|
||||
model=model,
|
||||
usage=usage,
|
||||
custom_llm_provider="bedrock",
|
||||
service_tier=service_tier,
|
||||
)
|
||||
|
|
|
|||
0
litellm/llms/watsonx/__init__.py
Normal file
0
litellm/llms/watsonx/__init__.py
Normal file
0
litellm/llms/watsonx/chat/__init__.py
Normal file
0
litellm/llms/watsonx/chat/__init__.py
Normal file
0
litellm/llms/watsonx/completion/__init__.py
Normal file
0
litellm/llms/watsonx/completion/__init__.py
Normal file
0
litellm/llms/watsonx/embed/__init__.py
Normal file
0
litellm/llms/watsonx/embed/__init__.py
Normal file
0
litellm/llms/watsonx/rerank/__init__.py
Normal file
0
litellm/llms/watsonx/rerank/__init__.py
Normal file
204
litellm/llms/watsonx/rerank/transformation.py
Normal file
204
litellm/llms/watsonx/rerank/transformation.py
Normal file
|
|
@ -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,
|
||||
)
|
||||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
92
tests/litellm/llms/bedrock/test_nova_imported_models.py
Normal file
92
tests/litellm/llms/bedrock/test_nova_imported_models.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
)
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
0
tests/test_litellm/llms/watsonx/rerank/__init__.py
Normal file
0
tests/test_litellm/llms/watsonx/rerank/__init__.py
Normal file
224
tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py
Normal file
224
tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue