From 69d07a3625a80f532e994838281f39b65834ec56 Mon Sep 17 00:00:00 2001 From: hugosmoreira Date: Wed, 12 Aug 2026 13:39:15 -0700 Subject: [PATCH] fix(cohere): support multimodal embed inputs --- litellm/llms/cohere/embed/transformation.py | 125 ++++++++++++------ litellm/types/llms/bedrock.py | 4 +- litellm/types/llms/cohere.py | 17 ++- .../llms/cohere/embed/test_transformation.py | 121 +++++++++++++++++ 4 files changed, 226 insertions(+), 41 deletions(-) create mode 100644 tests/test_litellm/llms/cohere/embed/test_transformation.py diff --git a/litellm/llms/cohere/embed/transformation.py b/litellm/llms/cohere/embed/transformation.py index eb3f65bec94..d6862c9af92 100644 --- a/litellm/llms/cohere/embed/transformation.py +++ b/litellm/llms/cohere/embed/transformation.py @@ -23,6 +23,7 @@ from litellm.types.llms.bedrock import ( CohereEmbeddingRequest, CohereEmbeddingRequestWithModel, ) +from litellm.types.llms.cohere import CohereEmbeddingInputList from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues from litellm.types.utils import EmbeddingResponse, PromptTokensDetailsWrapper, Usage from litellm.utils import is_base64_encoded @@ -91,68 +92,116 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig): return api_base or "https://api.cohere.ai/v2/embed" def _transform_request( - self, model: str, input: list[str], inference_params: dict + self, + model: str, + input: list[str] | CohereEmbeddingInputList, + inference_params: dict, ) -> CohereEmbeddingRequestWithModel: - is_encoded = False - for input_str in input: - is_encoded = is_base64_encoded(input_str) - - if is_encoded: # check if string is b64 encoded image or not - transformed_request = CohereEmbeddingRequestWithModel( + is_structured_input: Final = bool(input) and isinstance(input[0], dict) + transformed_request: Final = ( + CohereEmbeddingRequestWithModel( model=model, - images=input, - input_type="image", - ) - else: - transformed_request = CohereEmbeddingRequestWithModel( - model=model, - texts=input, + inputs=cast( # cast-ok: the first-item check narrows this homogeneous input list + CohereEmbeddingInputList, input + ), input_type=COHERE_DEFAULT_EMBEDDING_INPUT_TYPE, ) + if is_structured_input + else self._transform_string_request( + model=model, + input=cast( # cast-ok: the structured-input branch was excluded above + list[str], input + ), + ) + ) for k, v in inference_params.items(): transformed_request[k] = v return transformed_request + def _transform_string_request( + self, + model: str, + input: list[str], # mutable-ok: Cohere's JSON request schema requires an array + ) -> CohereEmbeddingRequestWithModel: + is_encoded: Final = bool(input) and is_base64_encoded(input[-1]) + return ( + CohereEmbeddingRequestWithModel( + model=model, + images=input, + input_type="image", + ) + if is_encoded + else CohereEmbeddingRequestWithModel( + model=model, + texts=input, + input_type=COHERE_DEFAULT_EMBEDDING_INPUT_TYPE, + ) + ) + + def _normalize_embedding_input( + self, + input: AllEmbeddingInputValues | CohereEmbeddingInputList, + ) -> list[str] | CohereEmbeddingInputList: # mutable-ok: provider transformation consumes JSON arrays + if isinstance(input, str): + return [input] + if not input: + raise ValueError("Input must not be empty") + if isinstance(input[0], dict): + return cast( # cast-ok: the first-item check narrows this homogeneous input list + CohereEmbeddingInputList, input + ) + if isinstance(input[0], list) or isinstance(input[0], int): + raise ValueError("Input must be a list of strings") + return cast(list[str], input) + def transform_embedding_request( self, model: str, - input: AllEmbeddingInputValues, + input: AllEmbeddingInputValues | CohereEmbeddingInputList, optional_params: dict, headers: dict, ) -> dict: - if isinstance(input, list) and (isinstance(input[0], list) or isinstance(input[0], int)): - raise ValueError("Input must be a list of strings") return cast( dict, self._transform_request( model=model, - input=cast(list[str], input) if isinstance(input, list) else [input], + input=self._normalize_embedding_input(input), inference_params=optional_params, ), ) - def _calculate_usage(self, input: list[str], encoding: Any, meta: dict) -> Usage: - input_tokens = 0 - + def _calculate_usage( + self, + input: list[str] | CohereEmbeddingInputList, + encoding: Any, + meta: dict, + ) -> Usage: text_tokens: Final[int | None] = meta.get("billed_units", {}).get("input_tokens") - image_tokens: Final[int | None] = meta.get("billed_units", {}).get("images") - - prompt_tokens_details: PromptTokensDetailsWrapper | None = None - if image_tokens is None and text_tokens is None: - for text in input: - input_tokens += len(encoding.encode(text)) - else: - prompt_tokens_details = PromptTokensDetailsWrapper( + fallback_texts: Final = tuple( + text + for item in input + for text in ( + (item,) + if isinstance(item, str) + else tuple(content["text"] for content in item["content"] if content["type"] == "text") + ) + ) + input_tokens: Final = ( + sum(len(encoding.encode(text)) for text in fallback_texts) + if image_tokens is None and text_tokens is None + else (image_tokens or 0) + (text_tokens or 0) + ) + prompt_tokens_details: Final = ( + None + if image_tokens is None and text_tokens is None + else PromptTokensDetailsWrapper( image_tokens=image_tokens, text_tokens=text_tokens, ) - if image_tokens: - input_tokens += image_tokens - if text_tokens: - input_tokens += text_tokens + ) return Usage( prompt_tokens=input_tokens, @@ -170,7 +219,7 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig): model_response: EmbeddingResponse, model: str, encoding: Any, - input: list, + input: list[str] | CohereEmbeddingInputList, ) -> EmbeddingResponse: response_json: Final = response.json() ## LOGGING @@ -199,9 +248,6 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig): model_response.object = "list" model_response.data = output_data model_response.model = model - input_tokens = 0 - for text in input: - input_tokens += len(encoding.encode(text)) setattr( model_response, @@ -230,7 +276,10 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig): model_response=model_response, model=model, encoding=litellm.encoding, - input=logging_obj.model_call_details["input"], + input=cast( # cast-ok: logging preserves the original provider input without a precise static type + list[str] | CohereEmbeddingInputList, + logging_obj.model_call_details["input"], + ), ) def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 5665aa3277a..1ca2f9b6148 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -2,8 +2,9 @@ import json from enum import Enum from typing import TYPE_CHECKING, Any, Final, Literal -from typing_extensions import Required, TypedDict, override +from typing_extensions import ReadOnly, Required, TypedDict, override +from .cohere import CohereEmbeddingInputList from .openai import ChatCompletionToolCallChunk @@ -484,6 +485,7 @@ COHERE_EMBEDDING_INPUT_TYPES = Literal["search_document", "search_query", "class class CohereEmbeddingRequest(TypedDict, total=False): texts: list[str] images: list[str] + inputs: ReadOnly[CohereEmbeddingInputList] input_type: Required[COHERE_EMBEDDING_INPUT_TYPES] truncate: Literal["NONE", "START", "END"] embedding_types: Literal["float", "int8", "uint8", "binary", "ubinary"] diff --git a/litellm/types/llms/cohere.py b/litellm/types/llms/cohere.py index 92d1a8a573d..95cdb40f416 100644 --- a/litellm/types/llms/cohere.py +++ b/litellm/types/llms/cohere.py @@ -1,6 +1,19 @@ -from typing import Literal +from typing import Literal, TypeAlias -from typing_extensions import Required, TypedDict +from typing_extensions import ReadOnly, Required, TypedDict + +from .openai import ChatCompletionImageObject, ChatCompletionTextObject + + +class CohereEmbeddingInput(TypedDict): + content: ReadOnly[ # mutable-ok: Cohere's JSON request schema requires an array + list[ChatCompletionTextObject | ChatCompletionImageObject] + ] + + +CohereEmbeddingInputList: TypeAlias = ( + list[CohereEmbeddingInput] # mutable-ok: Cohere's JSON request schema requires an array +) class CallObject(TypedDict): diff --git a/tests/test_litellm/llms/cohere/embed/test_transformation.py b/tests/test_litellm/llms/cohere/embed/test_transformation.py new file mode 100644 index 00000000000..73e24c0d876 --- /dev/null +++ b/tests/test_litellm/llms/cohere/embed/test_transformation.py @@ -0,0 +1,121 @@ +from typing import Final +from unittest.mock import MagicMock, patch + +import httpx + +from litellm.llms.cohere.embed.transformation import CohereEmbeddingConfig +from litellm.types.llms.cohere import CohereEmbeddingInput +from litellm.types.utils import EmbeddingResponse + + +def test_transform_embedding_request_preserves_mixed_inputs() -> None: + config: Final = CohereEmbeddingConfig() + inputs: Final[list[CohereEmbeddingInput]] = [ + { + "content": [ + {"type": "text", "text": "a red shoe"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,AAAA"}, + }, + ] + } + ] + + request: Final = config.transform_embedding_request( + model="embed-v4.0", + input=inputs, + optional_params={ + "input_type": "search_document", + "output_dimension": 1536, + "embedding_types": ["float"], + }, + headers={}, + ) + + assert request == { + "model": "embed-v4.0", + "inputs": inputs, + "input_type": "search_document", + "output_dimension": 1536, + "embedding_types": ["float"], + } + assert "texts" not in request + assert "images" not in request + + +def test_transform_embedding_response_uses_multimodal_billing_metadata() -> None: + config: Final = CohereEmbeddingConfig() + inputs: Final[list[CohereEmbeddingInput]] = [ + { + "content": [ + {"type": "text", "text": "a red shoe"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,AAAA"}, + }, + ] + } + ] + response: Final = httpx.Response( + 200, + json={ + "embeddings": {"float": [[0.1, 0.2, 0.3]]}, + "meta": {"billed_units": {"input_tokens": 3, "images": 1}}, + }, + ) + logging_obj: Final = MagicMock() + logging_obj.model_call_details = {"input": inputs} + encoding: Final = MagicMock() + + with patch("litellm.encoding", encoding): + result: Final = config.transform_embedding_response( + model="embed-v4.0", + raw_response=response, + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + api_key="test-api-key", + request_data={ + "model": "embed-v4.0", + "inputs": inputs, + "input_type": "search_document", + }, + optional_params={}, + litellm_params={}, + ) + + assert result.data == [ + { + "object": "embedding", + "index": 0, + "embedding": [0.1, 0.2, 0.3], + } + ] + assert result.usage.prompt_tokens == 4 + assert result.usage.prompt_tokens_details.text_tokens == 3 + assert result.usage.prompt_tokens_details.image_tokens == 1 + encoding.encode.assert_not_called() + + +def test_multimodal_usage_fallback_counts_text_content() -> None: + config: Final = CohereEmbeddingConfig() + inputs: Final[list[CohereEmbeddingInput]] = [ + { + "content": [ + {"type": "text", "text": "a red shoe"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,AAAA"}, + }, + ] + } + ] + encoding: Final = MagicMock() + encoding.encode.return_value = [1, 2, 3] + + usage: Final = config._calculate_usage(inputs, encoding, {}) + + assert usage.prompt_tokens == 3 + assert usage.total_tokens == 3 + assert usage.prompt_tokens_details is None + encoding.encode.assert_called_once_with("a red shoe")