tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-06 18:23:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #14700 from BerriAI/litellm_contributor_prs_09_18_2025_p2

Browse Source 
Update Bedrock documentation for Titan V2 encoding_format support + Anthropic - account for 1h vs. 5m cache creation token cost difference + UI - add langsmith_sampling_rate as a dynamic param
This commit is contained in:
Krish Dholakia
2025-09-18 23:38:29 -07:00
committed by GitHub
parent b6d5a65fed 92e841e311
commit d5a839d971
18 changed files with 799 additions and 292 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/providers/bedrock.md
+53
View File
@@ -1822,6 +1822,59 @@ Here's an example of using a bedrock model with LiteLLM. For a complete list, re
| Mixtral 8x7B Instruct | `completion(model='bedrock/mistral.mixtral-8x7b-instruct-v0:1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
## Bedrock Embedding
### API keys
This can be set as env variables or passed as **params to litellm.embedding()**
```python
import os
os.environ["AWS_ACCESS_KEY_ID"] = "" # Access key
os.environ["AWS_SECRET_ACCESS_KEY"] = "" # Secret access key
os.environ["AWS_REGION_NAME"] = "" # us-east-1, us-east-2, us-west-1, us-west-2
```
### Usage
```python
from litellm import embedding
response = embedding(
model="bedrock/amazon.titan-embed-text-v1",
input=["good morning from litellm"],
)
print(response)
```
#### Titan V2 - encoding_format support
```python
from litellm import embedding
# Float format (default)
response = embedding(
model="bedrock/amazon.titan-embed-text-v2:0",
input=["good morning from litellm"],
encoding_format="float" # Returns float array
)
# Binary format
response = embedding(
model="bedrock/amazon.titan-embed-text-v2:0",
input=["good morning from litellm"],
encoding_format="base64" # Returns base64 encoded binary
)
```
## Supported AWS Bedrock Embedding Models
| Model Name | Usage | Supported Additional OpenAI params |
|----------------------|---------------------------------------------|-----|
| Titan Embeddings V2 | `embedding(model="bedrock/amazon.titan-embed-text-v2:0", input=input)` | `dimensions`, `encoding_format` |
| Titan Embeddings - V1 | `embedding(model="bedrock/amazon.titan-embed-text-v1", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py#L53)
| Titan Multimodal Embeddings | `embedding(model="bedrock/amazon.titan-embed-image-v1", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_multimodal_transformation.py#L28) |
| Cohere Embeddings - English | `embedding(model="bedrock/cohere.embed-english-v3", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/cohere_transformation.py#L18)
| Cohere Embeddings - Multilingual | `embedding(model="bedrock/cohere.embed-multilingual-v3", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/cohere_transformation.py#L18)
### Advanced - [Drop Unsupported Params](https://docs.litellm.ai/docs/completion/drop_params#openai-proxy-usage)
### Advanced - [Pass model/provider-specific Params](https://docs.litellm.ai/docs/completion/provider_specific_params#proxy-usage)
## Image Generation
Use this for stable diffusion, and amazon nova canvas on bedrock
litellm/integrations/langsmith.py
+34 -24
View File
@@ -39,6 +39,7 @@ class LangsmithLogger(CustomBatchLogger):
langsmith_api_key: Optional[str] = None,
langsmith_project: Optional[str] = None,
langsmith_base_url: Optional[str] = None,
langsmith_sampling_rate: Optional[float] = None,
**kwargs,
):
self.flush_lock = asyncio.Lock()
@@ -49,7 +50,8 @@ class LangsmithLogger(CustomBatchLogger):
langsmith_base_url=langsmith_base_url,
)
self.sampling_rate: float = (
float(os.getenv("LANGSMITH_SAMPLING_RATE")) # type: ignore
langsmith_sampling_rate
or float(os.getenv("LANGSMITH_SAMPLING_RATE")) # type: ignore
if os.getenv("LANGSMITH_SAMPLING_RATE") is not None
and os.getenv("LANGSMITH_SAMPLING_RATE").strip().isdigit() # type: ignore
else 1.0
@@ -76,26 +78,14 @@ class LangsmithLogger(CustomBatchLogger):
langsmith_base_url: Optional[str] = None,
) -> LangsmithCredentialsObject:
_credentials_api_key = langsmith_api_key or os.getenv("LANGSMITH_API_KEY")
if _credentials_api_key is None:
raise Exception(
"Invalid Langsmith API Key given. _credentials_api_key=None."
)
_credentials_project = (
langsmith_project or os.getenv("LANGSMITH_PROJECT") or "litellm-completion"
)
if _credentials_project is None:
raise Exception(
"Invalid Langsmith API Key given. _credentials_project=None."
)
_credentials_base_url = (
langsmith_base_url
or os.getenv("LANGSMITH_BASE_URL")
or "https://api.smith.langchain.com"
)
if _credentials_base_url is None:
raise Exception(
"Invalid Langsmith API Key given. _credentials_base_url=None."
)
return LangsmithCredentialsObject(
LANGSMITH_API_KEY=_credentials_api_key,
@@ -200,12 +190,7 @@ class LangsmithLogger(CustomBatchLogger):
def log_success_event(self, kwargs, response_obj, start_time, end_time):
try:
sampling_rate = (
float(os.getenv("LANGSMITH_SAMPLING_RATE")) # type: ignore
if os.getenv("LANGSMITH_SAMPLING_RATE") is not None
and os.getenv("LANGSMITH_SAMPLING_RATE").strip().isdigit() # type: ignore
else 1.0
)
sampling_rate = self._get_sampling_rate_to_use_for_request(kwargs=kwargs)
random_sample = random.random()
if random_sample > sampling_rate:
verbose_logger.info(
@@ -219,6 +204,7 @@ class LangsmithLogger(CustomBatchLogger):
kwargs,
response_obj,
)
credentials = self._get_credentials_to_use_for_request(kwargs=kwargs)
data = self._prepare_log_data(
kwargs=kwargs,
@@ -245,7 +231,7 @@ class LangsmithLogger(CustomBatchLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
try:
sampling_rate = self.sampling_rate
sampling_rate = self._get_sampling_rate_to_use_for_request(kwargs=kwargs)
random_sample = random.random()
if random_sample > sampling_rate:
verbose_logger.info(
@@ -286,7 +272,7 @@ class LangsmithLogger(CustomBatchLogger):
)
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
sampling_rate = self.sampling_rate
sampling_rate = self._get_sampling_rate_to_use_for_request(kwargs=kwargs)
random_sample = random.random()
if random_sample > sampling_rate:
verbose_logger.info(
@@ -417,6 +403,17 @@ class LangsmithLogger(CustomBatchLogger):
for queue_object in self.log_queue:
credentials = queue_object["credentials"]
# if credential missing, skip - log warning
if (
credentials["LANGSMITH_API_KEY"] is None
or credentials["LANGSMITH_PROJECT"] is None
):
verbose_logger.warning(
"Langsmith Logging - credentials missing - api_key: %s, project: %s",
credentials["LANGSMITH_API_KEY"],
credentials["LANGSMITH_PROJECT"],
)
continue
key = CredentialsKey(
api_key=credentials["LANGSMITH_API_KEY"],
project=credentials["LANGSMITH_PROJECT"],
@@ -432,6 +429,19 @@ class LangsmithLogger(CustomBatchLogger):
return log_queue_by_credentials
def _get_sampling_rate_to_use_for_request(self, kwargs: Dict[str, Any]) -> float:
standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = (
kwargs.get("standard_callback_dynamic_params", None)
)
sampling_rate: float = self.sampling_rate
if standard_callback_dynamic_params is not None:
_sampling_rate = standard_callback_dynamic_params.get(
"langsmith_sampling_rate"
)
if _sampling_rate is not None:
sampling_rate = float(_sampling_rate)
return sampling_rate
def _get_credentials_to_use_for_request(
self, kwargs: Dict[str, Any]
) -> LangsmithCredentialsObject:
@@ -442,9 +452,9 @@ class LangsmithLogger(CustomBatchLogger):
Otherwise, use the default credentials.
"""
standard_callback_dynamic_params: Optional[
StandardCallbackDynamicParams
] = kwargs.get("standard_callback_dynamic_params", None)
standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = (
kwargs.get("standard_callback_dynamic_params", None)
)
if standard_callback_dynamic_params is not None:
credentials = self.get_credentials_from_env(
langsmith_api_key=standard_callback_dynamic_params.get(
litellm/litellm_core_utils/llm_cost_calc/utils.py
+238 -106
View File
@@ -1,11 +1,12 @@
# What is this?
## Helper utilities for cost_per_token()
from typing import Any, Literal, Optional, Tuple, cast
from typing import Any, Literal, Optional, Tuple, TypedDict, cast
import litellm
from litellm._logging import verbose_logger
from litellm.types.utils import (
CacheCreationTokenDetails,
CallTypes,
ImageResponse,
ModelInfo,
@@ -115,7 +116,7 @@ def _generic_cost_per_character(
def _get_token_base_cost(
model_info: ModelInfo, usage: Usage
) -> Tuple[float, float, float, float]:
) -> Tuple[float, float, float, float, float]:
"""
Return prompt cost, completion cost, and cache costs for a given model and usage.
@@ -134,6 +135,10 @@ def _get_token_base_cost(
cache_creation_cost = cast(
float, _get_cost_per_unit(model_info, "cache_creation_input_token_cost")
)
cache_creation_cost_above_1hr = cast(
float,
_get_cost_per_unit(model_info, "cache_creation_input_token_cost_above_1hr"),
)
cache_read_cost = cast(
float, _get_cost_per_unit(model_info, "cache_read_input_token_cost")
)
@@ -194,7 +199,13 @@ def _get_token_base_cost(
except Exception:
continue
return prompt_base_cost, completion_base_cost, cache_creation_cost, cache_read_cost
return (
prompt_base_cost,
completion_base_cost,
cache_creation_cost,
cache_creation_cost_above_1hr,
cache_read_cost,
)
def calculate_cost_component(
@@ -241,6 +252,196 @@ def _get_cost_per_unit(
return default_value
def calculate_cache_writing_cost(
cache_creation_tokens: int,
cache_creation_token_details: Optional[CacheCreationTokenDetails],
cache_creation_cost_above_1hr: float,
cache_creation_cost: float,
) -> float:
"""
Adjust cost of cache creation tokens based on the cache creation token details.
"""
total_cost: float = 0.0
if cache_creation_token_details is not None:
# get the number of 5m and 1h cache creation tokens
cache_creation_tokens_5m = (
cache_creation_token_details.ephemeral_5m_input_tokens
)
cache_creation_tokens_1h = (
cache_creation_token_details.ephemeral_1h_input_tokens
)
# add the number of 5m and 1h cache creation tokens to the cache creation tokens
total_cost += (
cache_creation_tokens_5m * cache_creation_cost
if cache_creation_tokens_5m is not None
else 0.0
)
total_cost += (
cache_creation_tokens_1h * cache_creation_cost_above_1hr
if cache_creation_tokens_1h is not None
else 0.0
)
else:
total_cost += cache_creation_tokens * cache_creation_cost
return total_cost
class PromptTokensDetailsResult(TypedDict):
cache_hit_tokens: int
cache_creation_tokens: int
cache_creation_token_details: Optional[CacheCreationTokenDetails]
text_tokens: int
audio_tokens: int
character_count: int
image_count: int
video_length_seconds: int
def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
cache_hit_tokens = (
cast(Optional[int], getattr(usage.prompt_tokens_details, "cached_tokens", 0))
or 0
)
cache_creation_tokens = (
cast(
Optional[int],
getattr(usage.prompt_tokens_details, "cache_creation_tokens", 0),
)
or 0
)
cache_creation_token_details = (
cast(
Optional[CacheCreationTokenDetails],
getattr(usage.prompt_tokens_details, "cache_creation_token_details", None),
)
or None
)
text_tokens = (
cast(Optional[int], getattr(usage.prompt_tokens_details, "text_tokens", None))
or 0 # default to prompt tokens, if this field is not set
)
audio_tokens = (
cast(Optional[int], getattr(usage.prompt_tokens_details, "audio_tokens", 0))
or 0
)
character_count = (
cast(
Optional[int],
getattr(usage.prompt_tokens_details, "character_count", 0),
)
or 0
)
image_count = (
cast(Optional[int], getattr(usage.prompt_tokens_details, "image_count", 0)) or 0
)
video_length_seconds = (
cast(
Optional[int],
getattr(usage.prompt_tokens_details, "video_length_seconds", 0),
)
or 0
)
return PromptTokensDetailsResult(
cache_hit_tokens=cache_hit_tokens,
cache_creation_tokens=cache_creation_tokens,
cache_creation_token_details=cache_creation_token_details,
text_tokens=text_tokens,
audio_tokens=audio_tokens,
character_count=character_count,
image_count=image_count,
video_length_seconds=video_length_seconds,
)
class CompletionTokensDetailsResult(TypedDict):
audio_tokens: int
text_tokens: int
reasoning_tokens: int
def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResult:
audio_tokens = (
cast(
Optional[int],
getattr(usage.completion_tokens_details, "audio_tokens", 0),
)
or 0
)
text_tokens = (
cast(
Optional[int],
getattr(usage.completion_tokens_details, "text_tokens", None),
)
or 0 # default to completion tokens, if this field is not set
)
reasoning_tokens = (
cast(
Optional[int],
getattr(usage.completion_tokens_details, "reasoning_tokens", 0),
)
or 0
)
return CompletionTokensDetailsResult(
audio_tokens=audio_tokens,
text_tokens=text_tokens,
reasoning_tokens=reasoning_tokens,
)
def _calculate_input_cost(
prompt_tokens_details: PromptTokensDetailsResult,
model_info: ModelInfo,
prompt_base_cost: float,
cache_read_cost: float,
cache_creation_cost: float,
cache_creation_cost_above_1hr: float,
) -> float:
"""
Calculates the input cost for a given model, prompt tokens, and completion tokens.
"""
prompt_cost = float(prompt_tokens_details["text_tokens"]) * prompt_base_cost
### CACHE READ COST - Now uses tiered pricing
prompt_cost += float(prompt_tokens_details["cache_hit_tokens"]) * cache_read_cost
### AUDIO COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_audio_token", prompt_tokens_details["audio_tokens"]
)
### CACHE WRITING COST - Now uses tiered pricing
prompt_cost += calculate_cache_writing_cost(
cache_creation_tokens=prompt_tokens_details["cache_creation_tokens"],
cache_creation_token_details=prompt_tokens_details[
"cache_creation_token_details"
],
cache_creation_cost_above_1hr=cache_creation_cost_above_1hr,
cache_creation_cost=cache_creation_cost,
)
### CHARACTER COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_character", prompt_tokens_details["character_count"]
)
### IMAGE COUNT COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_image", prompt_tokens_details["image_count"]
)
### VIDEO LENGTH COST
prompt_cost += calculate_cost_component(
model_info,
"input_cost_per_video_per_second",
prompt_tokens_details["video_length_seconds"],
)
return prompt_cost
def generic_cost_per_token(
model: str, usage: Usage, custom_llm_provider: str
) -> Tuple[float, float]:
@@ -264,97 +465,45 @@ def generic_cost_per_token(
### Cost of processing (non-cache hit + cache hit) + Cost of cache-writing (cache writing)
prompt_cost = 0.0
### PROCESSING COST
text_tokens = usage.prompt_tokens
cache_hit_tokens = 0
cache_creation_tokens = 0
audio_tokens = 0
character_count = 0
image_count = 0
video_length_seconds = 0
prompt_tokens_details = PromptTokensDetailsResult(
cache_hit_tokens=0,
cache_creation_tokens=0,
cache_creation_token_details=None,
text_tokens=usage.prompt_tokens,
audio_tokens=0,
character_count=0,
image_count=0,
video_length_seconds=0,
)
if usage.prompt_tokens_details:
cache_hit_tokens = (
cast(
Optional[int], getattr(usage.prompt_tokens_details, "cached_tokens", 0)
)
or 0
)
cache_creation_tokens = (
cast(
Optional[int],
getattr(usage.prompt_tokens_details, "cache_creation_tokens", 0),
)
or 0
)
text_tokens = (
cast(
Optional[int], getattr(usage.prompt_tokens_details, "text_tokens", None)
)
or 0 # default to prompt tokens, if this field is not set
)
audio_tokens = (
cast(Optional[int], getattr(usage.prompt_tokens_details, "audio_tokens", 0))
or 0
)
character_count = (
cast(
Optional[int],
getattr(usage.prompt_tokens_details, "character_count", 0),
)
or 0
)
image_count = (
cast(Optional[int], getattr(usage.prompt_tokens_details, "image_count", 0))
or 0
)
video_length_seconds = (
cast(
Optional[int],
getattr(usage.prompt_tokens_details, "video_length_seconds", 0),
)
or 0
)
prompt_tokens_details = _parse_prompt_tokens_details(usage)
## EDGE CASE - text tokens not set inside PromptTokensDetails
if text_tokens == 0:
if prompt_tokens_details["text_tokens"] == 0:
text_tokens = (
usage.prompt_tokens
- cache_hit_tokens
- audio_tokens
- cache_creation_tokens
- prompt_tokens_details["cache_hit_tokens"]
- prompt_tokens_details["audio_tokens"]
- prompt_tokens_details["cache_creation_tokens"]
)
prompt_tokens_details["text_tokens"] = text_tokens
prompt_base_cost, completion_base_cost, cache_creation_cost, cache_read_cost = (
_get_token_base_cost(model_info=model_info, usage=usage)
)
(
prompt_base_cost,
completion_base_cost,
cache_creation_cost,
cache_creation_cost_above_1hr,
cache_read_cost,
) = _get_token_base_cost(model_info=model_info, usage=usage)
prompt_cost = float(text_tokens) * prompt_base_cost
### CACHE READ COST - Now uses tiered pricing
prompt_cost += float(cache_hit_tokens) * cache_read_cost
### AUDIO COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_audio_token", audio_tokens
)
### CACHE WRITING COST - Now uses tiered pricing
prompt_cost += float(cache_creation_tokens) * cache_creation_cost
### CHARACTER COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_character", character_count
)
### IMAGE COUNT COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_image", image_count
)
### VIDEO LENGTH COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_video_per_second", video_length_seconds
prompt_cost = _calculate_input_cost(
prompt_tokens_details=prompt_tokens_details,
model_info=model_info,
prompt_base_cost=prompt_base_cost,
cache_read_cost=cache_read_cost,
cache_creation_cost=cache_creation_cost,
cache_creation_cost_above_1hr=cache_creation_cost_above_1hr,
)
## CALCULATE OUTPUT COST
@@ -363,27 +512,10 @@ def generic_cost_per_token(
reasoning_tokens = 0
is_text_tokens_total = False
if usage.completion_tokens_details is not None:
audio_tokens = (
cast(
Optional[int],
getattr(usage.completion_tokens_details, "audio_tokens", 0),
)
or 0
)
text_tokens = (
cast(
Optional[int],
getattr(usage.completion_tokens_details, "text_tokens", None),
)
or 0 # default to completion tokens, if this field is not set
)
reasoning_tokens = (
cast(
Optional[int],
getattr(usage.completion_tokens_details, "reasoning_tokens", 0),
)
or 0
)
completion_tokens_details = _parse_completion_tokens_details(usage)
audio_tokens = completion_tokens_details["audio_tokens"]
text_tokens = completion_tokens_details["text_tokens"]
reasoning_tokens = completion_tokens_details["reasoning_tokens"]
if text_tokens == 0:
text_tokens = usage.completion_tokens
litellm/llms/anthropic/chat/transformation.py
+17 -1
View File
@@ -45,7 +45,10 @@ from litellm.types.llms.openai import (
OpenAIMcpServerTool,
OpenAIWebSearchOptions,
)
from litellm.types.utils import CompletionTokensDetailsWrapper
from litellm.types.utils import (
CacheCreationTokenDetails,
CompletionTokensDetailsWrapper,
)
from litellm.types.utils import Message as LitellmMessage
from litellm.types.utils import PromptTokensDetailsWrapper, ServerToolUse
from litellm.utils import (
@@ -820,6 +823,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
_usage = usage_object
cache_creation_input_tokens: int = 0
cache_read_input_tokens: int = 0
cache_creation_token_details: Optional[CacheCreationTokenDetails] = None
web_search_requests: Optional[int] = None
if (
"cache_creation_input_tokens" in _usage
@@ -842,8 +846,20 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
int, _usage["server_tool_use"]["web_search_requests"]
)
if "cache_creation" in _usage and _usage["cache_creation"] is not None:
cache_creation_token_details = CacheCreationTokenDetails(
ephemeral_5m_input_tokens=_usage["cache_creation"].get(
"ephemeral_5m_input_tokens"
),
ephemeral_1h_input_tokens=_usage["cache_creation"].get(
"ephemeral_1h_input_tokens"
),
)
prompt_tokens_details = PromptTokensDetailsWrapper(
cached_tokens=cache_read_input_tokens,
cache_creation_tokens=cache_read_input_tokens,
cache_creation_token_details=cache_creation_token_details,
)
completion_token_details = (
CompletionTokensDetailsWrapper(
litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py
+37 -15
View File
@@ -10,7 +10,7 @@ Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-tit
"""
import types
from typing import List, Optional
from typing import List, Optional, Union
from litellm.types.llms.bedrock import (
AmazonTitanV2EmbeddingRequest,
@@ -30,9 +30,7 @@ class AmazonTitanV2Config:
normalize: Optional[bool] = None
dimensions: Optional[int] = None
def __init__(
self, normalize: Optional[bool] = None, dimensions: Optional[int] = None
) -> None:
def __init__(self, normalize: Optional[bool] = None, dimensions: Optional[int] = None) -> None:
locals_ = locals().copy()
for key, value in locals_.items():
if key != "self" and value is not None:
@@ -57,32 +55,56 @@ class AmazonTitanV2Config:
}
def get_supported_openai_params(self) -> List[str]:
return ["dimensions"]
return ["dimensions", "encoding_format"]
def map_openai_params(
self, non_default_params: dict, optional_params: dict
) -> dict:
def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict:
for k, v in non_default_params.items():
if k == "dimensions":
optional_params["dimensions"] = v
elif k == "encoding_format":
# Map OpenAI encoding_format to AWS embeddingTypes
if v == "float":
optional_params["embeddingTypes"] = ["float"]
elif v == "base64":
# base64 maps to binary format in AWS
optional_params["embeddingTypes"] = ["binary"]
else:
# For any other encoding format, default to float
optional_params["embeddingTypes"] = ["float"]
return optional_params
def _transform_request(
self, input: str, inference_params: dict
) -> AmazonTitanV2EmbeddingRequest:
def _transform_request(self, input: str, inference_params: dict) -> AmazonTitanV2EmbeddingRequest:
return AmazonTitanV2EmbeddingRequest(inputText=input, **inference_params) # type: ignore
def _transform_response(
self, response_list: List[dict], model: str
) -> EmbeddingResponse:
def _transform_response(self, response_list: List[dict], model: str) -> EmbeddingResponse:
total_prompt_tokens = 0
transformed_responses: List[Embedding] = []
for index, response in enumerate(response_list):
_parsed_response = AmazonTitanV2EmbeddingResponse(**response) # type: ignore
# According to AWS docs, embeddingsByType is always present
# If binary was requested (encoding_format="base64"), use binary data
# Otherwise, use float data from embeddingsByType or fallback to embedding field
embedding_data: Union[List[float], List[int]]
if ("embeddingsByType" in _parsed_response and
"binary" in _parsed_response["embeddingsByType"]):
# Use binary data if available (for encoding_format="base64")
embedding_data = _parsed_response["embeddingsByType"]["binary"]
elif ("embeddingsByType" in _parsed_response and
"float" in _parsed_response["embeddingsByType"]):
# Use float data from embeddingsByType
embedding_data = _parsed_response["embeddingsByType"]["float"]
elif "embedding" in _parsed_response:
# Fallback to legacy embedding field
embedding_data = _parsed_response["embedding"]
else:
raise ValueError(f"No embedding data found in response: {response}")
transformed_responses.append(
Embedding(
embedding=_parsed_response["embedding"],
embedding=embedding_data,
index=index,
object="embedding",
)
litellm/proxy/_new_secret_config.yaml
-15
View File
@@ -15,18 +15,3 @@ model_list:
model: hosted_vllm/whisper-v3
api_base: "https://webhook.site/2f385e05-00aa-402b-86d1-efc9261471a5"
api_key: dummy
guardrails:
- guardrail_name: "intel-bedrock-guard-cfg"
litellm_params:
guardrail: bedrock
mode: [pre_call, post_call]
guardrailIdentifier: "1234"
guardrailVersion: "1"
aws_access_key_id: "os.environ/AWS_ACCESS_KEY_ID"
aws_secret_access_key: "os.environ/AWS_SECRET_ACCESS_KEY"
aws_bedrock_runtime_endpoint: "os.environ/AWS_BEDROCK_RUNTIME_ENDPOINT"
default_on: true
litellm/proxy/litellm_pre_call_utils.py
+26 -18
View File
@@ -21,7 +21,9 @@ from litellm.proxy._types import (
)
# Cache special headers as a frozenset for O(1) lookup performance
_SPECIAL_HEADERS_CACHE = frozenset(v.value.lower() for v in SpecialHeaders._member_map_.values())
_SPECIAL_HEADERS_CACHE = frozenset(
v.value.lower() for v in SpecialHeaders._member_map_.values()
)
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.router import Router
from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS
@@ -64,6 +66,7 @@ LITELLM_METADATA_ROUTES = (
"files",
)
def _get_metadata_variable_name(request: Request) -> str:
"""
Helper to return what the "metadata" field should be called in the request data
@@ -157,6 +160,7 @@ class KeyAndTeamLoggingSettings:
@staticmethod
def get_team_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth):
if (
user_api_key_dict.team_metadata is not None
and "logging" in user_api_key_dict.team_metadata
@@ -169,12 +173,12 @@ def _get_dynamic_logging_metadata(
user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig
) -> Optional[TeamCallbackMetadata]:
callback_settings_obj: Optional[TeamCallbackMetadata] = None
key_dynamic_logging_settings: Optional[
dict
] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict)
team_dynamic_logging_settings: Optional[
dict
] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict)
key_dynamic_logging_settings: Optional[dict] = (
KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict)
)
team_dynamic_logging_settings: Optional[dict] = (
KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict)
)
#########################################################################################
# Key-based callbacks
#########################################################################################
@@ -234,12 +238,16 @@ def clean_headers(
Removes litellm api key from headers
"""
clean_headers = {}
litellm_key_lower = litellm_key_header_name.lower() if litellm_key_header_name is not None else None
litellm_key_lower = (
litellm_key_header_name.lower() if litellm_key_header_name is not None else None
)
for header, value in headers.items():
header_lower = header.lower()
# 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 (litellm_key_lower is None or header_lower != litellm_key_lower)):
if header_lower not in _SPECIAL_HEADERS_CACHE and (
litellm_key_lower is None or header_lower != litellm_key_lower
):
clean_headers[header] = value
return clean_headers
@@ -614,11 +622,11 @@ class LiteLLMProxyRequestSetup:
## KEY-LEVEL SPEND LOGS / TAGS
if "tags" in key_metadata and key_metadata["tags"] is not None:
data[_metadata_variable_name][
"tags"
] = LiteLLMProxyRequestSetup._merge_tags(
request_tags=data[_metadata_variable_name].get("tags"),
tags_to_add=key_metadata["tags"],
data[_metadata_variable_name]["tags"] = (
LiteLLMProxyRequestSetup._merge_tags(
request_tags=data[_metadata_variable_name].get("tags"),
tags_to_add=key_metadata["tags"],
)
)
if "spend_logs_metadata" in key_metadata and isinstance(
key_metadata["spend_logs_metadata"], dict
@@ -847,9 +855,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915
data[_metadata_variable_name]["litellm_api_version"] = version
if general_settings is not None:
data[_metadata_variable_name][
"global_max_parallel_requests"
] = general_settings.get("global_max_parallel_requests", None)
data[_metadata_variable_name]["global_max_parallel_requests"] = (
general_settings.get("global_max_parallel_requests", None)
)
### KEY-LEVEL Controls
key_metadata = user_api_key_dict.metadata
litellm/proxy/management_endpoints/team_callback_endpoints.py
+27 -64
View File
@@ -6,7 +6,7 @@ Use this when each team should control its own callbacks
import json
import traceback
from typing import Optional
from typing import List, Optional
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
@@ -79,10 +79,14 @@ async def add_team_callbacks(
"""
try:
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
# Check if team_id exists already
_existing_team = await prisma_client.get_data(
@@ -98,65 +102,30 @@ async def add_team_callbacks(
# store team callback settings in metadata
team_metadata = _existing_team.metadata
team_callback_settings = team_metadata.get("callback_settings", {})
# expect callback settings to be
team_callback_settings_obj = TeamCallbackMetadata(**team_callback_settings)
if data.callback_type == "success":
if team_callback_settings_obj.success_callback is None:
team_callback_settings_obj.success_callback = []
team_callback_settings: List[dict] = team_metadata.get(
"logging"
) # will be dict of type AddTeamCallback
if team_callback_settings is None or not isinstance(
team_callback_settings, list
):
team_callback_settings = []
if data.callback_name in team_callback_settings_obj.success_callback:
## check if it already exists, for the same callback event
for callback in team_callback_settings:
if (
callback.get("callback_name") == data.callback_name
and callback.get("callback_type") == data.callback_type
):
raise ProxyException(
message=f"callback_name = {data.callback_name} already exists in failure_callback, for team_id = {team_id}. \n Existing failure_callback = {team_callback_settings_obj.success_callback}",
message=f"callback_name = {data.callback_name} already exists in team_callback_settings, for team_id = {team_id} and event = {data.callback_type}",
code=status.HTTP_400_BAD_REQUEST,
type=ProxyErrorTypes.bad_request_error,
param="callback_name",
)
team_callback_settings_obj.success_callback.append(data.callback_name)
elif data.callback_type == "failure":
if team_callback_settings_obj.failure_callback is None:
team_callback_settings_obj.failure_callback = []
team_callback_settings.append(data.model_dump())
if data.callback_name in team_callback_settings_obj.failure_callback:
raise ProxyException(
message=f"callback_name = {data.callback_name} already exists in failure_callback, for team_id = {team_id}. \n Existing failure_callback = {team_callback_settings_obj.failure_callback}",
code=status.HTTP_400_BAD_REQUEST,
type=ProxyErrorTypes.bad_request_error,
param="callback_name",
)
team_callback_settings_obj.failure_callback.append(data.callback_name)
elif data.callback_type == "success_and_failure":
if team_callback_settings_obj.success_callback is None:
team_callback_settings_obj.success_callback = []
if team_callback_settings_obj.failure_callback is None:
team_callback_settings_obj.failure_callback = []
if data.callback_name in team_callback_settings_obj.success_callback:
raise ProxyException(
message=f"callback_name = {data.callback_name} already exists in success_callback, for team_id = {team_id}. \n Existing success_callback = {team_callback_settings_obj.success_callback}",
code=status.HTTP_400_BAD_REQUEST,
type=ProxyErrorTypes.bad_request_error,
param="callback_name",
)
if data.callback_name in team_callback_settings_obj.failure_callback:
raise ProxyException(
message=f"callback_name = {data.callback_name} already exists in failure_callback, for team_id = {team_id}. \n Existing failure_callback = {team_callback_settings_obj.failure_callback}",
code=status.HTTP_400_BAD_REQUEST,
type=ProxyErrorTypes.bad_request_error,
param="callback_name",
)
team_callback_settings_obj.success_callback.append(data.callback_name)
team_callback_settings_obj.failure_callback.append(data.callback_name)
for var, value in data.callback_vars.items():
if team_callback_settings_obj.callback_vars is None:
team_callback_settings_obj.callback_vars = {}
team_callback_settings_obj.callback_vars[var] = value
team_callback_settings_obj_dict = team_callback_settings_obj.model_dump()
team_metadata["callback_settings"] = team_callback_settings_obj_dict
team_metadata["logging"] = team_callback_settings
team_metadata_json = json.dumps(team_metadata) # update team_metadata
new_team_row = await prisma_client.db.litellm_teamtable.update(
@@ -168,22 +137,16 @@ async def add_team_callbacks(
"data": new_team_row,
}
except HTTPException as e:
raise e
except ProxyException as e:
raise e
except Exception as e:
verbose_proxy_logger.error(
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.add_team_callbacks(): Exception occured - {}".format(
str(e)
)
)
verbose_proxy_logger.debug(traceback.format_exc())
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Internal Server Error({str(e)})"),
type=ProxyErrorTypes.internal_server_error.value,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
)
elif isinstance(e, ProxyException):
raise e
raise ProxyException(
message="Internal Server Error, " + str(e),
type=ProxyErrorTypes.internal_server_error.value,
litellm/router.py
+36 -33
View File
@@ -360,9 +360,9 @@ class Router:
) # names of models under litellm_params. ex. azure/chatgpt-v-2
self.deployment_latency_map = {}
### CACHING ###
cache_type: Literal[
"local", "redis", "redis-semantic", "s3", "disk"
] = "local" # default to an in-memory cache
cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = (
"local" # default to an in-memory cache
)
redis_cache = None
cache_config: Dict[str, Any] = {}
@@ -404,9 +404,9 @@ class Router:
self.default_max_parallel_requests = default_max_parallel_requests
self.provider_default_deployment_ids: List[str] = []
self.pattern_router = PatternMatchRouter()
self.team_pattern_routers: Dict[
str, PatternMatchRouter
] = {} # {"TEAM_ID": PatternMatchRouter}
self.team_pattern_routers: Dict[str, PatternMatchRouter] = (
{}
) # {"TEAM_ID": PatternMatchRouter}
self.auto_routers: Dict[str, "AutoRouter"] = {}
if model_list is not None:
@@ -588,9 +588,9 @@ class Router:
)
)
self.model_group_retry_policy: Optional[
Dict[str, RetryPolicy]
] = model_group_retry_policy
self.model_group_retry_policy: Optional[Dict[str, RetryPolicy]] = (
model_group_retry_policy
)
self.allowed_fails_policy: Optional[AllowedFailsPolicy] = None
if allowed_fails_policy is not None:
@@ -1212,7 +1212,10 @@ class Router:
async def _acompletion(
self, model: str, messages: List[Dict[str, str]], **kwargs
) -> Union[ModelResponse, CustomStreamWrapper,]:
) -> Union[
ModelResponse,
CustomStreamWrapper,
]:
"""
- Get an available deployment
- call it with a semaphore over the call
@@ -3156,9 +3159,9 @@ class Router:
healthy_deployments=healthy_deployments, responses=responses
)
returned_response = cast(OpenAIFileObject, responses[0])
returned_response._hidden_params[
"model_file_id_mapping"
] = model_file_id_mapping
returned_response._hidden_params["model_file_id_mapping"] = (
model_file_id_mapping
)
return returned_response
except Exception as e:
verbose_router_logger.exception(
@@ -3721,11 +3724,11 @@ class Router:
if isinstance(e, litellm.ContextWindowExceededError):
if context_window_fallbacks is not None:
context_window_fallback_model_group: Optional[
List[str]
] = self._get_fallback_model_group_from_fallbacks(
fallbacks=context_window_fallbacks,
model_group=model_group,
context_window_fallback_model_group: Optional[List[str]] = (
self._get_fallback_model_group_from_fallbacks(
fallbacks=context_window_fallbacks,
model_group=model_group,
)
)
if context_window_fallback_model_group is None:
raise original_exception
@@ -3757,11 +3760,11 @@ class Router:
e.message += "\n{}".format(error_message)
elif isinstance(e, litellm.ContentPolicyViolationError):
if content_policy_fallbacks is not None:
content_policy_fallback_model_group: Optional[
List[str]
] = self._get_fallback_model_group_from_fallbacks(
fallbacks=content_policy_fallbacks,
model_group=model_group,
content_policy_fallback_model_group: Optional[List[str]] = (
self._get_fallback_model_group_from_fallbacks(
fallbacks=content_policy_fallbacks,
model_group=model_group,
)
)
if content_policy_fallback_model_group is None:
raise original_exception
@@ -4415,7 +4418,7 @@ class Router:
return tpm_key
except Exception as e:
verbose_router_logger.exception(
verbose_router_logger.debug(
"litellm.router.Router::deployment_callback_on_success(): Exception occured - {}".format(
str(e)
)
@@ -4993,26 +4996,26 @@ class Router:
"""
from litellm.router_strategy.auto_router.auto_router import AutoRouter
auto_router_config_path: Optional[
str
] = deployment.litellm_params.auto_router_config_path
auto_router_config_path: Optional[str] = (
deployment.litellm_params.auto_router_config_path
)
auto_router_config: Optional[str] = deployment.litellm_params.auto_router_config
if auto_router_config_path is None and auto_router_config is None:
raise ValueError(
"auto_router_config_path or auto_router_config is required for auto-router deployments. Please set it in the litellm_params"
)
default_model: Optional[
str
] = deployment.litellm_params.auto_router_default_model
default_model: Optional[str] = (
deployment.litellm_params.auto_router_default_model
)
if default_model is None:
raise ValueError(
"auto_router_default_model is required for auto-router deployments. Please set it in the litellm_params"
)
embedding_model: Optional[
str
] = deployment.litellm_params.auto_router_embedding_model
embedding_model: Optional[str] = (
deployment.litellm_params.auto_router_embedding_model
)
if embedding_model is None:
raise ValueError(
"auto_router_embedding_model is required for auto-router deployments. Please set it in the litellm_params"
litellm/types/integrations/langsmith.py
+2 -2
View File
@@ -28,8 +28,8 @@ class LangsmithInputs(BaseModel):
class LangsmithCredentialsObject(TypedDict):
LANGSMITH_API_KEY: str
LANGSMITH_PROJECT: str
LANGSMITH_API_KEY: Optional[str]
LANGSMITH_PROJECT: Optional[str]
LANGSMITH_BASE_URL: str
litellm/types/llms/bedrock.py
+12 -5
View File
@@ -328,15 +328,22 @@ class CohereEmbeddingResponse(TypedDict):
texts: List[str]
class AmazonTitanV2EmbeddingRequest(TypedDict):
inputText: str
class AmazonTitanV2EmbeddingRequest(TypedDict, total=False):
inputText: Required[str]
dimensions: int
normalize: bool
embeddingTypes: List[Literal["float", "binary"]]
class AmazonTitanV2EmbeddingResponse(TypedDict):
embedding: List[float]
inputTextTokenCount: int
class AmazonTitanV2EmbeddingsByType(TypedDict, total=False):
binary: List[int] # Array of integers for binary format
float: List[float] # Array of floats for float format
class AmazonTitanV2EmbeddingResponse(TypedDict, total=False):
embedding: List[float] # Legacy field - array of floats (backward compatibility)
embeddingsByType: AmazonTitanV2EmbeddingsByType # New format per AWS schema
inputTextTokenCount: Required[int] # Always present in AWS response
class AmazonTitanG1EmbeddingRequest(TypedDict):
litellm/types/utils.py
+11
View File
@@ -862,6 +862,11 @@ class CompletionTokensDetailsWrapper(
"""Text tokens generated by the model."""
class CacheCreationTokenDetails(BaseModel):
ephemeral_5m_input_tokens: Optional[int] = None
ephemeral_1h_input_tokens: Optional[int] = None
class PromptTokensDetailsWrapper(
PromptTokensDetails
): # wrapper for older openai versions
@@ -886,6 +891,9 @@ class PromptTokensDetailsWrapper(
cache_creation_tokens: Optional[int] = None
"""Number of cache creation tokens sent to the model. Used for Anthropic prompt caching."""
cache_creation_token_details: Optional[CacheCreationTokenDetails] = None
"""Details of cache creation tokens sent to the model. Used for tracking 5m/1h cache creation tokens for Anthropic prompt caching."""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
if self.character_count is None:
@@ -898,6 +906,8 @@ class PromptTokensDetailsWrapper(
del self.web_search_requests
if self.cache_creation_tokens is None:
del self.cache_creation_tokens
if self.cache_creation_token_details is None:
del self.cache_creation_token_details
class ServerToolUse(BaseModel):
@@ -2131,6 +2141,7 @@ class StandardCallbackDynamicParams(TypedDict, total=False):
langsmith_api_key: Optional[str]
langsmith_project: Optional[str]
langsmith_base_url: Optional[str]
langsmith_sampling_rate: Optional[float]
# Humanloop dynamic params
humanloop_api_key: Optional[str]
litellm/utils.py
+3
View File
@@ -4880,6 +4880,9 @@ def _get_model_info_helper( # noqa: PLR0915
cache_read_input_token_cost=_model_info.get(
"cache_read_input_token_cost", None
),
cache_creation_input_token_cost_above_1hr=_model_info.get(
"cache_creation_input_token_cost_above_1hr", None
),
input_cost_per_character=_model_info.get(
"input_cost_per_character", None
),
tests/local_testing/test_anthropic_prompt_caching.py
+2 -3
View File
@@ -320,7 +320,7 @@ async def test_anthropic_api_prompt_caching_basic_with_cache_creation():
random_id
)
* 400,
"cache_control": {"type": "ephemeral"},
"cache_control": {"type": "ephemeral", "ttl": "1h"},
}
],
},
@@ -331,7 +331,7 @@ async def test_anthropic_api_prompt_caching_basic_with_cache_creation():
{
"type": "text",
"text": "What are the key terms and conditions in this agreement?",
"cache_control": {"type": "ephemeral"},
"cache_control": {"type": "ephemeral", "ttl": "5m"},
}
],
},
@@ -580,7 +580,6 @@ async def test_anthropic_api_prompt_caching_streaming():
if hasattr(chunk, "usage") and hasattr(
chunk.usage, "cache_creation_input_tokens"
):
print("chunk.usage", chunk.usage)
is_cache_creation_input_tokens_in_usage = True
idx += 1
tests/local_testing/test_get_optional_params_embeddings.py
+1 -1
View File
@@ -68,4 +68,4 @@ def test_bedrock_embed_v2_with_drop_params():
custom_llm_provider=custom_llm_provider,
)
print(f"received optional_params: {optional_params}")
assert optional_params == {"dimensions": 512}
assert optional_params == {"dimensions": 512, "embeddingTypes": ["binary"]}
tests/test_litellm/integrations/test_langsmith_init.py
+134
View File
@@ -0,0 +1,134 @@
import os
import sys
from unittest.mock import MagicMock, patch
import pytest
sys.path.insert(0, os.path.abspath("../.."))
from litellm.integrations.langsmith import LangsmithLogger
class TestLangsmithLoggerInit:
"""Test cases for LangSmith logger initialization, particularly sampling rate handling.
These tests verify that the sampling_rate attribute is set during initialization.
Note: The current implementation has some edge cases in the sampling rate logic.
"""
@patch("asyncio.create_task")
@patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "1"}, clear=False)
def test_langsmith_sampling_rate_parameter_respected_with_valid_env(
self, mock_create_task
):
"""Test that langsmith_sampling_rate parameter is properly set when env var condition is met."""
# When there's a valid integer in env var, the parameter should be used due to 'or' logic
sampling_rate = 0.5
logger = LangsmithLogger(
langsmith_api_key="test-key",
langsmith_project="test-project",
langsmith_sampling_rate=sampling_rate,
)
# With the current 'or' logic and valid env var, the parameter should be used
assert (
logger.sampling_rate == sampling_rate
), f"Expected sampling_rate to be {sampling_rate}, got {logger.sampling_rate}"
@patch("asyncio.create_task")
@patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "1"}, clear=False)
def test_langsmith_sampling_rate_zero_parameter_falls_back_to_env(
self, mock_create_task
):
"""Test that 0.0 parameter falls back to env var due to falsy value."""
# This demonstrates the current behavior where 0.0 is falsy and falls back to env
logger = LangsmithLogger(
langsmith_api_key="test-key",
langsmith_project="test-project",
langsmith_sampling_rate=0.0, # This is falsy!
)
# Due to current 'or' logic, 0.0 falls back to env var
assert (
logger.sampling_rate == 1.0
), f"Expected sampling_rate to fall back to 1.0 from env, got {logger.sampling_rate}"
@patch("asyncio.create_task")
@patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "1"}, clear=False)
def test_langsmith_sampling_rate_from_integer_env_var(self, mock_create_task):
"""Test that sampling rate uses environment variable when parameter not provided and env var is integer."""
logger = LangsmithLogger(
langsmith_api_key="test-key", langsmith_project="test-project"
)
# Should use env var since it's a valid integer
assert (
logger.sampling_rate == 1.0
), f"Expected sampling_rate to be 1.0 from env var, got {logger.sampling_rate}"
@patch("asyncio.create_task")
@patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "0.8"}, clear=False)
def test_langsmith_sampling_rate_decimal_env_var_ignored(self, mock_create_task):
"""Test that decimal environment variables are ignored due to isdigit() check."""
logger = LangsmithLogger(
langsmith_api_key="test-key", langsmith_project="test-project"
)
# Decimal env vars are ignored due to isdigit() check, falls back to 1.0
assert (
logger.sampling_rate == 1.0
), f"Expected sampling_rate to default to 1.0 (decimal env ignored), got {logger.sampling_rate}"
@patch("asyncio.create_task")
@patch.dict(os.environ, {}, clear=True)
def test_langsmith_sampling_rate_default_value(self, mock_create_task):
"""Test that sampling rate defaults to 1.0 when no parameter or env var provided."""
logger = LangsmithLogger(
langsmith_api_key="test-key", langsmith_project="test-project"
)
assert (
logger.sampling_rate == 1.0
), f"Expected default sampling_rate to be 1.0, got {logger.sampling_rate}"
@patch("asyncio.create_task")
@patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "invalid"}, clear=False)
def test_langsmith_sampling_rate_invalid_env_var_defaults(self, mock_create_task):
"""Test that invalid environment variable falls back to default value."""
logger = LangsmithLogger(
langsmith_api_key="test-key", langsmith_project="test-project"
)
assert (
logger.sampling_rate == 1.0
), f"Expected sampling_rate to default to 1.0 with invalid env var, got {logger.sampling_rate}"
@patch("asyncio.create_task")
@patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": ""}, clear=False)
def test_langsmith_sampling_rate_empty_env_var_defaults(self, mock_create_task):
"""Test that empty environment variable falls back to default value."""
logger = LangsmithLogger(
langsmith_api_key="test-key", langsmith_project="test-project"
)
assert (
logger.sampling_rate == 1.0
), f"Expected sampling_rate to default to 1.0 with empty env var, got {logger.sampling_rate}"
@patch("asyncio.create_task")
def test_langsmith_sampling_rate_attribute_exists(self, mock_create_task):
"""Test that the sampling_rate attribute is always set on the logger instance."""
logger = LangsmithLogger(
langsmith_api_key="test-key", langsmith_project="test-project"
)
# Verify the attribute exists and is a float
assert hasattr(
logger, "sampling_rate"
), "LangsmithLogger should have sampling_rate attribute"
assert isinstance(
logger.sampling_rate, float
), f"sampling_rate should be a float, got {type(logger.sampling_rate)}"
assert (
logger.sampling_rate >= 0.0
), f"sampling_rate should be non-negative, got {logger.sampling_rate}"
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
+81 -2
View File
@@ -22,8 +22,11 @@ sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
from litellm.types.utils import Usage
from litellm.litellm_core_utils.llm_cost_calc.utils import (
calculate_cache_writing_cost,
generic_cost_per_token,
)
from litellm.types.utils import CacheCreationTokenDetails, Usage
def test_reasoning_tokens_no_price_set():
@@ -385,3 +388,79 @@ def test_string_cost_values_with_threshold():
assert round(prompt_cost, 12) == round(expected_prompt_cost, 12)
assert round(completion_cost, 12) == round(expected_completion_cost, 12)
def test_calculate_cache_writing_cost():
"""Test the calculate_cache_writing_cost function with detailed cache creation token breakdown."""
# Test case 1: With cache creation token details (matching the provided input)
cache_creation_tokens = 14055
cache_creation_token_details = CacheCreationTokenDetails(
ephemeral_5m_input_tokens=56, ephemeral_1h_input_tokens=13999
)
cache_creation_cost_above_1hr = 6e-06
cache_creation_cost = 3.75e-06
result = calculate_cache_writing_cost(
cache_creation_tokens=cache_creation_tokens,
cache_creation_token_details=cache_creation_token_details,
cache_creation_cost_above_1hr=cache_creation_cost_above_1hr,
cache_creation_cost=cache_creation_cost,
)
# Expected calculation:
# 5m tokens: 56 * 3.75e-06 = 0.00021
# 1h tokens: 13999 * 6e-06 = 0.083994
# Total: 0.00021 + 0.083994 = 0.084204
expected_cost = (56 * 3.75e-06) + (13999 * 6e-06)
assert round(result, 6) == round(expected_cost, 6)
assert round(result, 6) == 0.084204
# Test case 2: Without cache creation token details (fallback behavior)
cache_creation_tokens_no_details = 1000
cache_creation_token_details_none = None
cache_creation_cost_fallback = 5e-06
result_no_details = calculate_cache_writing_cost(
cache_creation_tokens=cache_creation_tokens_no_details,
cache_creation_token_details=cache_creation_token_details_none,
cache_creation_cost_above_1hr=cache_creation_cost_above_1hr,
cache_creation_cost=cache_creation_cost_fallback,
)
# Expected calculation when no details: 1000 * 5e-06 = 0.005
expected_cost_no_details = 1000 * 5e-06
assert round(result_no_details, 6) == round(expected_cost_no_details, 6)
assert result_no_details == 0.005
# Test case 3: With cache creation token details but None values
cache_creation_token_details_partial = CacheCreationTokenDetails(
ephemeral_5m_input_tokens=None, ephemeral_1h_input_tokens=100
)
result_partial = calculate_cache_writing_cost(
cache_creation_tokens=500,
cache_creation_token_details=cache_creation_token_details_partial,
cache_creation_cost_above_1hr=6e-06,
cache_creation_cost=3e-06,
)
# Expected calculation: 0 (for None 5m tokens) + (100 * 6e-06) = 0.0006
expected_cost_partial = (0.0) + (100 * 6e-06)
assert round(result_partial, 6) == round(expected_cost_partial, 6)
assert round(result_partial, 6) == 0.0006
# Test case 4: Zero costs
result_zero = calculate_cache_writing_cost(
cache_creation_tokens=1000,
cache_creation_token_details=CacheCreationTokenDetails(
ephemeral_5m_input_tokens=50, ephemeral_1h_input_tokens=950
),
cache_creation_cost_above_1hr=0.0,
cache_creation_cost=0.0,
)
assert result_zero == 0.0
tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py
+85 -3
View File
@@ -2,11 +2,12 @@ import json
import os
import sys
from unittest.mock import Mock, patch
import pytest
sys.path.insert(0, os.path.abspath("../../../../..")) # Adds the parent directory to the system path
import litellm
from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
# Mock responses for different embedding models
titan_embedding_response = {
@@ -146,7 +147,7 @@ def test_bedrock_embedding_with_sigv4():
"""Test embedding falls back to SigV4 auth when no bearer token is provided"""
litellm.set_verbose = True
model = "bedrock/amazon.titan-embed-text-v1"
with patch("litellm.llms.bedrock.embed.embedding.BedrockEmbedding.embeddings") as mock_bedrock_embed:
mock_embedding_response = litellm.EmbeddingResponse()
mock_embedding_response.data = [{"embedding": [0.1, 0.2, 0.3]}]
@@ -159,4 +160,85 @@ def test_bedrock_embedding_with_sigv4():
)
assert isinstance(response, litellm.EmbeddingResponse)
mock_bedrock_embed.assert_called_once()
mock_bedrock_embed.assert_called_once()
def test_bedrock_titan_v2_encoding_format_float():
"""Test amazon.titan-embed-text-v2:0 with encoding_format=float parameter"""
litellm.set_verbose = True
client = HTTPHandler()
test_api_key = "test-bearer-token-12345"
model = "bedrock/amazon.titan-embed-text-v2:0"
# Mock response with embeddingsByType for binary format (addressing issue #14680)
titan_v2_response = {
"embedding": [0.1, 0.2, 0.3],
"inputTextTokenCount": 10
}
with patch.object(client, "post") as mock_post:
mock_response = Mock()
mock_response.status_code = 200
mock_response.text = json.dumps(titan_v2_response)
mock_response.json = lambda: json.loads(mock_response.text)
mock_post.return_value = mock_response
response = litellm.embedding(
model=model,
input=test_input,
encoding_format="float", # This should work but currently throws UnsupportedParamsError
client=client,
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
api_key=test_api_key
)
assert isinstance(response, litellm.EmbeddingResponse)
assert isinstance(response.data[0]['embedding'], list)
assert len(response.data[0]['embedding']) == 3
# Verify that the request contains embeddingTypes: ["float"] instead of encoding_format
request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}"))
assert "embeddingTypes" in request_body
assert request_body["embeddingTypes"] == ["float"]
assert "encoding_format" not in request_body
def test_bedrock_titan_v2_encoding_format_base64():
"""Test amazon.titan-embed-text-v2:0 with encoding_format=base64 parameter (maps to binary)"""
litellm.set_verbose = True
client = HTTPHandler()
test_api_key = "test-bearer-token-12345"
model = "bedrock/amazon.titan-embed-text-v2:0"
# Mock response with embeddingsByType for binary format
titan_v2_binary_response = {
"embeddingsByType": {
"binary": "YmluYXJ5X2VtYmVkZGluZ19kYXRh" # base64 encoded binary data
},
"inputTextTokenCount": 10
}
with patch.object(client, "post") as mock_post:
mock_response = Mock()
mock_response.status_code = 200
mock_response.text = json.dumps(titan_v2_binary_response)
mock_response.json = lambda: json.loads(mock_response.text)
mock_post.return_value = mock_response
response = litellm.embedding(
model=model,
input=test_input,
encoding_format="base64", # This should map to embeddingTypes: ["binary"]
client=client,
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
api_key=test_api_key
)
assert isinstance(response, litellm.EmbeddingResponse)
# Verify that the request contains embeddingTypes: ["binary"] for base64 encoding
request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}"))
assert "embeddingTypes" in request_body
assert request_body["embeddingTypes"] == ["binary"]