diff --git a/litellm/constants.py b/litellm/constants.py index 82ca92475ec..b2d90e20dbb 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1374,6 +1374,7 @@ bedrock_embedding_models: Final[set] = set( "cohere.embed-multilingual-v3", "cohere.embed-v4:0", "twelvelabs.marengo-embed-2-7-v1:0", + "twelvelabs.marengo-embed-3-0-v1:0", ] ) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 68dc27ec25e..c05d4c29a5e 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -780,6 +780,7 @@ class PromptTokensDetailsResult(TypedDict): image_count: int video_length_seconds: float audio_length_seconds: float + query_count: int def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: @@ -828,6 +829,7 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: ) or 0.0 ) + query_count: Final = _coerce_token_count(getattr(usage.prompt_tokens_details, "query_count", 0)) return PromptTokensDetailsResult( cache_hit_tokens=cache_hit_tokens, @@ -841,6 +843,7 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: image_count=image_count, video_length_seconds=float(video_length_seconds), audio_length_seconds=float(audio_length_seconds), + query_count=query_count, ) @@ -978,6 +981,11 @@ def _calculate_input_cost( prompt_tokens_details["audio_length_seconds"], ) + if prompt_tokens_details["query_count"]: + prompt_cost += calculate_cost_component( + model_info, "input_cost_per_query", prompt_tokens_details["query_count"] + ) + return prompt_cost @@ -1149,6 +1157,7 @@ def generic_cost_per_token( image_count=0, video_length_seconds=0.0, audio_length_seconds=0.0, + query_count=0, ) if usage.prompt_tokens_details: prompt_tokens_details = parse_prompt_tokens_details(usage) diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index d3725434498..e249feb9ff3 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -35,7 +35,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 +from .twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig, drop_params_enabled if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -239,7 +239,7 @@ class BedrockEmbedding(BaseAWSLLM): returned_response = AmazonTitanG1Config()._transform_response(response_list=response_list, model=model) elif provider == "twelvelabs": returned_response = TwelveLabsMarengoEmbeddingConfig()._transform_response( - response_list=response_list, model=model + response_list=response_list, model=model, batch_data=batch_data ) elif provider == "nova": returned_response = AmazonNovaEmbeddingConfig()._transform_response( @@ -484,12 +484,13 @@ class BedrockEmbedding(BaseAWSLLM): elif provider == "twelvelabs": batch_data = [] for i in input: - twelvelabs_request = TwelveLabsMarengoEmbeddingConfig()._transform_request( + twelvelabs_request = TwelveLabsMarengoEmbeddingConfig(model=model)._transform_request( input=i, inference_params=inference_params, async_invoke_route=has_async_invoke, model_id=modelId, output_s3_uri=inference_params.get("output_s3_uri"), + drop_params=drop_params_enabled(litellm_params), ) batch_data.append(twelvelabs_request) elif provider == "nova": diff --git a/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py b/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py new file mode 100644 index 00000000000..4aac6f22bae --- /dev/null +++ b/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py @@ -0,0 +1,239 @@ +""" +Request builder for Bedrock TwelveLabs Marengo Embed 3.0, whose payload nests the input under a key named after +``inputType`` instead of the flat 2.7 layout. + +Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo-3.html +""" + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.llms.bedrock.common_utils import BedrockError +from litellm.types.llms.bedrock import ( + TWELVELABS_MARENGO_3_EMBEDDING_OPTIONS, + TWELVELABS_MARENGO_3_EMBEDDING_SCOPES, + TWELVELABS_MARENGO_3_EMBEDDING_TYPES, + TWELVELABS_MARENGO_3_INPUT_TYPES, + TwelveLabsMarengo3AudioRequest, + TwelveLabsMarengo3EmbeddingRequest, + TwelveLabsMarengo3ImageRequest, + TwelveLabsMarengo3MultiInputRequest, + TwelveLabsMarengo3NamedMediaSource, + TwelveLabsMarengo3RequestBase, + TwelveLabsMarengo3Segmentation, + TwelveLabsMarengo3TextImageRequest, + TwelveLabsMarengo3TextRequest, + TwelveLabsMarengo3TimedMediaInput, + TwelveLabsMarengo3TimedMediaOptions, + TwelveLabsMarengo3VideoRequest, + TwelveLabsMediaSource, + TwelveLabsS3Location, +) +from litellm.utils import get_base64_str + +MARENGO_3_MODEL_MARKER: Final = "marengo-embed-3-" +S3_URI_PREFIX: Final = "s3://" +TIMED_MEDIA_OPTION_FIELDS: Final = MappingProxyType( + { + "startSec": True, + "endSec": True, + "segmentation": True, + "embeddingOption": True, + "embeddingType": True, + "embeddingScope": True, + } +) +TIMED_MEDIA_OPTIONS: Final = TypeAdapter(TwelveLabsMarengo3TimedMediaOptions) +TIMED_INPUT_TYPES: Final = frozenset({"video", "audio"}) +MARENGO_2_7_ONLY_PARAMS: Final = ("textTruncate", "lengthSec", "useFixedLengthSec", "minClipSec") +MARENGO_2_7_ONLY_FIELDS: Final = MappingProxyType({name: True for name in MARENGO_2_7_ONLY_PARAMS}) + + +def is_marengo_3_model(model: str | None) -> bool: + return MARENGO_3_MODEL_MARKER in (model or "") + + +class Marengo3Params(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + inputType: TWELVELABS_MARENGO_3_INPUT_TYPES | None = None + input_type: TWELVELABS_MARENGO_3_INPUT_TYPES | None = None + media_source: str | None = None + media_sources: Mapping[str, str] | None = None + bucketOwner: str | None = None + startSec: float | None = None + endSec: float | None = None + segmentation: TwelveLabsMarengo3Segmentation | None = None + embeddingOption: tuple[TWELVELABS_MARENGO_3_EMBEDDING_OPTIONS, ...] | None = None + embeddingType: tuple[TWELVELABS_MARENGO_3_EMBEDDING_TYPES, ...] | None = None + embeddingScope: tuple[TWELVELABS_MARENGO_3_EMBEDDING_SCOPES, ...] | None = None + inferenceId: str | None = None + textTruncate: object = None + lengthSec: object = None + useFixedLengthSec: object = None + minClipSec: object = None + + @property + def resolved_input_type(self) -> TWELVELABS_MARENGO_3_INPUT_TYPES: + return self.inputType or self.input_type or "text" + + def timed_media_options(self) -> TwelveLabsMarengo3TimedMediaOptions: + return TIMED_MEDIA_OPTIONS.validate_python(self.given_timed_media_options()) + + def given_timed_media_options(self) -> dict[str, object]: + return self.model_dump(include=TIMED_MEDIA_OPTION_FIELDS, exclude_none=True) + + def given_2_7_only_params(self) -> dict[str, object]: + return self.model_dump(include=MARENGO_2_7_ONLY_FIELDS, exclude_none=True) + + +def _require_bucket_owner(bucket_owner: str | None) -> str: + if bucket_owner is None: + raise BedrockError( + status_code=400, + message="s3:// media requires the 'bucketOwner' parameter, the account id that owns the bucket", + ) + return bucket_owner + + +def _media_source(media: str, bucket_owner: str | None) -> TwelveLabsMediaSource: + if not media.startswith(S3_URI_PREFIX): + inline: Final[TwelveLabsMediaSource] = {"base64String": get_base64_str(media)} + return inline + s3_location: Final[TwelveLabsS3Location] = {"uri": media, "bucketOwner": _require_bucket_owner(bucket_owner)} + remote: Final[TwelveLabsMediaSource] = {"s3Location": s3_location} + return remote + + +def _named_media_source(name: str, media: str, bucket_owner: str | None) -> TwelveLabsMarengo3NamedMediaSource: + named: Final[TwelveLabsMarengo3NamedMediaSource] = { + "name": name, + "mediaType": "image", + **_media_source(media, bucket_owner), + } + return named + + +def _timed_media_input(media: str, params: Marengo3Params) -> TwelveLabsMarengo3TimedMediaInput: + timed: Final[TwelveLabsMarengo3TimedMediaInput] = { + "mediaSource": _media_source(media, params.bucketOwner), + **params.timed_media_options(), + } + return timed + + +def _describe(error: ValidationError) -> str: + return "; ".join( + f"{'.'.join(str(part) for part in problem['loc'])}: {problem['msg']}" for problem in error.errors() + ) + + +def _validated_params(inference_params: Mapping[str, object]) -> Marengo3Params: + try: + return Marengo3Params.model_validate(inference_params) + except ValidationError as error: + raise BedrockError(status_code=400, message=f"Invalid Marengo 3.0 parameters: {_describe(error)}") from error + + +def _reject_unless_dropped(given: Mapping[str, object], drop_params: bool, reason: str) -> None: + if not given or drop_params: + return + raise BedrockError(status_code=400, message=f"{reason} {', '.join(given)}; set drop_params to drop them") + + +def _require(value: str | None, input_type: str, param_name: str) -> str: + if value is None: + raise BedrockError(status_code=400, message=f"Input type '{input_type}' requires the '{param_name}' parameter") + return value + + +def _require_media_sources(value: Mapping[str, str] | None) -> Mapping[str, str]: + if not value: + raise BedrockError( + status_code=400, + message="Input type 'multi_input' requires a non-empty 'media_sources' mapping of name to media", + ) + return value + + +def _request_base(inference_id: str | None) -> TwelveLabsMarengo3RequestBase: + if inference_id is None: + anonymous: Final[TwelveLabsMarengo3RequestBase] = {} + return anonymous + identified: Final[TwelveLabsMarengo3RequestBase] = {"inferenceId": inference_id} + return identified + + +def build_marengo_3_request( + input: str, inference_params: Mapping[str, object], drop_params: bool = False +) -> TwelveLabsMarengo3EmbeddingRequest: + params: Final = _validated_params(inference_params) + base: Final = _request_base(params.inferenceId) + input_type: Final = params.resolved_input_type + _reject_unless_dropped( + params.given_2_7_only_params(), drop_params, "Marengo 3.0 does not accept the Marengo 2.7 parameters" + ) + if input_type not in TIMED_INPUT_TYPES: + _reject_unless_dropped( + params.given_timed_media_options(), drop_params, f"Input type '{input_type}' does not accept" + ) + match input_type: + case "text": + text_request: Final[TwelveLabsMarengo3TextRequest] = { + **base, + "inputType": "text", + "text": {"inputText": input}, + } + return text_request + case "image": + image_request: Final[TwelveLabsMarengo3ImageRequest] = { + **base, + "inputType": "image", + "image": {"mediaSource": _media_source(input, params.bucketOwner)}, + } + return image_request + case "video": + video_request: Final[TwelveLabsMarengo3VideoRequest] = { + **base, + "inputType": "video", + "video": _timed_media_input(input, params), + } + return video_request + case "audio": + audio_request: Final[TwelveLabsMarengo3AudioRequest] = { + **base, + "inputType": "audio", + "audio": _timed_media_input(input, params), + } + return audio_request + case "text_image": + text_image_request: Final[TwelveLabsMarengo3TextImageRequest] = { + **base, + "inputType": "text_image", + "text_image": { + "inputText": input, + "mediaSource": _media_source( + _require(params.media_source, input_type, "media_source"), params.bucketOwner + ), + }, + } + return text_image_request + case "multi_input": + media_sources: Final = tuple( + _named_media_source(name, media, params.bucketOwner) + for name, media in _require_media_sources(params.media_sources).items() + ) + multi_input_request: Final[TwelveLabsMarengo3MultiInputRequest] = { + **base, + "inputType": "multi_input", + "multi_input": {"inputText": input, "mediaSources": media_sources} + if input + else {"mediaSources": media_sources}, + } + return multi_input_request + case _: + assert_never(input_type) diff --git a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py index a39c59b0efd..65ca2be191f 100644 --- a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py +++ b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py @@ -4,19 +4,120 @@ Transformation logic from OpenAI /v1/embeddings format to Bedrock TwelveLabs Mar Why separate file? Make it easy to see how transformation works Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo.html +Marengo 3.0 docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo-3.html """ +from collections.abc import Mapping from typing import Final, cast +from pydantic import BaseModel, ConfigDict, TypeAdapter +from typing_extensions import assert_never + +import litellm +from litellm.llms.bedrock.embed.twelvelabs_marengo_3_transformation import ( + MARENGO_2_7_ONLY_PARAMS, + build_marengo_3_request, + is_marengo_3_model, +) from litellm.types.llms.bedrock import ( TWELVELABS_EMBEDDING_INPUT_TYPES, + TWELVELABS_MARENGO_3_INPUT_TYPES, TwelveLabsAsyncInvokeRequest, + TwelveLabsMarengo3EmbeddingRequest, TwelveLabsMarengoEmbeddingRequest, TwelveLabsOutputDataConfig, TwelveLabsS3Location, TwelveLabsS3OutputDataConfig, ) -from litellm.types.utils import Embedding, EmbeddingResponse, Usage +from litellm.types.utils import Embedding, EmbeddingResponse, PromptTokensDetailsWrapper, Usage + + +class MarengoEmbeddingItem(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + embedding: tuple[float, ...] | None = None + + +class MarengoInvokeResponse(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + data: tuple[MarengoEmbeddingItem, ...] = () + embedding: tuple[float, ...] | None = None + embeddings: tuple[MarengoEmbeddingItem, ...] = () + + def vectors(self) -> tuple[tuple[float, ...], ...]: + if self.data: + return tuple(item.embedding for item in self.data if item.embedding is not None) + if self.embedding is not None: + return (self.embedding,) + return tuple(item.embedding for item in self.embeddings if item.embedding is not None) + + +class MarengoBilledMultiInput(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + inputText: str | None = None + mediaSources: tuple[Mapping[str, object], ...] = () + + +class MarengoBilledRequest(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + inputType: TWELVELABS_MARENGO_3_INPUT_TYPES | None = None + multi_input: MarengoBilledMultiInput | None = None + + +INVOKE_RESPONSES: Final = TypeAdapter(tuple[MarengoInvokeResponse, ...]) +BILLED_REQUESTS: Final = TypeAdapter(tuple[MarengoBilledRequest, ...]) + + +def _billed_units(request: MarengoBilledRequest) -> tuple[int, int]: + input_type: Final = request.inputType + match input_type: + case "text": + return (1, 0) + case "image": + return (0, 1) + case "text_image": + return (1, 1) + case "multi_input": + multi_input: Final = request.multi_input or MarengoBilledMultiInput() + return (1 if multi_input.inputText else 0, len(multi_input.mediaSources)) + case "video" | "audio" | None: + return (0, 0) + case _: + assert_never(input_type) + + +def _billed_usage(batch_data: list[dict] | None) -> Usage: + units: Final = tuple(_billed_units(request) for request in BILLED_REQUESTS.validate_python(batch_data or ())) + query_count: Final = sum(text_requests for text_requests, _ in units) + image_count: Final = sum(images for _, images in units) + details: Final = ( + PromptTokensDetailsWrapper(query_count=query_count or None, image_count=image_count or None) + if query_count or image_count + else None + ) + return Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0, prompt_tokens_details=details) + + +MARENGO_SHARED_PARAMS: Final = ( + "encoding_format", + "embeddingOption", + "startSec", + "input_type", + "endSec", + "segmentation", + "embeddingType", + "embeddingScope", + "inferenceId", + "media_source", + "media_sources", +) + + +def drop_params_enabled(litellm_params: Mapping[str, object]) -> bool: + return litellm.drop_params is True or litellm_params.get("drop_params") is True class TwelveLabsMarengoEmbeddingConfig: @@ -26,28 +127,24 @@ class TwelveLabsMarengoEmbeddingConfig: Supports text, image, video, and audio inputs. - InvokeModel: text and image inputs - StartAsyncInvoke: video, audio, image, and text inputs + + Marengo 3.0 (model ids containing "marengo-embed-3") nests the input under a key named after inputType and + adds the text_image and multi_input input types; that payload is built by build_marengo_3_request. """ - def __init__(self) -> None: - pass + def __init__(self, model: str | None = None) -> None: + self.is_marengo_3: Final = is_marengo_3_model(model) def get_supported_openai_params(self) -> list[str]: - return [ - "encoding_format", - "textTruncate", - "embeddingOption", - "startSec", - "lengthSec", - "useFixedLengthSec", - "minClipSec", - "input_type", - ] + if self.is_marengo_3: + return list(MARENGO_SHARED_PARAMS) + return [*MARENGO_SHARED_PARAMS, *MARENGO_2_7_ONLY_PARAMS] 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": + if v == "float" and not self.is_marengo_3: optional_params["embeddingOption"] = ["visual-text", "visual-image"] elif k == "textTruncate": optional_params["textTruncate"] = v @@ -56,7 +153,19 @@ class TwelveLabsMarengoEmbeddingConfig: elif k == "input_type": # Map input_type to inputType for Bedrock optional_params["inputType"] = v - elif k in ["startSec", "lengthSec", "useFixedLengthSec", "minClipSec"]: + elif k in ( + "startSec", + "lengthSec", + "useFixedLengthSec", + "minClipSec", + "endSec", + "segmentation", + "embeddingType", + "embeddingScope", + "inferenceId", + "media_source", + "media_sources", + ): optional_params[k] = v return optional_params @@ -77,7 +186,8 @@ class TwelveLabsMarengoEmbeddingConfig: async_invoke_route: bool = False, model_id: str | None = None, output_s3_uri: str | None = None, - ) -> TwelveLabsMarengoEmbeddingRequest | TwelveLabsAsyncInvokeRequest: + drop_params: bool = False, + ) -> TwelveLabsMarengoEmbeddingRequest | TwelveLabsMarengo3EmbeddingRequest | TwelveLabsAsyncInvokeRequest: """ Transform OpenAI-style input to TwelveLabs Marengo format/async-invoke format. @@ -87,20 +197,29 @@ class TwelveLabsMarengoEmbeddingConfig: - Video inputs (async-invoke only) - Audio inputs (async-invoke only) - S3 URLs for all media types (async-invoke only) + - Marengo 3.0 only: text_image and multi_input inputs (nested payload) """ - # Get input_type or default to "text" input_type: Final = cast( TWELVELABS_EMBEDDING_INPUT_TYPES, inference_params.get("inputType") or inference_params.get("input_type") or "text", ) - # Validate that async-invoke is used for video/audio if input_type in ["video", "audio"] and not async_invoke_route: raise ValueError( f"Input type '{input_type}' requires async_invoke route. " f"Use model format: 'bedrock/async_invoke/model_id'" ) + if self.is_marengo_3: + marengo_3_request: Final = build_marengo_3_request( + input=input, inference_params=inference_params, drop_params=drop_params + ) + if async_invoke_route and model_id: + return self._wrap_async_invoke_request( + model_input=marengo_3_request, model_id=model_id, output_s3_uri=output_s3_uri + ) + return marengo_3_request + transformed_request: Final[TwelveLabsMarengoEmbeddingRequest] = {"inputType": input_type} if input_type == "text": @@ -154,7 +273,7 @@ class TwelveLabsMarengoEmbeddingConfig: def _wrap_async_invoke_request( self, - model_input: TwelveLabsMarengoEmbeddingRequest, + model_input: TwelveLabsMarengoEmbeddingRequest | TwelveLabsMarengo3EmbeddingRequest, model_id: str, output_s3_uri: str | None = None, ) -> TwelveLabsAsyncInvokeRequest: @@ -188,62 +307,16 @@ class TwelveLabsMarengoEmbeddingConfig: ), ) - def _transform_response(self, response_list: list[dict], model: str) -> EmbeddingResponse: - """ - Transform TwelveLabs response to OpenAI format. - Handles the actual TwelveLabs response format: {"data": [{"embedding": [...]}]} - """ - embeddings: Final[list[Embedding]] = [] - total_tokens = 0 - - for response in response_list: - # TwelveLabs response format has a "data" field containing the embeddings - if "data" in response and isinstance(response["data"], list): - for item in response["data"]: - if "embedding" in item: - # Single embedding response - embedding = Embedding( - embedding=item["embedding"], - index=len(embeddings), - object="embedding", - ) - embeddings.append(embedding) - - # Estimate token count (rough approximation) - if "inputTextTokenCount" in item: - total_tokens += item["inputTextTokenCount"] - else: - # Rough estimate: 1 token per 4 characters for text, or use embedding size - total_tokens += len(item["embedding"]) // 4 - elif "embedding" in response: - # Direct embedding response (fallback for other formats) - 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: Final = Usage(prompt_tokens=total_tokens, total_tokens=total_tokens) - - return EmbeddingResponse(data=embeddings, model=model, usage=usage) + def _transform_response( + self, response_list: list[dict], model: str, batch_data: list[dict] | None = None + ) -> EmbeddingResponse: + vectors: Final = tuple( + vector for response in INVOKE_RESPONSES.validate_python(response_list) for vector in response.vectors() + ) + embeddings: Final = [ + Embedding(embedding=list(vector), index=index, object="embedding") for index, vector in enumerate(vectors) + ] + return EmbeddingResponse(data=embeddings, model=model, usage=_billed_usage(batch_data)) def _transform_async_invoke_response(self, response: dict, model: str) -> EmbeddingResponse: """ diff --git a/litellm/main.py b/litellm/main.py index 75b7f7f10a5..56f9cb2c0d0 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6545,7 +6545,7 @@ def embedding( client=client, timeout=timeout, aembedding=aembedding, - litellm_params={}, + litellm_params=litellm_params_dict, api_base=api_base, print_verbose=print_verbose, extra_headers=headers, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b1ffc1583e4..7784ed2a6ac 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -650,7 +650,10 @@ }, "twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, "litellm_provider": "bedrock", "max_input_tokens": 77, "max_tokens": 77, @@ -662,7 +665,7 @@ }, "us.twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "input_cost_per_image": 0.0001, @@ -677,7 +680,7 @@ }, "eu.twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "input_cost_per_image": 0.0001, @@ -690,6 +693,48 @@ "supports_embedding_image_input": true, "supports_image_input": true }, + "twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, + "us.twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, + "eu.twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, "twelvelabs.pegasus-1-2-v1:0": { "input_cost_per_video_per_second": 0.00049, "output_cost_per_token": 7.5e-06, diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index bed0ba3dc08..9f93886a9c6 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1,7 +1,7 @@ import json from collections.abc import Sequence from enum import Enum -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias from typing_extensions import ReadOnly, Required, TypedDict, override @@ -557,7 +557,7 @@ class AmazonTitanMultimodalEmbeddingResponse(TypedDict): message: str # Specifies any errors that occur during generation. -# TwelveLabs Marengo Embed 2.7 types +# TwelveLabs Marengo Embed types TWELVELABS_EMBEDDING_INPUT_TYPES = Literal["text", "image", "video", "audio"] TWELVELABS_EMBEDDING_OPTIONS = Literal["visual-text", "visual-image", "audio"] @@ -591,6 +591,113 @@ class TwelveLabsMarengoEmbeddingResponse(TypedDict): endSec: float +TWELVELABS_MARENGO_3_INPUT_TYPES: TypeAlias = Literal["text", "image", "video", "audio", "text_image", "multi_input"] +TWELVELABS_MARENGO_3_EMBEDDING_OPTIONS: TypeAlias = Literal["visual", "audio", "transcription"] +TWELVELABS_MARENGO_3_EMBEDDING_TYPES: TypeAlias = Literal["separate_embedding", "fused_embedding"] +TWELVELABS_MARENGO_3_EMBEDDING_SCOPES: TypeAlias = Literal["clip", "asset"] + + +class TwelveLabsMarengo3FixedSegmentationConfig(TypedDict): + durationSec: ReadOnly[int] + + +class TwelveLabsMarengo3FixedSegmentation(TypedDict): + method: ReadOnly[Literal["fixed"]] + fixed: ReadOnly[TwelveLabsMarengo3FixedSegmentationConfig] + + +class TwelveLabsMarengo3DynamicSegmentationConfig(TypedDict): + minDurationSec: ReadOnly[int] + + +class TwelveLabsMarengo3DynamicSegmentation(TypedDict): + method: ReadOnly[Literal["dynamic"]] + dynamic: ReadOnly[TwelveLabsMarengo3DynamicSegmentationConfig] + + +TwelveLabsMarengo3Segmentation: TypeAlias = TwelveLabsMarengo3FixedSegmentation | TwelveLabsMarengo3DynamicSegmentation + + +class TwelveLabsMarengo3TextInput(TypedDict): + inputText: ReadOnly[str] + + +class TwelveLabsMarengo3ImageInput(TypedDict): + mediaSource: ReadOnly[TwelveLabsMediaSource] + + +class TwelveLabsMarengo3TimedMediaOptions(TypedDict, total=False): + startSec: ReadOnly[float] + endSec: ReadOnly[float] + segmentation: ReadOnly[TwelveLabsMarengo3Segmentation] + embeddingOption: ReadOnly[Sequence[TWELVELABS_MARENGO_3_EMBEDDING_OPTIONS]] + embeddingType: ReadOnly[Sequence[TWELVELABS_MARENGO_3_EMBEDDING_TYPES]] + embeddingScope: ReadOnly[Sequence[TWELVELABS_MARENGO_3_EMBEDDING_SCOPES]] + + +class TwelveLabsMarengo3TimedMediaInput(TwelveLabsMarengo3TimedMediaOptions): + mediaSource: Required[ReadOnly[TwelveLabsMediaSource]] + + +class TwelveLabsMarengo3TextImageInput(TypedDict): + inputText: ReadOnly[str] + mediaSource: ReadOnly[TwelveLabsMediaSource] + + +class TwelveLabsMarengo3NamedMediaSource(TwelveLabsMediaSource): + name: Required[ReadOnly[str]] + mediaType: Required[ReadOnly[Literal["image"]]] + + +class TwelveLabsMarengo3MultiInput(TypedDict, total=False): + inputText: ReadOnly[str] + mediaSources: Required[ReadOnly[Sequence[TwelveLabsMarengo3NamedMediaSource]]] + + +class TwelveLabsMarengo3RequestBase(TypedDict, total=False): + inferenceId: ReadOnly[str] + + +class TwelveLabsMarengo3TextRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["text"]] + text: ReadOnly[TwelveLabsMarengo3TextInput] + + +class TwelveLabsMarengo3ImageRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["image"]] + image: ReadOnly[TwelveLabsMarengo3ImageInput] + + +class TwelveLabsMarengo3VideoRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["video"]] + video: ReadOnly[TwelveLabsMarengo3TimedMediaInput] + + +class TwelveLabsMarengo3AudioRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["audio"]] + audio: ReadOnly[TwelveLabsMarengo3TimedMediaInput] + + +class TwelveLabsMarengo3TextImageRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["text_image"]] + text_image: ReadOnly[TwelveLabsMarengo3TextImageInput] + + +class TwelveLabsMarengo3MultiInputRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["multi_input"]] + multi_input: ReadOnly[TwelveLabsMarengo3MultiInput] + + +TwelveLabsMarengo3EmbeddingRequest: TypeAlias = ( + TwelveLabsMarengo3TextRequest + | TwelveLabsMarengo3ImageRequest + | TwelveLabsMarengo3VideoRequest + | TwelveLabsMarengo3AudioRequest + | TwelveLabsMarengo3TextImageRequest + | TwelveLabsMarengo3MultiInputRequest +) + + class TwelveLabsS3OutputDataConfig(TypedDict): s3Uri: str @@ -601,7 +708,7 @@ class TwelveLabsOutputDataConfig(TypedDict): class TwelveLabsAsyncInvokeRequest(TypedDict): modelId: str - modelInput: TwelveLabsMarengoEmbeddingRequest + modelInput: ReadOnly[TwelveLabsMarengoEmbeddingRequest | TwelveLabsMarengo3EmbeddingRequest] outputDataConfig: TwelveLabsOutputDataConfig diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 40a6d77482f..d62f00f3676 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -272,7 +272,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_token_above_272k_tokens_flex: float | None input_cost_per_token_above_512k_tokens: float | None # MiniMax-M3: prompts >512K priced at 2x input input_cost_per_character_above_128k_tokens: float | None # only for vertex ai models - input_cost_per_query: float | None # only for rerank models + input_cost_per_query: float | None # per-request pricing: rerank, search, and Bedrock Marengo embeddings input_cost_per_image: float | None # only for vertex ai models input_cost_per_image_token: float | None # for gpt-image-1 and similar models input_cost_per_video_token: float | None # for gemini omni models with video input @@ -1694,6 +1694,9 @@ class PromptTokensDetailsWrapper( audio_length_seconds: float | None = None """Length of audio sent to the model. Used for multimodal embeddings priced per audio-second.""" + query_count: int | None = None + """Number of billable requests sent to the model. Used for embeddings priced per request, such as Bedrock Marengo.""" + cache_write_tokens: int | None = None """Number of cache write (creation) tokens sent to the model. OpenAI naming (prompt_tokens_details.cache_write_tokens); this is the canonical field.""" @@ -1735,6 +1738,8 @@ class PromptTokensDetailsWrapper( del self.video_length_seconds if self.audio_length_seconds is None: del self.audio_length_seconds + if self.query_count is None: + del self.query_count if self.web_search_requests is None: del self.web_search_requests if self.google_maps_grounding_requests is None: diff --git a/litellm/utils.py b/litellm/utils.py index 33ef57e1837..36b48d3b8d8 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3632,7 +3632,7 @@ def get_optional_params_embeddings( elif "cohere.embed" in model: object = litellm.BedrockCohereEmbeddingConfig() elif "twelvelabs" in model or "marengo" in model: - object = litellm.TwelveLabsMarengoEmbeddingConfig() + object = litellm.TwelveLabsMarengoEmbeddingConfig(model=model) elif "nova" in model.lower(): object = litellm.AmazonNovaEmbeddingConfig() else: # unmapped model @@ -6043,7 +6043,7 @@ def get_model_info( input_cost_per_character_above_128k_tokens: Optional[ float ] # only for vertex ai models - input_cost_per_query: Optional[float] # only for rerank models + input_cost_per_query: Optional[float] # per-request pricing: rerank, search, and Bedrock Marengo embeddings input_cost_per_image: Optional[float] # only for vertex ai models input_cost_per_audio_token: Optional[float] input_cost_per_audio_per_second: Optional[float] # only for vertex ai models diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b1ffc1583e4..7784ed2a6ac 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -650,7 +650,10 @@ }, "twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, "litellm_provider": "bedrock", "max_input_tokens": 77, "max_tokens": 77, @@ -662,7 +665,7 @@ }, "us.twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "input_cost_per_image": 0.0001, @@ -677,7 +680,7 @@ }, "eu.twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "input_cost_per_image": 0.0001, @@ -690,6 +693,48 @@ "supports_embedding_image_input": true, "supports_image_input": true }, + "twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, + "us.twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, + "eu.twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, "twelvelabs.pegasus-1-2-v1:0": { "input_cost_per_video_per_second": 0.00049, "output_cost_per_token": 7.5e-06, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 59f0938e338..65a6dd2a4ca 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2658,6 +2658,7 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): "image_count": 0, "video_length_seconds": 0.0, "audio_length_seconds": 0.0, + "query_count": 0, } model_info: ModelInfo = {} @@ -3239,6 +3240,37 @@ def test_image_count_prevents_text_tokens_fallback(_local_model_cost_map): assert completion_cost == 0.0 +def test_query_count_bills_input_cost_per_query(_local_model_cost_map): + usage = Usage( + prompt_tokens=0, + completion_tokens=0, + total_tokens=0, + prompt_tokens_details=PromptTokensDetailsWrapper(query_count=3, image_count=1), + ) + + prompt_cost, completion_cost = generic_cost_per_token( + model="us.twelvelabs.marengo-embed-3-0-v1:0", + usage=usage, + custom_llm_provider="bedrock", + ) + + assert prompt_cost == pytest.approx(3 * 7e-05 + 1e-04) + assert completion_cost == 0.0 + + +def test_query_count_is_free_without_a_per_query_price(_local_model_cost_map): + usage = Usage( + prompt_tokens=0, + completion_tokens=0, + total_tokens=0, + prompt_tokens_details=PromptTokensDetailsWrapper(query_count=1), + ) + + prompt_cost, _ = generic_cost_per_token(model="text-embedding-3-small", usage=usage, custom_llm_provider="openai") + + assert prompt_cost == 0.0 + + # --------------------------------------------------------------------------- # Data-residency (OpenAI regional processing) tests # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py index 74a55cc1ef2..ddbd3a2e9ba 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py @@ -184,6 +184,45 @@ class TestBedrockAsyncInvokeEmbedding: request_url = mock_post.call_args.kwargs.get("url", "") assert "/async-invoke" in request_url + def test_async_invoke_marengo_3_wraps_the_nested_payload_with_the_base_model_id(self): + client = HTTPHandler() + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(async_invoke_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model="bedrock/async_invoke/twelvelabs.marengo-embed-3-0-v1:0", + input="s3://test-bucket/clip.mp4", + client=client, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + api_key="test-bearer-token-12345", + input_type="video", + embeddingOption=["visual", "audio"], + segmentation={"method": "fixed", "fixed": {"durationSec": 6}}, + bucketOwner="123456789012", + output_s3_uri="s3://test-bucket/async-invoke-output/", + ) + + assert response._hidden_params._invocation_arn == async_invoke_response["invocationArn"] + assert mock_post.call_args.kwargs["url"].endswith("/async-invoke") + assert json.loads(mock_post.call_args.kwargs["data"]) == { + "modelId": "twelvelabs.marengo-embed-3-0-v1:0", + "modelInput": { + "inputType": "video", + "video": { + "mediaSource": {"s3Location": {"uri": "s3://test-bucket/clip.mp4", "bucketOwner": "123456789012"}}, + "segmentation": {"method": "fixed", "fixed": {"durationSec": 6}}, + "embeddingOption": ["visual", "audio"], + }, + }, + "outputDataConfig": {"s3OutputDataConfig": {"s3Uri": "s3://test-bucket/async-invoke-output/"}}, + } + @pytest.mark.asyncio async def test_async_invoke_twelvelabs_embedding_async_with_mock(self): """Test async invoke embedding with async calls.""" diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index 50f8bbcf584..bcd1a29d0e8 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -5,6 +5,7 @@ from unittest.mock import Mock, patch import pytest import litellm +from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler # Mock responses for different embedding models @@ -1059,3 +1060,182 @@ def test_bedrock_embedding_bearer_token_never_runs_the_sigv4_credential_chain(mo assert response.data[0]["embedding"] == titan_embedding_response["embedding"] assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" + + +marengo_3_embedding_response = {"data": [{"embedding": [0.01 * i for i in range(512)]}]} +MARENGO_3_DUCK = "data:image/png;base64,ZHVjaw==" + + +@pytest.mark.parametrize( + "model,kwargs,expected_body,expected_usage_details", + [ + ( + "bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", + {"input_type": "text"}, + {"inputType": "text", "text": {"inputText": "a duck on water"}}, + {"query_count": 1}, + ), + ( + "bedrock/twelvelabs.marengo-embed-3-0-v1:0", + {"input_type": "text"}, + {"inputType": "text", "text": {"inputText": "a duck on water"}}, + {"query_count": 1}, + ), + ( + "bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", + {"input_type": "text_image", "media_source": MARENGO_3_DUCK}, + { + "inputType": "text_image", + "text_image": {"inputText": "a duck on water", "mediaSource": {"base64String": "ZHVjaw=="}}, + }, + {"query_count": 1, "image_count": 1}, + ), + ( + "bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", + {"input_type": "multi_input", "media_sources": {"bird": MARENGO_3_DUCK}}, + { + "inputType": "multi_input", + "multi_input": { + "inputText": "a duck on water", + "mediaSources": [{"name": "bird", "mediaType": "image", "base64String": "ZHVjaw=="}], + }, + }, + {"query_count": 1, "image_count": 1}, + ), + ], +) +def test_marengo_3_embedding_sends_the_nested_payload_and_parses_512_dims( + model, kwargs, expected_body, expected_usage_details +): + client = HTTPHandler() + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(marengo_3_embedding_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model=model, + input="a duck on water", + client=client, + aws_region_name="us-east-1", + api_key="test-bearer-token-12345", + **kwargs, + ) + + assert json.loads(mock_post.call_args.kwargs["data"]) == expected_body + assert mock_post.call_args.kwargs["url"].endswith(f"/model/{model.removeprefix('bedrock/').replace(':', '%3A')}/invoke") + assert len(response.data[0]["embedding"]) == 512 + assert response.data[0]["embedding"][:2] == [0.0, 0.01] + assert response.usage.prompt_tokens == 0 + assert response.usage.total_tokens == 0 + assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == expected_usage_details + + +def test_marengo_3_image_embedding_sends_the_media_under_the_image_key(): + client = HTTPHandler() + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(marengo_3_embedding_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model="bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", + input=MARENGO_3_DUCK, + client=client, + aws_region_name="us-east-1", + api_key="test-bearer-token-12345", + input_type="image", + ) + + assert json.loads(mock_post.call_args.kwargs["data"]) == { + "inputType": "image", + "image": {"mediaSource": {"base64String": "ZHVjaw=="}}, + } + assert len(response.data[0]["embedding"]) == 512 + assert response.data[0]["embedding"][:2] == [0.0, 0.01] + assert response.usage.prompt_tokens == 0 + assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == {"image_count": 1} + + +def test_marengo_2_7_embedding_keeps_the_flat_payload(): + client = HTTPHandler() + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(twelvelabs_embedding_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model="bedrock/us.twelvelabs.marengo-embed-2-7-v1:0", + input="a duck on water", + client=client, + aws_region_name="us-east-1", + api_key="test-bearer-token-12345", + input_type="text", + ) + + assert json.loads(mock_post.call_args.kwargs["data"]) == { + "inputType": "text", + "inputText": "a duck on water", + "textTruncate": "end", + } + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + assert response.usage.prompt_tokens == 0 + assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == {"query_count": 1} + + +def test_marengo_usage_counts_text_requests_and_images_across_a_batch(): + duck = {"mediaType": "image", "base64String": "ZHVjaw=="} + response = TwelveLabsMarengoEmbeddingConfig()._transform_response( + response_list=[marengo_3_embedding_response, marengo_3_embedding_response, marengo_3_embedding_response], + model="us.twelvelabs.marengo-embed-3-0-v1:0", + batch_data=[ + {"inputType": "text", "text": {"inputText": "a duck"}}, + {"inputType": "image", "image": {"mediaSource": {"base64String": "ZHVjaw=="}}}, + {"inputType": "multi_input", "multi_input": {"mediaSources": [{"name": "a", **duck}, {"name": "b", **duck}]}}, + ], + ) + + assert [item["index"] for item in response.data] == [0, 1, 2] + assert response.usage.prompt_tokens == 0 + assert response.usage.total_tokens == 0 + assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == {"query_count": 1, "image_count": 3} + + +def test_marengo_usage_without_request_data_bills_nothing(): + response = TwelveLabsMarengoEmbeddingConfig()._transform_response( + response_list=[marengo_3_embedding_response], model="us.twelvelabs.marengo-embed-3-0-v1:0" + ) + + assert len(response.data[0]["embedding"]) == 512 + assert response.usage.prompt_tokens == 0 + assert response.usage.prompt_tokens_details is None + + +def test_marengo_response_items_without_an_embedding_are_skipped(): + response = TwelveLabsMarengoEmbeddingConfig()._transform_response( + response_list=[{"data": [{"embeddingOption": "visual-text", "startSec": 0.0}, {"embedding": [0.1, 0.2, 0.3]}]}], + model="us.twelvelabs.marengo-embed-3-0-v1:0", + ) + + assert [item["embedding"] for item in response.data] == [[0.1, 0.2, 0.3]] + assert response.data[0]["index"] == 0 + + +def test_marengo_3_text_image_without_media_source_is_a_bad_request(): + with pytest.raises(litellm.BadRequestError, match=r"text_image.*media_source"): + litellm.embedding( + model="bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", + input="a duck on water", + aws_region_name="us-east-1", + api_key="test-bearer-token-12345", + input_type="text_image", + ) diff --git a/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py b/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py new file mode 100644 index 00000000000..f149953b6f1 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py @@ -0,0 +1,416 @@ +import json +from unittest.mock import Mock, patch + +import pytest + +import litellm +from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.bedrock.embed.twelvelabs_marengo_3_transformation import ( + MARENGO_2_7_ONLY_PARAMS, + build_marengo_3_request, + is_marengo_3_model, +) +from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import ( + TwelveLabsMarengoEmbeddingConfig, + drop_params_enabled, +) + +MARENGO_3_BASE = "twelvelabs.marengo-embed-3-0-v1:0" +MARENGO_3_US = "us.twelvelabs.marengo-embed-3-0-v1:0" +MARENGO_27_US = "us.twelvelabs.marengo-embed-2-7-v1:0" +DUCK_DATA_URL = "data:image/png;base64,ZHVjaw==" +OUTPUT_S3_URI = "s3://out-bucket/marengo/" + + +@pytest.mark.parametrize( + "model,expected", + [ + (MARENGO_3_BASE, True), + (MARENGO_3_US, True), + ("eu.twelvelabs.marengo-embed-3-0-v1:0", True), + ("async_invoke/twelvelabs.marengo-embed-3-0-v1:0", True), + (MARENGO_27_US, False), + ("twelvelabs.marengo-embed-2-7-v1:0", False), + ("twelvelabs.marengo-embed-30-v1:0", False), + (None, False), + ], +) +def test_is_marengo_3_model(model, expected): + assert is_marengo_3_model(model) is expected + + +def wire(request: object) -> object: + return json.loads(json.dumps(request)) + + +def test_text_request_nests_input_text_under_text(): + assert build_marengo_3_request("a dog on the beach", {"input_type": "text"}) == { + "inputType": "text", + "text": {"inputText": "a dog on the beach"}, + } + + +def test_missing_input_type_defaults_to_text(): + assert build_marengo_3_request("hello", {})["inputType"] == "text" + + +def test_camel_case_input_type_wins_over_snake_case(): + request = build_marengo_3_request(DUCK_DATA_URL, {"inputType": "image", "input_type": "text"}) + assert request["inputType"] == "image" + + +def test_image_request_strips_data_url_prefix(): + assert build_marengo_3_request(DUCK_DATA_URL, {"input_type": "image"}) == { + "inputType": "image", + "image": {"mediaSource": {"base64String": "ZHVjaw=="}}, + } + + +def test_image_request_from_s3_carries_bucket_owner(): + request = build_marengo_3_request("s3://media/duck.png", {"input_type": "image", "bucketOwner": "123456789012"}) + assert request == { + "inputType": "image", + "image": {"mediaSource": {"s3Location": {"uri": "s3://media/duck.png", "bucketOwner": "123456789012"}}}, + } + + +@pytest.mark.parametrize( + "input_media,params", + [ + ("s3://media/duck.png", {"input_type": "image"}), + ("s3://media/clip.mp4", {"input_type": "video"}), + ("a duck", {"input_type": "text_image", "media_source": "s3://media/duck.png"}), + ("a duck", {"input_type": "multi_input", "media_sources": {"img1": "s3://media/duck.png"}}), + ], +) +def test_s3_media_without_bucket_owner_is_rejected_naming_it(input_media, params): + with pytest.raises(BedrockError) as excinfo: + build_marengo_3_request(input_media, params) + assert excinfo.value.status_code == 400 + assert excinfo.value.message == ( + "s3:// media requires the 'bucketOwner' parameter, the account id that owns the bucket" + ) + + +def test_text_image_request_pairs_text_with_media_source(): + request = build_marengo_3_request( + "a duck", {"input_type": "text_image", "media_source": DUCK_DATA_URL, "output_s3_uri": OUTPUT_S3_URI} + ) + assert request == { + "inputType": "text_image", + "text_image": {"inputText": "a duck", "mediaSource": {"base64String": "ZHVjaw=="}}, + } + + +def test_text_image_request_requires_media_source(): + with pytest.raises(BedrockError, match=r"text_image.*media_source") as excinfo: + build_marengo_3_request("a duck", {"input_type": "text_image"}) + assert excinfo.value.status_code == 400 + + +def test_multi_input_request_names_each_media_source(): + request = build_marengo_3_request( + "a photo of <@bird> next to <@dog>", + { + "input_type": "multi_input", + "media_sources": {"bird": DUCK_DATA_URL, "dog": "s3://media/dog.png"}, + "bucketOwner": "123456789012", + }, + ) + assert wire(request) == { + "inputType": "multi_input", + "multi_input": { + "inputText": "a photo of <@bird> next to <@dog>", + "mediaSources": [ + {"name": "bird", "mediaType": "image", "base64String": "ZHVjaw=="}, + { + "name": "dog", + "mediaType": "image", + "s3Location": {"uri": "s3://media/dog.png", "bucketOwner": "123456789012"}, + }, + ], + }, + } + + +def test_multi_input_without_text_omits_input_text(): + request = build_marengo_3_request("", {"input_type": "multi_input", "media_sources": {"bird": DUCK_DATA_URL}}) + assert "inputText" not in request["multi_input"] + assert request["multi_input"]["mediaSources"][0]["name"] == "bird" + + +@pytest.mark.parametrize("params", [{"input_type": "multi_input"}, {"input_type": "multi_input", "media_sources": {}}]) +def test_multi_input_request_requires_media_sources(params): + with pytest.raises(BedrockError, match=r"multi_input.*media_sources") as excinfo: + build_marengo_3_request("<@bird>", params) + assert excinfo.value.status_code == 400 + + +@pytest.mark.parametrize("input_type", ["video", "audio"]) +def test_timed_media_request_nests_every_option_under_the_media_key(input_type): + request = build_marengo_3_request( + "s3://media/clip.mp4", + { + "input_type": input_type, + "startSec": 2, + "endSec": 12.5, + "segmentation": {"method": "dynamic", "dynamic": {"minDurationSec": 4}}, + "embeddingOption": ["visual", "audio"], + "embeddingType": ["fused_embedding"], + "embeddingScope": ["clip", "asset"], + "inferenceId": "req-42", + "bucketOwner": "123456789012", + }, + ) + assert wire(request) == { + "inputType": input_type, + input_type: { + "mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4", "bucketOwner": "123456789012"}}, + "startSec": 2.0, + "endSec": 12.5, + "segmentation": {"method": "dynamic", "dynamic": {"minDurationSec": 4}}, + "embeddingOption": ["visual", "audio"], + "embeddingType": ["fused_embedding"], + "embeddingScope": ["clip", "asset"], + }, + "inferenceId": "req-42", + } + + +def test_timed_media_request_without_options_carries_only_the_media_source(): + request = build_marengo_3_request("s3://media/clip.mp4", {"input_type": "video", "bucketOwner": "123456789012"}) + assert request["video"] == { + "mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4", "bucketOwner": "123456789012"}} + } + + +@pytest.mark.parametrize( + "params", + [ + {"input_type": "clip"}, + {"input_type": "video", "embeddingOption": ["visual-text"]}, + {"input_type": "video", "segmentation": {"method": "fixed", "dynamic": {"minDurationSec": 4}}}, + {"input_type": "multi_input", "media_sources": ["not", "a", "mapping"]}, + ], +) +def test_invalid_marengo_3_params_are_rejected_before_the_request_is_sent(params): + with pytest.raises(BedrockError, match=r"Invalid Marengo 3\.0 parameters") as excinfo: + build_marengo_3_request("s3://media/clip.mp4", params) + assert excinfo.value.status_code == 400 + + +def test_config_sends_the_nested_payload_for_marengo_3_and_the_flat_one_for_2_7(): + nested = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US)._transform_request( + input="hello", inference_params={"input_type": "text"} + ) + flat = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_27_US)._transform_request( + input="hello", inference_params={"input_type": "text"} + ) + assert nested == {"inputType": "text", "text": {"inputText": "hello"}} + assert flat == {"inputType": "text", "inputText": "hello", "textTruncate": "end"} + + +def test_config_without_a_model_keeps_the_2_7_payload(): + request = TwelveLabsMarengoEmbeddingConfig()._transform_request(input="hello", inference_params={}) + assert request == {"inputType": "text", "inputText": "hello", "textTruncate": "end"} + + +@pytest.mark.parametrize("input_type", ["video", "audio"]) +def test_marengo_3_video_and_audio_still_require_the_async_route(input_type): + with pytest.raises(ValueError, match=f"Input type '{input_type}' requires async_invoke route"): + TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_BASE)._transform_request( + input="s3://media/clip.mp4", inference_params={"input_type": input_type} + ) + + +def test_marengo_3_async_invoke_wraps_the_nested_payload_with_the_base_model_id(): + request = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_BASE)._transform_request( + input="s3://media/clip.mp4", + inference_params={ + "input_type": "video", + "embeddingOption": ["visual"], + "bucketOwner": "123456789012", + "output_s3_uri": OUTPUT_S3_URI, + }, + async_invoke_route=True, + model_id="async_invoke%2Ftwelvelabs.marengo-embed-3-0-v1%3A0", + output_s3_uri=OUTPUT_S3_URI, + ) + assert wire(request) == { + "modelId": MARENGO_3_BASE, + "modelInput": { + "inputType": "video", + "video": { + "mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4", "bucketOwner": "123456789012"}}, + "embeddingOption": ["visual"], + }, + }, + "outputDataConfig": {"s3OutputDataConfig": {"s3Uri": OUTPUT_S3_URI}}, + } + + +def test_marengo_3_async_invoke_requires_an_output_s3_uri(): + with pytest.raises(ValueError, match="output_s3_uri cannot be empty"): + TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_BASE)._transform_request( + input="hello", + inference_params={"input_type": "text"}, + async_invoke_route=True, + model_id=MARENGO_3_BASE, + output_s3_uri="", + ) + + +def test_encoding_format_float_no_longer_injects_2_7_embedding_options_for_marengo_3(): + marengo_3 = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US).map_openai_params( + non_default_params={"encoding_format": "float"}, optional_params={} + ) + marengo_27 = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_27_US).map_openai_params( + non_default_params={"encoding_format": "float"}, optional_params={} + ) + assert marengo_3 == {} + assert marengo_27 == {"embeddingOption": ["visual-text", "visual-image"]} + + +def test_marengo_3_only_params_are_forwarded_by_map_openai_params(): + mapped = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US).map_openai_params( + non_default_params={ + "input_type": "text_image", + "media_source": DUCK_DATA_URL, + "media_sources": {"bird": DUCK_DATA_URL}, + "endSec": 5, + "segmentation": {"method": "fixed", "fixed": {"durationSec": 6}}, + "embeddingType": ["separate_embedding"], + "embeddingScope": ["clip"], + "inferenceId": "req-1", + }, + optional_params={}, + ) + assert mapped == { + "inputType": "text_image", + "media_source": DUCK_DATA_URL, + "media_sources": {"bird": DUCK_DATA_URL}, + "endSec": 5, + "segmentation": {"method": "fixed", "fixed": {"durationSec": 6}}, + "embeddingType": ["separate_embedding"], + "embeddingScope": ["clip"], + "inferenceId": "req-1", + } + + +@pytest.mark.parametrize( + "params,problem", + [ + ( + {"input_type": "clip"}, + "input_type: Input should be 'text', 'image', 'video', 'audio', 'text_image' or 'multi_input'", + ), + ({"input_type": "video", "embeddingOption": "visual"}, "embeddingOption: Input should be a valid tuple"), + ( + {"input_type": "multi_input", "media_sources": ["not", "a", "mapping"]}, + "media_sources: Input should be a valid dictionary", + ), + ], +) +def test_invalid_marengo_3_params_name_the_field_and_the_reason(params, problem): + with pytest.raises(BedrockError) as excinfo: + build_marengo_3_request("s3://media/clip.mp4", params) + assert excinfo.value.message == f"Invalid Marengo 3.0 parameters: {problem}" + + +MARENGO_2_7_ONLY_VALUES = {"textTruncate": "end", "lengthSec": 5, "useFixedLengthSec": True, "minClipSec": 2} + + +@pytest.mark.parametrize("name", MARENGO_2_7_ONLY_PARAMS) +def test_marengo_2_7_only_params_are_rejected_on_3_0_unless_dropped(name): + params = {"input_type": "text", name: MARENGO_2_7_ONLY_VALUES[name]} + with pytest.raises(BedrockError) as excinfo: + build_marengo_3_request("hello", params) + assert excinfo.value.status_code == 400 + assert excinfo.value.message == ( + f"Marengo 3.0 does not accept the Marengo 2.7 parameters {name}; set drop_params to drop them" + ) + assert build_marengo_3_request("hello", params, drop_params=True) == { + "inputType": "text", + "text": {"inputText": "hello"}, + } + + +def test_marengo_2_7_only_params_are_advertised_only_for_2_7(): + marengo_3 = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US).get_supported_openai_params() + marengo_27 = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_27_US).get_supported_openai_params() + assert set(MARENGO_2_7_ONLY_PARAMS).isdisjoint(marengo_3) + assert set(MARENGO_2_7_ONLY_PARAMS) <= set(marengo_27) + assert set(marengo_3) <= set(marengo_27) + + +def test_drop_params_comes_from_the_call_or_the_global(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + assert drop_params_enabled({}) is False + assert drop_params_enabled({"drop_params": True}) is True + monkeypatch.setattr(litellm, "drop_params", True) + assert drop_params_enabled({}) is True + + +def test_config_drops_marengo_2_7_only_params_only_when_asked(): + config = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US) + with pytest.raises(BedrockError, match=r"Marengo 2\.7 parameters textTruncate"): + config._transform_request("hello", {"textTruncate": "end"}) + assert config._transform_request("hello", {"textTruncate": "end"}, drop_params=True) == { + "inputType": "text", + "text": {"inputText": "hello"}, + } + + +@pytest.mark.parametrize( + "params", + [ + {"input_type": "text"}, + {"input_type": "image"}, + {"input_type": "text_image", "media_source": DUCK_DATA_URL}, + {"input_type": "multi_input", "media_sources": {"bird": DUCK_DATA_URL}}, + ], +) +def test_timed_media_options_are_rejected_on_untimed_input_types_unless_dropped(params): + timed = {**params, "startSec": 0, "embeddingOption": ["visual"]} + with pytest.raises(BedrockError) as excinfo: + build_marengo_3_request(DUCK_DATA_URL, timed) + assert excinfo.value.status_code == 400 + assert excinfo.value.message == ( + f"Input type '{params['input_type']}' does not accept startSec, embeddingOption; set drop_params to drop them" + ) + assert build_marengo_3_request(DUCK_DATA_URL, timed, drop_params=True) == build_marengo_3_request( + DUCK_DATA_URL, params + ) + + +def _embed_marengo_3_us(client: HTTPHandler, **params: object): + return litellm.embedding( + model=f"bedrock/{MARENGO_3_US}", + input="hello", + client=client, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + api_key="test-bearer-token", + **params, + ) + + +def test_per_request_drop_params_reaches_the_marengo_3_builder(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + client = HTTPHandler() + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps({"data": [{"embedding": [0.1, 0.2]}]}) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + with pytest.raises(litellm.BadRequestError, match=r"Marengo 2\.7 parameters textTruncate"): + _embed_marengo_3_us(client, textTruncate="end") + assert mock_post.call_count == 0 + + response = _embed_marengo_3_us(client, textTruncate="end", drop_params=True) + + assert response.data[0]["embedding"] == [0.1, 0.2] + assert json.loads(mock_post.call_args.kwargs["data"]) == {"inputType": "text", "text": {"inputText": "hello"}} diff --git a/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py b/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py new file mode 100644 index 00000000000..0bb99339435 --- /dev/null +++ b/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py @@ -0,0 +1,117 @@ +import json +from pathlib import Path + +import pytest + +import litellm +from litellm.constants import bedrock_embedding_models +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.types.utils import PromptTokensDetailsWrapper, Usage + +REPO_ROOT = Path(__file__).parents[2] +MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" + +BASE_MODEL = "twelvelabs.marengo-embed-3-0-v1:0" +PROFILE_MODELS = ("us.twelvelabs.marengo-embed-3-0-v1:0", "eu.twelvelabs.marengo-embed-3-0-v1:0") +ALL_MODELS = (BASE_MODEL, *PROFILE_MODELS) +MARENGO_2_7_MODELS = ( + "twelvelabs.marengo-embed-2-7-v1:0", + "us.twelvelabs.marengo-embed-2-7-v1:0", + "eu.twelvelabs.marengo-embed-2-7-v1:0", +) +PER_REQUEST_MODELS = (*ALL_MODELS, *MARENGO_2_7_MODELS) + +TEXT_REQUEST_COST = 7e-05 +IMAGE_REQUEST_COST = 0.0001 +VIDEO_COST_PER_SECOND = 0.0007 +AUDIO_COST_PER_SECOND = 0.00014 + + +def _load(path): + with open(path) as f: + return json.load(f) + + +@pytest.mark.parametrize("model", ALL_MODELS) +def test_marengo_embed_3_specs(model): + info = _load(MAIN_PATH).get(model) + assert info is not None, f"{model} missing from model_prices_and_context_window.json" + + assert info["litellm_provider"] == "bedrock" + assert info["mode"] == "embedding" + assert info["input_cost_per_query"] == TEXT_REQUEST_COST + assert info["output_cost_per_token"] == 0.0 + assert info["max_input_tokens"] == 500 + assert info["max_tokens"] == 500 + assert info["output_vector_size"] == 512 + assert info["supports_embedding_image_input"] is True + assert info["supports_image_input"] is True + assert "deprecation_date" not in info + + routed_model, provider, _, _ = get_llm_provider(model=f"bedrock/{model}") + assert routed_model == model + assert provider == "bedrock" + + +@pytest.mark.parametrize("model", PER_REQUEST_MODELS) +def test_marengo_prices_are_per_request_not_per_token(model): + info = _load(MAIN_PATH)[model] + assert "input_cost_per_token" not in info + assert info["input_cost_per_query"] == TEXT_REQUEST_COST + assert info["input_cost_per_image"] == IMAGE_REQUEST_COST + assert info["input_cost_per_video_per_second"] == VIDEO_COST_PER_SECOND + assert info["input_cost_per_audio_per_second"] == AUDIO_COST_PER_SECOND + + +@pytest.mark.parametrize("model", ALL_MODELS) +def test_marengo_embed_3_is_visible_to_callers(model, local_model_cost_map): + info = litellm.get_model_info(model=model, custom_llm_provider="bedrock") + assert info["mode"] == "embedding" + assert info["output_vector_size"] == 512 + assert info["max_input_tokens"] == 500 + + +@pytest.mark.parametrize("model", PER_REQUEST_MODELS) +@pytest.mark.parametrize( + "details,expected_cost", + [ + (PromptTokensDetailsWrapper(query_count=1), TEXT_REQUEST_COST), + (PromptTokensDetailsWrapper(image_count=1), IMAGE_REQUEST_COST), + (PromptTokensDetailsWrapper(query_count=1, image_count=1), TEXT_REQUEST_COST + IMAGE_REQUEST_COST), + (PromptTokensDetailsWrapper(query_count=1, image_count=2), TEXT_REQUEST_COST + 2 * IMAGE_REQUEST_COST), + (PromptTokensDetailsWrapper(video_length_seconds=10), 10 * VIDEO_COST_PER_SECOND), + (PromptTokensDetailsWrapper(audio_length_seconds=10), 10 * AUDIO_COST_PER_SECOND), + ], +) +def test_marengo_requests_are_billed_per_request(model, details, expected_cost, local_model_cost_map): + usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0, prompt_tokens_details=details) + prompt_cost, completion_cost = litellm.cost_per_token( + model=model, usage_object=usage, custom_llm_provider="bedrock" + ) + assert prompt_cost == pytest.approx(expected_cost) + assert completion_cost == 0.0 + + +@pytest.mark.parametrize("model", PER_REQUEST_MODELS) +def test_marengo_token_counts_bill_nothing(model, local_model_cost_map): + usage = Usage(prompt_tokens=128, completion_tokens=0, total_tokens=128) + prompt_cost, completion_cost = litellm.cost_per_token( + model=model, usage_object=usage, custom_llm_provider="bedrock" + ) + assert prompt_cost == 0.0 + assert completion_cost == 0.0 + + +def test_marengo_embed_3_is_a_known_bedrock_embedding_model(): + assert BASE_MODEL in bedrock_embedding_models + + +@pytest.mark.parametrize("model", PER_REQUEST_MODELS) +def test_backup_matches_main(model): + main_cost = _load(MAIN_PATH) + backup_cost = _load(BACKUP_PATH) + + assert model in main_cost, f"{model} missing from model_prices_and_context_window.json" + assert model in backup_cost, f"{model} missing from model_prices_and_context_window_backup.json" + assert backup_cost[model] == main_cost[model], f"{model} differs between main and backup model cost maps"