scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge branch 'BerriAI:main' into LangfuseUsageDetails

Browse source
This commit is contained in:
Fabrício Ceschin 2025-09-18 15:21:11 -04:00 • committed by GitHub
commit 56f7111a57
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
22 changed files with 325 additions and 776 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/__init__.py
View file

@ -1045,7 +1045,6 @@ from .llms.databricks.chat.transformation import DatabricksConfig
from .llms.databricks.embed.transformation import DatabricksEmbeddingConfig
from .llms.predibase.chat.transformation import PredibaseConfig
from .llms.replicate.chat.transformation import ReplicateConfig
from .llms.cohere.completion.transformation import CohereTextConfig as CohereConfig
from .llms.snowflake.chat.transformation import SnowflakeConfig
from .llms.cohere.rerank.transformation import CohereRerankConfig
from .llms.cohere.rerank_v2.transformation import CohereRerankV2Config

5
litellm/constants.py
View file

@ -822,6 +822,7 @@ bedrock_embedding_models: set = set(
"amazon.titan-embed-text-v1",
"cohere.embed-english-v3",
"cohere.embed-multilingual-v3",
"twelvelabs.marengo-embed-2-7-v1:0",
]
)
@ -1065,4 +1066,6 @@ SENTRY_PII_DENYLIST = [
]
# CoroutineChecker cache configuration
COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY = int(os.getenv("COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY", 1000))
COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY = int(
os.getenv("COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY", 1000)
)

4
litellm/litellm_core_utils/get_supported_openai_params.py
View file

@ -94,9 +94,7 @@ def get_supported_openai_params( # noqa: PLR0915
return litellm.VLLMConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "deepseek":
return litellm.DeepSeekChatConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "cohere":
return litellm.CohereConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "cohere_chat":
elif custom_llm_provider == "cohere_chat" or custom_llm_provider == "cohere":
return litellm.CohereChatConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "maritalk":
return litellm.MaritalkConfig().get_supported_openai_params(model=model)

11
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -3079,7 +3079,6 @@ class BedrockConverseMessagesProcessor:
messages.append(DEFAULT_USER_CONTINUE_MESSAGE)
return messages
@staticmethod
async def _bedrock_converse_messages_pt_async( # noqa: PLR0915
messages: List,
@ -3124,9 +3123,9 @@ class BedrockConverseMessagesProcessor:
_part = BedrockContentBlock(text=element["text"])
_parts.append(_part)
elif element["type"] == "guarded_text":
# Wrap guarded_text in guardrailConverseContent block
# Wrap guarded_text in guardContent block
_part = BedrockContentBlock(
guardrailConverseContent={"text": element["text"]}
guardContent={"text": {"text": element["text"]}}
)
_parts.append(_part)
elif element["type"] == "image_url":
@ -3171,7 +3170,6 @@ class BedrockConverseMessagesProcessor:
msg_i += 1
if user_content:
if len(contents) > 0 and contents[-1]["role"] == "user":
if (
assistant_continue_message is not None
@ -3506,9 +3504,9 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915
_part = BedrockContentBlock(text=element["text"])
_parts.append(_part)
elif element["type"] == "guarded_text":
# Wrap guarded_text in guardrailConverseContent block
# Wrap guarded_text in guardContent block
_part = BedrockContentBlock(
guardrailConverseContent={"text": element["text"]}
guardContent={"text": {"text": element["text"]}}
)
_parts.append(_part)
elif element["type"] == "image_url":
@ -3554,7 +3552,6 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915
msg_i += 1
if user_content:
if len(contents) > 0 and contents[-1]["role"] == "user":
if (
assistant_continue_message is not None

28
litellm/llms/bedrock/count_tokens/handler.py
View file

@ -8,8 +8,10 @@ from typing import Any, Dict
from fastapi import HTTPException
import litellm
from litellm._logging import verbose_logger
from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
class BedrockCountTokensHandler(BedrockCountTokensConfig):
@ -78,28 +80,26 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
model=resolved_model,
)
# Make HTTP request
import httpx
async_client = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK)
async with httpx.AsyncClient() as client:
response = await client.post(
response = await async_client.post(
endpoint_url,
headers=signed_headers,
content=signed_body,
data=signed_body,
timeout=30.0,
)
verbose_logger.debug(f"Response status: {response.status_code}")
verbose_logger.debug(f"Response status: {response.status_code}")
if response.status_code != 200:
error_text = response.text
verbose_logger.error(f"AWS Bedrock error: {error_text}")
raise HTTPException(
status_code=400,
detail={"error": f"AWS Bedrock error: {error_text}"},
)
if response.status_code != 200:
error_text = response.text
verbose_logger.error(f"AWS Bedrock error: {error_text}")
raise HTTPException(
status_code=400,
detail={"error": f"AWS Bedrock error: {error_text}"},
)
bedrock_response = response.json()
bedrock_response = response.json()
verbose_logger.debug(f"Bedrock response: {bedrock_response}")

60
litellm/llms/bedrock/embed/embedding.py
View file

@ -4,8 +4,8 @@ Handles embedding calls to Bedrock's `/invoke` endpoint
import copy
import json
from typing import Any, Callable, List, Optional, Tuple, Union
import urllib.parse
from typing import Any, Callable, List, Optional, Tuple, Union
import httpx
@ -18,7 +18,11 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
)
from litellm.secret_managers.main import get_secret
from litellm.types.llms.bedrock import AmazonEmbeddingRequest, CohereEmbeddingRequest
from litellm.types.llms.bedrock import (
AmazonEmbeddingRequest,
CohereEmbeddingRequest,
TwelveLabsMarengoEmbeddingRequest,
)
from litellm.types.utils import EmbeddingResponse
from ..base_aws_llm import BaseAWSLLM
@ -29,6 +33,7 @@ from .amazon_titan_multimodal_transformation import (
)
from .amazon_titan_v2_transformation import AmazonTitanV2Config
from .cohere_transformation import BedrockCohereEmbeddingConfig
from .twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig
class BedrockEmbedding(BaseAWSLLM):
@ -164,16 +169,16 @@ class BedrockEmbedding(BaseAWSLLM):
headers = {"Content-Type": "application/json"}
if extra_headers is not None:
headers = {"Content-Type": "application/json", **extra_headers}
prepped = self.get_request_headers(
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
endpoint_url=endpoint_url,
data=json.dumps(data),
headers=headers,
api_key=api_key
)
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
endpoint_url=endpoint_url,
data=json.dumps(data),
headers=headers,
api_key=api_key,
)
## LOGGING
logging_obj.pre_call(
@ -248,16 +253,16 @@ class BedrockEmbedding(BaseAWSLLM):
headers = {"Content-Type": "application/json"}
if extra_headers is not None:
headers = {"Content-Type": "application/json", **extra_headers}
prepped = self.get_request_headers(
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
endpoint_url=endpoint_url,
data=json.dumps(data),
headers=headers,
api_key=api_key,
)
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
endpoint_url=endpoint_url,
data=json.dumps(data),
headers=headers,
api_key=api_key,
)
## LOGGING
logging_obj.pre_call(
@ -336,7 +341,7 @@ class BedrockEmbedding(BaseAWSLLM):
### TRANSFORMATION ###
unencoded_model_id = (
optional_params.pop("model_id", None) or model
) # default to model if not passed
) # default to model if not passed
modelId = urllib.parse.quote(unencoded_model_id, safe="")
aws_region_name = self._get_aws_region_name(
optional_params=optional_params,
@ -394,6 +399,17 @@ class BedrockEmbedding(BaseAWSLLM):
)
)
batch_data.append(transformed_request)
elif provider == "twelvelabs" and model in [
"twelvelabs.marengo-embed-2-7-v1:0",
]:
batch_data = []
for i in input:
twelvelabs_request: (
TwelveLabsMarengoEmbeddingRequest
) = TwelveLabsMarengoEmbeddingConfig()._transform_request(
input=i, inference_params=inference_params
)
batch_data.append(twelvelabs_request)
### SET RUNTIME ENDPOINT ###
endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint(
@ -445,7 +461,7 @@ class BedrockEmbedding(BaseAWSLLM):
headers = {"Content-Type": "application/json"}
if extra_headers is not None:
headers = {"Content-Type": "application/json", **extra_headers}
prepped = self.get_request_headers(
credentials=credentials,
aws_region_name=aws_region_name,

131
litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py Normal file
View file

@ -0,0 +1,131 @@
"""
Transformation logic from OpenAI /v1/embeddings format to Bedrock TwelveLabs Marengo /invoke format.
Why separate file? Make it easy to see how transformation works
Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo.html
"""
from typing import List
from litellm.types.llms.bedrock import (
TwelveLabsMarengoEmbeddingRequest,
)
from litellm.types.utils import Embedding, EmbeddingResponse, Usage
from litellm.utils import get_base64_str, is_base64_encoded
class TwelveLabsMarengoEmbeddingConfig:
"""
Reference - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo.html
Supports text and image inputs for Phase 1.
Video and audio support will be added in Phase 2.
"""
def __init__(self) -> None:
pass
def get_supported_openai_params(self) -> List[str]:
return ["encoding_format", "textTruncate", "embeddingOption"]
def map_openai_params(
self, non_default_params: dict, optional_params: dict
) -> dict:
for k, v in non_default_params.items():
if k == "encoding_format":
# TwelveLabs doesn't have encoding_format, but we can map it to embeddingOption
if v == "float":
optional_params["embeddingOption"] = ["visual-text", "visual-image"]
elif k == "textTruncate":
optional_params["textTruncate"] = v
elif k == "embeddingOption":
optional_params["embeddingOption"] = v
return optional_params
def _transform_request(
self, input: str, inference_params: dict
) -> TwelveLabsMarengoEmbeddingRequest:
"""
Transform OpenAI-style input to TwelveLabs Marengo format.
Phase 1: Supports text and image inputs only.
"""
# Check if input is base64 encoded image
is_encoded = is_base64_encoded(input)
if is_encoded:
# Image input
b64_str = get_base64_str(input)
transformed_request = TwelveLabsMarengoEmbeddingRequest(
inputType="image", mediaSource={"base64String": b64_str}
)
else:
# Text input
transformed_request = TwelveLabsMarengoEmbeddingRequest(
inputType="text", inputText=input
)
# Set default textTruncate if not specified
if "textTruncate" not in inference_params:
transformed_request["textTruncate"] = "end"
# Set default embedding options for Phase 1 (text and image)
if "embeddingOption" not in inference_params:
if is_encoded:
# For images, return both visual-text and visual-image embeddings
transformed_request["embeddingOption"] = ["visual-text", "visual-image"]
else:
# For text, return visual-text embedding
transformed_request["embeddingOption"] = ["visual-text"]
# Apply any additional inference parameters
for k, v in inference_params.items():
if k not in [
"inputType",
"inputText",
"mediaSource",
]: # Don't override core fields
transformed_request[k] = v # type: ignore
return transformed_request
def _transform_response(
self, response_list: List[dict], model: str
) -> EmbeddingResponse:
"""
Transform TwelveLabs response to OpenAI format.
Handles multiple embedding types in the response.
"""
embeddings: List[Embedding] = []
total_tokens = 0
for response in response_list:
if "embedding" in response:
# Single embedding response
embedding = Embedding(
embedding=response["embedding"],
index=len(embeddings),
object="embedding",
)
embeddings.append(embedding)
# Estimate token count (rough approximation)
if "inputTextTokenCount" in response:
total_tokens += response["inputTextTokenCount"]
else:
# Rough estimate: 1 token per 4 characters for text
total_tokens += len(response.get("inputText", "")) // 4
elif "embeddings" in response:
# Multiple embeddings response (from video/audio)
for i, emb in enumerate(response["embeddings"]):
embedding = Embedding(
embedding=emb["embedding"],
index=len(embeddings),
object="embedding",
)
embeddings.append(embedding)
total_tokens += len(emb["embedding"]) // 4 # Rough estimate
usage = Usage(prompt_tokens=total_tokens, total_tokens=total_tokens)
return EmbeddingResponse(data=embeddings, model=model, usage=usage)

5
litellm/llms/cohere/completion/handler.py
View file

@ -1,5 +0,0 @@
"""
Cohere /generate API - uses `llm_http_handler.py` to make httpx requests
Request/Response transformation is handled in `transformation.py`
"""

265
litellm/llms/cohere/completion/transformation.py
View file

@ -1,265 +0,0 @@
import time
from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, List, Optional, Union
import httpx
import litellm
from litellm.litellm_core_utils.prompt_templates.common_utils import (
convert_content_list_to_str,
)
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import Choices, Message, ModelResponse, Usage
from ..common_utils import CohereError
from ..common_utils import ModelResponseIterator as CohereModelResponseIterator
from ..common_utils import validate_environment as cohere_validate_environment
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class CohereTextConfig(BaseConfig):
"""
Reference: https://docs.cohere.com/reference/generate
The class `CohereConfig` provides configuration for the Cohere's API interface. Below are the parameters:
- `num_generations` (integer): Maximum number of generations returned. Default is 1, with a minimum value of 1 and a maximum value of 5.
- `max_tokens` (integer): Maximum number of tokens the model will generate as part of the response. Default value is 20.
- `truncate` (string): Specifies how the API handles inputs longer than maximum token length. Options include NONE, START, END. Default is END.
- `temperature` (number): A non-negative float controlling the randomness in generation. Lower temperatures result in less random generations. Default is 0.75.
- `preset` (string): Identifier of a custom preset, a combination of parameters such as prompt, temperature etc.
- `end_sequences` (array of strings): The generated text gets cut at the beginning of the earliest occurrence of an end sequence, which will be excluded from the text.
- `stop_sequences` (array of strings): The generated text gets cut at the end of the earliest occurrence of a stop sequence, which will be included in the text.
- `k` (integer): Limits generation at each step to top `k` most likely tokens. Default is 0.
- `p` (number): Limits generation at each step to most likely tokens with total probability mass of `p`. Default is 0.
- `frequency_penalty` (number): Reduces repetitiveness of generated tokens. Higher values apply stronger penalties to previously occurred tokens.
- `presence_penalty` (number): Reduces repetitiveness of generated tokens. Similar to frequency_penalty, but this penalty applies equally to all tokens that have already appeared.
- `return_likelihoods` (string): Specifies how and if token likelihoods are returned with the response. Options include GENERATION, ALL and NONE.
- `logit_bias` (object): Used to prevent the model from generating unwanted tokens or to incentivize it to include desired tokens. e.g. {"hello_world": 1233}
"""
num_generations: Optional[int] = None
max_tokens: Optional[int] = None
truncate: Optional[str] = None
temperature: Optional[int] = None
preset: Optional[str] = None
end_sequences: Optional[list] = None
stop_sequences: Optional[list] = None
k: Optional[int] = None
p: Optional[int] = None
frequency_penalty: Optional[int] = None
presence_penalty: Optional[int] = None
return_likelihoods: Optional[str] = None
logit_bias: Optional[dict] = None
def __init__(
self,
num_generations: Optional[int] = None,
max_tokens: Optional[int] = None,
truncate: Optional[str] = None,
temperature: Optional[int] = None,
preset: Optional[str] = None,
end_sequences: Optional[list] = None,
stop_sequences: Optional[list] = None,
k: Optional[int] = None,
p: Optional[int] = None,
frequency_penalty: Optional[int] = None,
presence_penalty: Optional[int] = None,
return_likelihoods: Optional[str] = None,
logit_bias: Optional[dict] = None,
) -> None:
locals_ = locals().copy()
for key, value in locals_.items():
if key != "self" and value is not None:
setattr(self.__class__, key, value)
@classmethod
def get_config(cls):
return super().get_config()
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
return cohere_validate_environment(
headers=headers,
model=model,
messages=messages,
optional_params=optional_params,
api_key=api_key,
)
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BaseLLMException:
return CohereError(status_code=status_code, message=error_message)
def get_supported_openai_params(self, model: str) -> List:
return [
"stream",
"temperature",
"max_tokens",
"logit_bias",
"top_p",
"frequency_penalty",
"presence_penalty",
"stop",
"n",
"extra_headers",
]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
for param, value in non_default_params.items():
if param == "stream":
optional_params["stream"] = value
elif param == "temperature":
optional_params["temperature"] = value
elif param == "max_tokens":
optional_params["max_tokens"] = value
elif param == "n":
optional_params["num_generations"] = value
elif param == "logit_bias":
optional_params["logit_bias"] = value
elif param == "top_p":
optional_params["p"] = value
elif param == "frequency_penalty":
optional_params["frequency_penalty"] = value
elif param == "presence_penalty":
optional_params["presence_penalty"] = value
elif param == "stop":
optional_params["stop_sequences"] = value
return optional_params
def transform_request(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
prompt = " ".join(
convert_content_list_to_str(message=message) for message in messages
)
## Load Config
config = litellm.CohereConfig.get_config()
for k, v in config.items():
if (
k not in optional_params
): # completion(top_k=3) > cohere_config(top_k=3) <- allows for dynamic variables to be passed in
optional_params[k] = v
## Handle Tool Calling
if "tools" in optional_params:
_is_function_call = True
tool_calling_system_prompt = self._construct_cohere_tool_for_completion_api(
tools=optional_params["tools"]
)
optional_params["tools"] = tool_calling_system_prompt
data = {
"model": model,
"prompt": prompt,
**optional_params,
}
return data
def transform_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ModelResponse,
logging_obj: LiteLLMLoggingObj,
request_data: dict,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> ModelResponse:
prompt = " ".join(
convert_content_list_to_str(message=message) for message in messages
)
completion_response = raw_response.json()
choices_list = []
for idx, item in enumerate(completion_response["generations"]):
if len(item["text"]) > 0:
message_obj = Message(content=item["text"])
else:
message_obj = Message(content=None)
choice_obj = Choices(
finish_reason=item["finish_reason"],
index=idx + 1,
message=message_obj,
)
choices_list.append(choice_obj)
model_response.choices = choices_list # type: ignore
## CALCULATING USAGE
prompt_tokens = len(encoding.encode(prompt))
completion_tokens = len(
encoding.encode(model_response["choices"][0]["message"].get("content", ""))
)
model_response.created = int(time.time())
model_response.model = model
usage = Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
)
setattr(model_response, "usage", usage)
return model_response
def _construct_cohere_tool_for_completion_api(
self,
tools: Optional[List] = None,
) -> dict:
if tools is None:
tools = []
return {"tools": tools}
def get_model_response_iterator(
self,
streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse],
sync_stream: bool,
json_mode: Optional[bool] = False,
):
return CohereModelResponseIterator(
streaming_response=streaming_response,
sync_stream=sync_stream,
json_mode=json_mode,
)

42
litellm/main.py
View file

@ -2395,47 +2395,7 @@ def completion( # type: ignore # noqa: PLR0915
)
return response
response = model_response
elif custom_llm_provider == "cohere":
cohere_key = (
api_key
or litellm.cohere_key
or get_secret("COHERE_API_KEY")
or get_secret("CO_API_KEY")
or litellm.api_key
)
api_base = (
api_base
or litellm.api_base
or get_secret("COHERE_API_BASE")
or "https://api.cohere.ai/v1/generate"
)
headers = headers or litellm.headers or {}
if headers is None:
headers = {}
if extra_headers is not None:
headers.update(extra_headers)
response = base_llm_http_handler.completion(
model=model,
stream=stream,
messages=messages,
acompletion=acompletion,
api_base=api_base,
model_response=model_response,
optional_params=optional_params,
litellm_params=litellm_params,
custom_llm_provider="cohere",
timeout=timeout,
headers=headers,
encoding=encoding,
api_key=cohere_key,
logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements
client=client,
)
elif custom_llm_provider == "cohere_chat":
elif custom_llm_provider == "cohere_chat" or custom_llm_provider == "cohere":
cohere_key = (
api_key
or litellm.cohere_key

37
litellm/types/llms/bedrock.py
View file

@ -88,10 +88,14 @@ class BedrockConverseReasoningContentBlockDelta(TypedDict, total=False):
text: str
class GuardrailConverseTextBlock(TypedDict, total=False):
text: str
class GuardrailConverseContentBlock(TypedDict, total=False):
"""Content block for selective guardrail evaluation in Bedrock Converse API"""
text: str
text: GuardrailConverseTextBlock
class ContentBlock(TypedDict, total=False):
@ -103,7 +107,7 @@ class ContentBlock(TypedDict, total=False):
toolUse: ToolUseBlock
cachePoint: CachePointBlock
reasoningContent: BedrockConverseReasoningContentBlock
guardrailConverseContent: GuardrailConverseContentBlock
guardContent: GuardrailConverseContentBlock
class MessageBlock(TypedDict):
@ -360,6 +364,35 @@ class AmazonTitanMultimodalEmbeddingResponse(TypedDict):
message: str # Specifies any errors that occur during generation.
# TwelveLabs Marengo Embed 2.7 types
TWELVELABS_EMBEDDING_INPUT_TYPES = Literal["text", "image", "video", "audio"]
TWELVELABS_EMBEDDING_OPTIONS = Literal["visual-text", "visual-image", "audio"]
class TwelveLabsMediaSource(TypedDict, total=False):
base64String: str
s3Location: dict # {"uri": str, "bucketOwner": str}
class TwelveLabsMarengoEmbeddingRequest(TypedDict, total=False):
inputType: Required[TWELVELABS_EMBEDDING_INPUT_TYPES]
inputText: str
mediaSource: TwelveLabsMediaSource
textTruncate: Literal["end", "none"]
startSec: float
lengthSec: float
useFixedLengthSec: float
minClipSec: int
embeddingOption: List[TWELVELABS_EMBEDDING_OPTIONS]
class TwelveLabsMarengoEmbeddingResponse(TypedDict):
embedding: List[float]
embeddingOption: TWELVELABS_EMBEDDING_OPTIONS
startSec: float
endSec: float
AmazonEmbeddingRequest = Union[
AmazonTitanMultimodalEmbeddingRequest,
AmazonTitanV2EmbeddingRequest,

21
litellm/utils.py
View file

@ -524,8 +524,6 @@ def get_dynamic_callbacks(
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
def function_setup( # noqa: PLR0915
original_function: str, rules_obj, start_time, *args, **kwargs
): # just run once to check if user wants to send their data anywhere - PostHog/Sentry/Slack/etc.
@ -3465,20 +3463,7 @@ def get_optional_params( # noqa: PLR0915
),
)
elif custom_llm_provider == "cohere":
## check if unsupported param passed in
# handle cohere params
optional_params = litellm.CohereConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(
drop_params
if drop_params is not None and isinstance(drop_params, bool)
else False
),
)
elif custom_llm_provider == "cohere_chat":
elif custom_llm_provider == "cohere_chat" or custom_llm_provider == "cohere":
# handle cohere params
optional_params = litellm.CohereChatConfig().map_openai_params(
non_default_params=non_default_params,
@ -6903,10 +6888,8 @@ class ProviderConfigManager:
return litellm.LlamaAPIConfig()
elif litellm.LlmProviders.TEXT_COMPLETION_OPENAI == provider:
return litellm.OpenAITextCompletionConfig()
elif litellm.LlmProviders.COHERE_CHAT == provider:
elif litellm.LlmProviders.COHERE_CHAT == provider or litellm.LlmProviders.COHERE == provider:
return litellm.CohereChatConfig()
elif litellm.LlmProviders.COHERE == provider:
return litellm.CohereConfig()
elif litellm.LlmProviders.SNOWFLAKE == provider:
return litellm.SnowflakeConfig()
elif litellm.LlmProviders.CLARIFAI == provider:

12
model_prices_and_context_window.json
View file

@ -296,6 +296,18 @@
"output_cost_per_token": 0.0,
"output_vector_size": 1024
},
"twelvelabs.marengo-embed-2-7-v1:0": {
"input_cost_per_token": 7e-05,
"litellm_provider": "bedrock",
"max_input_tokens": 77,
"max_tokens": 77,
"mode": "embedding",
"output_cost_per_token": 0.0,
"output_vector_size": 1024,
"supports_embedding_image_input": true,
"supports_image_input": true,
"supports_multimodal_embedding": true
},
"amazon.titan-text-express-v1": {
"input_cost_per_token": 1.3e-06,
"litellm_provider": "bedrock",

11
tests/litellm_utils_tests/test_utils.py
View file

@ -2326,3 +2326,14 @@ def test_get_whitelisted_models():
file.write(f"{model}\n")
print("whitelisted_models written to whitelisted_bedrock_models.txt")
def test_completion_with_no_model():
"""
Ensure error is raised when no model is provided
"""
# test on empty
with pytest.raises(TypeError):
response = litellm.completion(messages=[{"role": "user", "content": "Hello, how are you?"}])

16
tests/llm_translation/test_cohere.py
View file

@ -254,10 +254,17 @@ async def test_cohere_request_body_with_allowed_params():
}
}]
client = AsyncHTTPHandler()
# Create a mock response
mock_response = AsyncMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"text": "I am Command, a language model developed by Cohere.",
"generation_id": "mock-generation-id",
"finish_reason": "COMPLETE"
}
# Mock the post method
with patch.object(client, "post", new=AsyncMock()) as mock_post:
# Mock the AsyncHTTPHandler.post method at the module level
with patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", return_value=mock_response) as mock_post:
try:
await litellm.acompletion(
model="cohere/command",
@ -265,8 +272,7 @@ async def test_cohere_request_body_with_allowed_params():
allowed_openai_params=["tools", "response_format", "reasoning_effort"],
response_format=test_response_format,
reasoning_effort=test_reasoning_effort,
tools=test_tools,
client=client
tools=test_tools
)
except Exception:
pass # We only care about the request body validation

5
tests/local_testing/test_amazing_vertex_completion.py
View file

@ -3026,10 +3026,13 @@ def test_custom_api_base(api_base):
stream=stream,
auth_header=None,
url="my-fake-endpoint",
model="gemini-1.5-pro", # Required for Gemini custom API base URLs
)
if api_base:
assert url == api_base + ":"
# For Gemini with custom API base, URL should be constructed as api_base/models/model:endpoint
expected_url = f"{api_base}/models/gemini-1.5-pro:"
assert url == expected_url
else:
assert url == test_endpoint

106
tests/local_testing/test_bad_params.py
View file

@ -1,106 +0,0 @@
#### What this tests ####
# This tests chaos monkeys - if random parts of the system are broken / things aren't sent correctly - what happens.
# Expect to add more edge cases to this over time.
import os
import sys
import traceback
import pytest
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import litellm
from litellm import completion, embedding
from litellm.utils import Message
# litellm.set_verbose = True
user_message = "Hello, how are you?"
messages = [{"content": user_message, "role": "user"}]
model_val = None
def test_completion_with_no_model():
# test on empty
with pytest.raises(TypeError):
response = completion(messages=messages)
def test_completion_with_empty_model():
# test on empty
try:
response = completion(model=model_val, messages=messages)
except Exception as e:
print(f"error occurred: {e}")
pass
def test_completion_invalid_param_cohere():
try:
litellm.set_verbose = True
response = completion(model="command-nightly", messages=messages, seed=12)
pytest.fail(f"This should have failed cohere does not support `seed` parameter")
except Exception as e:
assert isinstance(e, litellm.UnsupportedParamsError)
print("got an exception=", str(e))
if "cohere does not support parameters: ['seed']" in str(e):
pass
else:
pytest.fail(f"An error occurred {e}")
def test_completion_function_call_cohere():
try:
response = completion(
model="command-nightly", messages=messages, functions=["TEST-FUNCTION"]
)
pytest.fail(f"An error occurred {e}")
except Exception as e:
print(e)
pass
def test_completion_function_call_openai():
try:
messages = [{"role": "user", "content": "What is the weather like in Boston?"}]
response = completion(
model="gpt-3.5-turbo",
messages=messages,
functions=[
{
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA",
},
"unit": {
"type": "string",
"enum": ["celsius", "fahrenheit"],
},
},
"required": ["location"],
},
}
],
)
print(f"response: {response}")
except Exception:
pass
# test_completion_function_call_openai()
def test_completion_with_no_provider():
# test on empty
try:
model = "cerebras/btlm-3b-8k-base"
response = completion(model=model, messages=messages)
except Exception as e:
print(f"error occurred: {e}")
pass

31
tests/logging_callback_tests/test_otel_logging.py
View file

@ -251,7 +251,8 @@ async def test_awesome_otel_with_message_logging_off(streaming, global_redact):
def validate_redacted_message_span_attributes(span):
expected_attributes = [
# Required non-metadata attributes that must be present
required_attributes = [
"gen_ai.request.model",
"gen_ai.system",
"llm.is_streaming",
@ -259,27 +260,8 @@ def validate_redacted_message_span_attributes(span):
"gen_ai.response.id",
"gen_ai.response.model",
"llm.usage.total_tokens",
"metadata.prompt_management_metadata",
"gen_ai.usage.completion_tokens",
"gen_ai.usage.prompt_tokens",
"metadata.user_api_key_hash",
"metadata.requester_ip_address",
"metadata.user_api_key_team_alias",
"metadata.requester_metadata",
"metadata.user_api_key_team_id",
"metadata.spend_logs_metadata",
"metadata.usage_object",
"metadata.user_api_key_alias",
"metadata.user_api_key_user_id",
"metadata.user_api_key_org_id",
"metadata.user_api_key_end_user_id",
"metadata.user_api_key_user_email",
"metadata.user_api_key_request_route",
"metadata.applied_guardrails",
"metadata.mcp_tool_call_metadata",
"metadata.vector_store_request_metadata",
"metadata.requester_custom_headers",
"metadata.cold_storage_object_key",
]
_all_attributes = set(
@ -293,6 +275,13 @@ def validate_redacted_message_span_attributes(span):
for attr in _all_attributes:
print(f"attr: {attr}, type: {type(attr)}")
assert _all_attributes == set(expected_attributes)
# Check that all required attributes are present
required_set = set(required_attributes)
assert required_set.issubset(_all_attributes), f"Missing required attributes: {required_set - _all_attributes}"
# Check that any additional attributes are metadata fields (start with "metadata.")
non_required_attrs = _all_attributes - required_set
for attr in non_required_attrs:
assert attr.startswith("metadata."), f"Non-metadata attribute found: {attr}"
pass

10
tests/otel_tests/test_prometheus.py
View file

@ -106,7 +106,7 @@ async def test_proxy_failure_metrics():
print("/metrics", metrics)
# Check if the failure metric is present and correct - use pattern matching for robustness
expected_metric_pattern = 'litellm_proxy_failed_requests_metric_total{api_key_alias="None",end_user="None",exception_class="Openai.RateLimitError",exception_status="429",hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",requested_model="fake-azure-endpoint",route="/chat/completions",team="None",team_alias="None",user="default_user_id"}'
expected_metric_pattern = 'litellm_proxy_failed_requests_metric_total{api_key_alias="None",end_user="None",exception_class="Openai.RateLimitError",exception_status="429",hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",requested_model="fake-azure-endpoint",route="/chat/completions",team="None",team_alias="None",user="default_user_id",user_email="None"}'
# Check if the pattern is in metrics (this metric doesn't include user_email field)
assert any(expected_metric_pattern in line for line in metrics.split('\n')), f"Expected failure metric pattern not found in /metrics. Pattern: {expected_metric_pattern}"
@ -576,8 +576,8 @@ async def test_user_email_in_all_required_metrics():
Test that user_email label is present in all the metrics that were requested to have it:
- litellm_proxy_total_requests_metric_total
- litellm_proxy_failed_requests_metric_total
- litellm_input_tokens_total
- litellm_output_tokens_total
- litellm_input_tokens_metric_total
- litellm_output_tokens_metric_total
- litellm_requests_metric_total
- litellm_spend_metric_total
"""
@ -608,8 +608,8 @@ async def test_user_email_in_all_required_metrics():
# Check that user_email appears in all the required metrics
required_metrics_with_user_email = [
"litellm_proxy_total_requests_metric_total",
"litellm_input_tokens_total",
"litellm_output_tokens_total",
"litellm_input_tokens_metric_total",
"litellm_output_tokens_metric_total",
"litellm_requests_metric_total",
"litellm_spend_metric_total"
]

26
tests/pass_through_unit_tests/test_pass_through_unit_tests.py
View file

@ -134,21 +134,19 @@ def test_init_kwargs_for_pass_through_endpoint_basic(
assert result["litellm_call_id"] == "test-call-id"
assert result["passthrough_logging_payload"] == passthrough_payload
#########################################################
# Check metadata
expected_metadata = {
"user_api_key": "test-key",
"user_api_key_hash": "test-key",
"user_api_key_alias": None,
"user_api_key_user_email": None,
"user_api_key_user_id": "test-user",
"user_api_key_team_id": "test-team",
"user_api_key_org_id": None,
"user_api_key_team_alias": None,
"user_api_key_end_user_id": "test-user",
"user_api_key_request_route": None,
}
assert result["litellm_params"]["metadata"] == expected_metadata
#########################################################
assert result["litellm_params"]["metadata"]["user_api_key"] == "test-key"
assert result["litellm_params"]["metadata"]["user_api_key_hash"] == "test-key"
assert result["litellm_params"]["metadata"]["user_api_key_alias"] is None
assert result["litellm_params"]["metadata"]["user_api_key_user_email"] is None
assert result["litellm_params"]["metadata"]["user_api_key_user_id"] == "test-user"
assert result["litellm_params"]["metadata"]["user_api_key_team_id"] == "test-team"
assert result["litellm_params"]["metadata"]["user_api_key_org_id"] is None
assert result["litellm_params"]["metadata"]["user_api_key_team_alias"] is None
assert result["litellm_params"]["metadata"]["user_api_key_end_user_id"] == "test-user"
assert result["litellm_params"]["metadata"]["user_api_key_request_route"] is None
def test_init_kwargs_with_litellm_metadata(mock_request, mock_user_api_key_dict):

265
tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
View file

@ -1597,7 +1597,7 @@ async def test_no_cache_control_no_cache_point():
# ============================================================================
def test_guarded_text_wraps_in_guardrail_converse_content():
"""Test that guarded_text content type gets wrapped in guardrailConverseContent blocks."""
"""Test that guarded_text content type gets wrapped in guardContent blocks."""
from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt
messages = [
@ -1631,9 +1631,9 @@ def test_guarded_text_wraps_in_guardrail_converse_content():
assert "text" in content[2]
assert content[2]["text"] == "More regular text"
# Second should be guardrailConverseContent
assert "guardrailConverseContent" in content[1]
assert content[1]["guardrailConverseContent"]["text"] == "This should be guarded"
# Second should be guardContent
assert "guardContent" in content[1]
assert content[1]["guardContent"]["text"]["text"] == "This should be guarded"
def test_guarded_text_with_system_messages():
@ -1685,9 +1685,9 @@ def test_guarded_text_with_system_messages():
assert "text" in content[0]
assert content[0]["text"] == "What is the main topic of this legal document?"
# Second should be guardrailConverseContent
assert "guardrailConverseContent" in content[1]
assert content[1]["guardrailConverseContent"]["text"] == "This is a set of very long instructions that you will follow. Here is a legal document that you will use to answer the user's question."
# Second should be guardContent
assert "guardContent" in content[1]
assert content[1]["guardContent"]["text"]["text"] == "This is a set of very long instructions that you will follow. Here is a legal document that you will use to answer the user's question."
def test_guarded_text_with_mixed_content_types():
@ -1726,9 +1726,9 @@ def test_guarded_text_with_mixed_content_types():
# Second should be image
assert "image" in content[1]
# Third should be guardrailConverseContent
assert "guardrailConverseContent" in content[2]
assert content[2]["guardrailConverseContent"]["text"] == "This sensitive content should be guarded"
# Third should be guardContent
assert "guardContent" in content[2]
assert content[2]["guardContent"]["text"]["text"] == "This sensitive content should be guarded"
@pytest.mark.asyncio
@ -1764,9 +1764,9 @@ async def test_async_guarded_text():
assert "text" in content[0]
assert content[0]["text"] == "Hello"
# Second should be guardrailConverseContent
assert "guardrailConverseContent" in content[1]
assert content[1]["guardrailConverseContent"]["text"] == "This should be guarded"
# Second should be guardContent
assert "guardContent" in content[1]
assert content[1]["guardContent"]["text"]["text"] == "This should be guarded"
def test_guarded_text_with_tool_calls():
@ -1818,15 +1818,15 @@ def test_guarded_text_with_tool_calls():
assert "text" in content[0]
assert content[0]["text"] == "What's the weather?"
# Second should be guardrailConverseContent
assert "guardrailConverseContent" in content[1]
assert content[1]["guardrailConverseContent"]["text"] == "Please be careful with sensitive information"
# Second should be guardContent
assert "guardContent" in content[1]
assert content[1]["guardContent"]["text"]["text"] == "Please be careful with sensitive information"
# Other messages should not have guardrailConverseContent
# Other messages should not have guardContent
for i in range(1, 3):
content = result[i]["content"]
for block in content:
assert "guardrailConverseContent" not in block
assert "guardContent" not in block
def test_guarded_text_guardrail_config_preserved():
@ -2066,234 +2066,11 @@ def test_auto_convert_in_full_transformation():
assert "messages" in result
assert len(result["messages"]) == 1
# The message should have guardrailConverseContent
# The message should have guardContent
message = result["messages"][0]
assert "content" in message
assert len(message["content"]) == 1
assert "guardrailConverseContent" in message["content"][0]
assert message["content"][0]["guardrailConverseContent"]["text"] == "What is the main topic of this legal document?"
assert "guardContent" in message["content"][0]
assert message["content"][0]["guardContent"]["text"]["text"] == "What is the main topic of this legal document?"
def test_convert_consecutive_user_messages_to_guarded_text():
"""Test that consecutive user messages at the end are converted to guarded_text."""
config = AmazonConverseConfig()
messages = [
{
"role": "user",
"content": [
{
"type": "text",
"text": "First user message"
}
]
},
{
"role": "assistant",
"content": "Assistant response"
},
{
"role": "user",
"content": [
{
"type": "text",
"text": "Second user message"
}
]
},
{
"role": "user",
"content": [
{
"type": "text",
"text": "Third user message"
}
]
}
]
optional_params = {
"guardrailConfig": {
"guardrailIdentifier": "gr-abc123",
"guardrailVersion": "1"
}
}
# Test the helper method directly
converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params)
# Verify the conversion - only the last two user messages should be converted
assert len(converted_messages) == 4
# First user message should remain unchanged
assert converted_messages[0]["role"] == "user"
assert converted_messages[0]["content"][0]["type"] == "text"
assert converted_messages[0]["content"][0]["text"] == "First user message"
# Assistant message should remain unchanged
assert converted_messages[1]["role"] == "assistant"
assert converted_messages[1]["content"] == "Assistant response"
# Second user message should be converted to guarded_text
assert converted_messages[2]["role"] == "user"
assert converted_messages[2]["content"][0]["type"] == "guarded_text"
assert converted_messages[2]["content"][0]["text"] == "Second user message"
# Third user message should be converted to guarded_text
assert converted_messages[3]["role"] == "user"
assert converted_messages[3]["content"][0]["type"] == "guarded_text"
assert converted_messages[3]["content"][0]["text"] == "Third user message"
def test_convert_all_user_messages_when_all_consecutive():
"""Test that all user messages are converted when they are all consecutive at the end."""
config = AmazonConverseConfig()
messages = [
{
"role": "user",
"content": [
{
"type": "text",
"text": "First user message"
}
]
},
{
"role": "user",
"content": [
{
"type": "text",
"text": "Second user message"
}
]
},
{
"role": "user",
"content": [
{
"type": "text",
"text": "Third user message"
}
]
}
]
optional_params = {
"guardrailConfig": {
"guardrailIdentifier": "gr-abc123",
"guardrailVersion": "1"
}
}
# Test the helper method directly
converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params)
# Verify all three user messages are converted
assert len(converted_messages) == 3
for i in range(3):
assert converted_messages[i]["role"] == "user"
assert converted_messages[i]["content"][0]["type"] == "guarded_text"
assert converted_messages[0]["content"][0]["text"] == "First user message"
assert converted_messages[1]["content"][0]["text"] == "Second user message"
assert converted_messages[2]["content"][0]["text"] == "Third user message"
def test_convert_consecutive_user_messages_with_string_content():
"""Test that consecutive user messages with string content are converted to guarded_text."""
config = AmazonConverseConfig()
messages = [
{
"role": "assistant",
"content": "Assistant response"
},
{
"role": "user",
"content": "First user message"
},
{
"role": "user",
"content": "Second user message"
}
]
optional_params = {
"guardrailConfig": {
"guardrailIdentifier": "gr-abc123",
"guardrailVersion": "1"
}
}
# Test the helper method directly
converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params)
# Verify the conversion
assert len(converted_messages) == 3
# Assistant message should remain unchanged
assert converted_messages[0]["role"] == "assistant"
assert converted_messages[0]["content"] == "Assistant response"
# Both user messages should be converted to guarded_text
assert converted_messages[1]["role"] == "user"
assert len(converted_messages[1]["content"]) == 1
assert converted_messages[1]["content"][0]["type"] == "guarded_text"
assert converted_messages[1]["content"][0]["text"] == "First user message"
assert converted_messages[2]["role"] == "user"
assert len(converted_messages[2]["content"]) == 1
assert converted_messages[2]["content"][0]["type"] == "guarded_text"
assert converted_messages[2]["content"][0]["text"] == "Second user message"
def test_skip_consecutive_user_messages_with_existing_guarded_text():
"""Test that consecutive user messages with existing guarded_text are skipped."""
config = AmazonConverseConfig()
messages = [
{
"role": "user",
"content": [
{
"type": "guarded_text",
"text": "Already guarded"
}
]
},
{
"role": "user",
"content": [
{
"type": "text",
"text": "Should be converted"
}
]
}
]
optional_params = {
"guardrailConfig": {
"guardrailIdentifier": "gr-abc123",
"guardrailVersion": "1"
}
}
# Test the helper method directly
converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params)
# Verify the conversion
assert len(converted_messages) == 2
# First message should remain unchanged (already has guarded_text)
assert converted_messages[0]["role"] == "user"
assert converted_messages[0]["content"][0]["type"] == "guarded_text"
assert converted_messages[0]["content"][0]["text"] == "Already guarded"
# Second message should be converted
assert converted_messages[1]["role"] == "user"
assert converted_messages[1]["content"][0]["type"] == "guarded_text"
assert converted_messages[1]["content"][0]["text"] == "Should be converted"

9
tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py
View file

@ -19,6 +19,13 @@ cohere_embedding_response = {
"inputTextTokenCount": 10
}
twelvelabs_embedding_response = {
"embedding": [0.1, 0.2, 0.3],
"embeddingOption": "visual-text",
"startSec": 0.0,
"endSec": 1.0
}
# Test data
test_input = "Hello world from litellm"
test_image_base64 = "data:image/png,test_image_base64_data"
@ -32,6 +39,8 @@ test_image_base64 = "data:image/png,test_image_base64_data"
("bedrock/amazon.titan-embed-image-v1", "image", titan_embedding_response),
("bedrock/cohere.embed-english-v3", "text", cohere_embedding_response),
("bedrock/cohere.embed-multilingual-v3", "text", cohere_embedding_response),
("bedrock/twelvelabs.marengo-embed-2-7-v1:0", "text", twelvelabs_embedding_response),
("bedrock/twelvelabs.marengo-embed-2-7-v1:0", "image", twelvelabs_embedding_response),
],
)
def test_bedrock_embedding_with_api_key_bearer_token(model, input_type, embed_response):