fix(cohere): support multimodal embed inputs

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:

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"]

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):

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")