From 3e3362d4615ad54fc471afdd6a7ea3c5079dfb03 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Thu, 29 Oct 2026 05:17:11 -0700 Subject: [PATCH] fix(cohere): price native Cohere embed calls on Foundry relays and the Cohere passthrough (#48978) * fix(azure_ai): price Foundry relays of the native Cohere v2 embed route A proxy relay of POST /azure_ai//providers/cohere/v2/embed logged spend 0 because the Foundry relay shapes had no entry for it, so the body fell through to the passthrough object with no usage. The Cohere v2 embed body-to-EmbeddingResponse transform is now one shared function that the direct Cohere path and the relay both call, priced from meta.billed_units and falling back to the echoed texts' token count when billed_units is absent. A body that is not a Cohere embed shape raises ValidationError, so the relay fails open as before. * fix(cohere): price native Cohere embeds on the v1 Foundry relay and the v2 passthrough, billing image tokens per caller The Foundry relay now also prices /providers/cohere/v1/embed, and the shared Cohere embed parser accepts the v1 plain list of float vectors next to the by-type object. Direct Cohere bills billed_units.image_tokens (or v1 images) apart from input_tokens, while Foundry already folds the image into input_tokens, so the parser takes an explicit CohereImageTokenBilling and each caller picks its own. The /cohere passthrough now routes and prices /v2/embed through the same parser as the cohere/ SDK path. The relay prices any body that carries billed units whatever texts holds, and the no-billed-units fallback counts special-token text as plain text instead of raising * fix(cohere): keep embed type labels, per-image v3 pricing, and the plain row for unpriceable passthrough embeds Passthrough and Foundry relay log rows keep each vector's embedding type, text-only usage keeps image_tokens null, a Cohere v3 billed image count is priced per image, and an embed call the cost map cannot price falls back to the plain passthrough row instead of raising inside the logger. * fix(cohere): warn when a passthrough embed cannot be priced and pin test timestamps * test(cohere): cover embed-jobs passthrough calls writing the plain passthrough row --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../azure_ai/passthrough/transformation.py | 31 ++- litellm/llms/cohere/embed/transformation.py | 192 ++++++++++------ .../cohere_passthrough_logging_handler.py | 37 ++-- .../pass_through_endpoints/success_handler.py | 11 +- ...est_azure_ai_passthrough_transformation.py | 158 +++++++++++++ .../test_cohere_embedding_transformation.py | 51 ++++- .../llms/cohere/embed/test_transformation.py | 74 ++++++- ...test_cohere_passthrough_logging_handler.py | 208 +++++++++++++++--- 8 files changed, 613 insertions(+), 149 deletions(-) diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index 35b7642d7de..2e5514b3975 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -5,8 +5,9 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError +from litellm._lazy_imports import get_default_encoding from litellm._logging import verbose_logger from litellm.llms.azure_ai.common_utils import ( AzureFoundryModelInfo, @@ -22,10 +23,15 @@ from litellm.llms.base_llm.passthrough.transformation import ( relayed_body, strip_leading_model_segment, ) +from litellm.llms.cohere.embed.transformation import ( + CohereEmbeddingRowShape, + CohereImageTokenBilling, + cohere_embedding_response, +) from litellm.rust_bridge.ocr.entrypoints import NATIVE_OCR_PASSTHROUGH_RESPONSE, NativeOcrPassthroughResponse from litellm.types.llms.openai import AllMessageValues from litellm.types.rerank import RerankResponse -from litellm.types.utils import CallTypes, ImageResponse, StandardPassThroughResponseObject +from litellm.types.utils import CallTypes, EmbeddingResponse, ImageResponse, StandardPassThroughResponseObject if TYPE_CHECKING: from httpx import URL, Response @@ -35,6 +41,7 @@ if TYPE_CHECKING: EMPTY_QUERY: Final[Mapping[str, object]] = MappingProxyType({}) +ECHOED_TEXTS: Final = TypeAdapter(tuple[str, ...], config=ConfigDict(hide_input_in_errors=True)) def api_version_from(litellm_params: Mapping[str, object]) -> str | None: @@ -85,9 +92,29 @@ def relay_query_params( return MappingProxyType({**(request_query_params or EMPTY_QUERY), "api-version": api_version}) +def echoed_texts(body: Mapping[str, object]) -> tuple[str, ...]: + try: + return ECHOED_TEXTS.validate_python(body.get("texts")) + except ValidationError: + return () + + +def cohere_embed_relay_response(body: Mapping[str, object]) -> EmbeddingResponse: + return cohere_embedding_response( + body, + None, + get_default_encoding(), + echoed_texts(body), + CohereImageTokenBilling.FOLDED_INTO_INPUT_TOKENS, + CohereEmbeddingRowShape.TYPED, + ) + + FOUNDRY_RELAY_SHAPES: Final = ( RelayShape("/rerank", CallTypes.arerank, RerankResponse.model_validate), RelayShape("/providers/blackforestlabs/v1/flux-2-pro", CallTypes.aimage_generation, ImageResponse.model_validate), + RelayShape("/providers/cohere/v1/embed", CallTypes.aembedding, cohere_embed_relay_response), + RelayShape("/providers/cohere/v2/embed", CallTypes.aembedding, cohere_embed_relay_response), ) diff --git a/litellm/llms/cohere/embed/transformation.py b/litellm/llms/cohere/embed/transformation.py index 590164ffa0f..1477e45997b 100644 --- a/litellm/llms/cohere/embed/transformation.py +++ b/litellm/llms/cohere/embed/transformation.py @@ -10,11 +10,13 @@ Convers Docs - https://docs.cohere.com/v2/reference/embed """ -from collections.abc import Iterable, Mapping, Sized -from typing import Final, Protocol, cast +from collections.abc import Iterable, Iterator, Mapping, Sequence, Sized +from enum import Enum +from itertools import chain, starmap +from typing import Final, Protocol, TypeAlias, cast import httpx -from pydantic import ConfigDict, TypeAdapter +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter import litellm from litellm import COHERE_DEFAULT_EMBEDDING_INPUT_TYPE @@ -32,14 +34,117 @@ from litellm.utils import is_base64_encoded from ..common_utils import CohereError _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) -_EMBEDDINGS_BY_TYPE: Final = TypeAdapter(Mapping[str, Iterable[object]], config=ConfigDict(hide_input_in_errors=True)) _DECODED_JSON: Final = TypeAdapter(object) class _SupportsEncode(Protocol): """Tokenizer handle: the embedding usage path only encodes text to measure its token length.""" - def encode(self, text: str, /) -> Sized: ... + def encode_ordinary(self, text: str, /) -> Sized: ... + + +class CohereImageTokenBilling(Enum): + SEPARATE_FROM_INPUT_TOKENS = "separate_from_input_tokens" + FOLDED_INTO_INPUT_TOKENS = "folded_into_input_tokens" + + +class CohereEmbeddingRowShape(Enum): + OPENAI = "openai" + TYPED = "typed" + + +class _CohereBilledUnits(BaseModel): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + + input_tokens: int | None = None + image_tokens: int | None = None + images: int | None = None + + +class _CohereEmbedMeta(BaseModel): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + + billed_units: _CohereBilledUnits = Field(default_factory=_CohereBilledUnits) + + +_CohereEmbeddings: TypeAlias = Mapping[str, Sequence[object]] | Sequence[object] + + +class _CohereEmbedBody(BaseModel): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + + embeddings: _CohereEmbeddings + meta: _CohereEmbedMeta = Field(default_factory=_CohereEmbedMeta) + + +def _embedding_rows(vectors: Sequence[object]) -> Iterator[Mapping[str, object]]: + return ({"object": "embedding", "index": index, "embedding": vector} for index, vector in enumerate(vectors)) + + +def _typed_embedding_rows(embedding_type: str, vectors: Sequence[object]) -> Iterator[Mapping[str, object]]: + return ({**row, "type": embedding_type} for row in _embedding_rows(vectors)) + + +def _response_rows(embeddings: _CohereEmbeddings, row_shape: CohereEmbeddingRowShape) -> Iterable[Mapping[str, object]]: + if not isinstance(embeddings, Mapping): + return _embedding_rows(embeddings) + match row_shape: + case CohereEmbeddingRowShape.OPENAI: + return chain.from_iterable(map(_embedding_rows, embeddings.values())) + case CohereEmbeddingRowShape.TYPED: + return chain.from_iterable(starmap(_typed_embedding_rows, embeddings.items())) + + +def _prompt_usage(prompt_tokens: int, prompt_tokens_details: PromptTokensDetailsWrapper | None) -> Usage: + return Usage( + prompt_tokens=prompt_tokens, + completion_tokens=0, + total_tokens=prompt_tokens, + prompt_tokens_details=prompt_tokens_details, + ) + + +def _billed_usage(billed_units: _CohereBilledUnits, image_billing: CohereImageTokenBilling) -> Usage: + match image_billing: + case CohereImageTokenBilling.SEPARATE_FROM_INPUT_TOKENS: + image_units: Final = (billed_units.image_tokens or 0) + (billed_units.images or 0) + return _prompt_usage( + (billed_units.input_tokens or 0) + image_units, + PromptTokensDetailsWrapper( + text_tokens=billed_units.input_tokens, + image_tokens=image_units or None, + image_count=billed_units.images, + ), + ) + case CohereImageTokenBilling.FOLDED_INTO_INPUT_TOKENS: + return _prompt_usage(billed_units.input_tokens or 0, None) + + +def _cohere_embed_usage( + billed_units: _CohereBilledUnits, + input_texts: Iterable[str], + encoding: _SupportsEncode, + image_billing: CohereImageTokenBilling, +) -> Usage: + if billed_units.input_tokens is None and billed_units.image_tokens is None and billed_units.images is None: + return _prompt_usage(sum(len(encoding.encode_ordinary(text)) for text in input_texts), None) + return _billed_usage(billed_units, image_billing) + + +def cohere_embedding_response( + body: Mapping[str, object], + model: str | None, + encoding: _SupportsEncode, + input_texts: Iterable[str], + image_billing: CohereImageTokenBilling, + row_shape: CohereEmbeddingRowShape, +) -> EmbeddingResponse: + cohere_body: Final = _CohereEmbedBody.model_validate(body) + return EmbeddingResponse( + model=model, + data=list(_response_rows(cohere_body.embeddings, row_shape)), + usage=_cohere_embed_usage(cohere_body.meta.billed_units, input_texts, encoding, image_billing), + ) class CohereEmbeddingConfig(BaseEmbeddingConfig): @@ -145,44 +250,16 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig): ), ) - def _calculate_usage(self, input: list[str], encoding: _SupportsEncode, meta: dict) -> Usage: - input_tokens = 0 - - 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( - 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, - completion_tokens=0, - total_tokens=input_tokens, - prompt_tokens_details=prompt_tokens_details, - ) - - def _transform_response( + def transform_native_embed_response( self, response: httpx.Response, api_key: str | None, logging_obj: LiteLLMLoggingObj, - data: dict | CohereEmbeddingRequest, - model_response: EmbeddingResponse, + data: Mapping[str, object] | CohereEmbeddingRequest, model: str, encoding: _SupportsEncode, - input: list, + input: Sequence[str], + row_shape: CohereEmbeddingRowShape, ) -> EmbeddingResponse: response_json: Final = _DECODED_JSON.validate_python(response.json()) ## LOGGING @@ -192,38 +269,15 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig): additional_args={"complete_input_dict": data}, original_response=response_json, ) - """ - response - { - 'object': "list", - 'data': [ - - ] - 'model', - 'usage' - } - """ - payload: Final = _JSON_OBJECT.validate_python(response_json) - embeddings: Final = _EMBEDDINGS_BY_TYPE.validate_python(payload["embeddings"]) - output_data: Final = [] - for k, embedding_list in embeddings.items(): - for idx, embedding in enumerate(embedding_list): - output_data.append({"object": "embedding", "index": idx, "embedding": embedding}) - 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, - "usage", - self._calculate_usage(input, encoding, dict(_JSON_OBJECT.validate_python(payload.get("meta", {})))), + return cohere_embedding_response( + _JSON_OBJECT.validate_python(response_json), + model, + encoding, + input, + CohereImageTokenBilling.SEPARATE_FROM_INPUT_TOKENS, + row_shape, ) - return model_response - def transform_embedding_response( self, model: str, @@ -235,15 +289,15 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig): optional_params: dict, litellm_params: dict, ) -> EmbeddingResponse: - return self._transform_response( + return self.transform_native_embed_response( response=raw_response, api_key=api_key, logging_obj=logging_obj, data=request_data, - model_response=model_response, model=model, encoding=litellm.encoding, input=logging_obj.model_call_details["input"], + row_shape=CohereEmbeddingRowShape.OPENAI, ) def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py index a11b0be3983..4eeb9b26174 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py @@ -1,10 +1,12 @@ from datetime import datetime from typing import Final +from urllib.parse import urlparse import httpx import litellm from litellm import stream_chunk_builder +from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.hidden_params import set_hidden_param from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import ( @@ -16,7 +18,7 @@ from litellm.llms.cohere.chat.v2_transformation import CohereV2ChatConfig from litellm.llms.cohere.common_utils import ( ModelResponseIterator as CohereModelResponseIterator, ) -from litellm.llms.cohere.embed.v1_transformation import CohereEmbeddingConfig +from litellm.llms.cohere.embed.transformation import CohereEmbeddingConfig, CohereEmbeddingRowShape from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.types.passthrough_endpoints.pass_through_endpoints import ( PassthroughStandardLoggingPayload, @@ -29,6 +31,12 @@ from litellm.types.utils import ( from .base_passthrough_logging_handler import BasePassthroughLoggingHandler +COHERE_EMBED_PATHS: Final = ("/v1/embed", "/v2/embed") + + +def is_cohere_embed_route(url_route: str) -> bool: + return urlparse(url_route).path.rstrip("/").endswith(COHERE_EMBED_PATHS) + class CoherePassthroughLoggingHandler(BasePassthroughLoggingHandler): @property @@ -82,28 +90,24 @@ class CoherePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ Handle Cohere passthrough logging with route detection and cost tracking. """ - # Check if this is an embed endpoint - if "/v1/embed" in url_route and "/v1/embeddings" not in url_route: + if is_cohere_embed_route(url_route): model: Final = request_body.get("model", response_body.get("model", "")) try: - cohere_embed_config: Final = CohereEmbeddingConfig() - litellm_model_response = litellm.EmbeddingResponse() handler_instance: Final = CoherePassthroughLoggingHandler() input_texts = request_body.get("texts", []) if not input_texts: input_texts = request_body.get("input", []) - # Transform the response - litellm_model_response = cohere_embed_config.transform_response( + litellm_model_response: Final = CohereEmbeddingConfig().transform_native_embed_response( response=httpx_response, api_key="", logging_obj=logging_obj, data=request_body, - model_response=litellm_model_response, model=model, encoding=litellm.encoding, input=input_texts, + row_shape=CohereEmbeddingRowShape.TYPED, ) # Calculate cost using LiteLLM's cost calculator @@ -153,20 +157,11 @@ class CoherePassthroughLoggingHandler(BasePassthroughLoggingHandler): "result": litellm_model_response, "kwargs": kwargs, } - except Exception: - # For other routes (e.g., /v2/chat), fall back to chat handler - return super().passthrough_chat_handler( - httpx_response=httpx_response, - response_body=response_body, - logging_obj=logging_obj, - url_route=url_route, - result=result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, - request_body=request_body, - **kwargs, + except Exception as e: + verbose_proxy_logger.warning( + "Could not price the Cohere passthrough embed response, logging it without a cost: %s", e ) + return {"result": None, "kwargs": kwargs} # For non-embed routes (e.g., /v2/chat), fall back to chat handler return super().passthrough_chat_handler( diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 9c6c3111209..db06a85fb3f 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -22,6 +22,7 @@ from .llm_provider_handlers.assembly_passthrough_logging_handler import ( ) from .llm_provider_handlers.cohere_passthrough_logging_handler import ( CoherePassthroughLoggingHandler, + is_cohere_embed_route, ) from .llm_provider_handlers.cursor_passthrough_logging_handler import ( CursorPassthroughLoggingHandler, @@ -91,7 +92,7 @@ class PassThroughEndpointLogging: self.TRACKED_ANTHROPIC_ROUTES = ["/messages", "/v1/messages/batches"] # Cohere - self.TRACKED_COHERE_ROUTES = ["/v2/chat", "/v1/embed"] + self.TRACKED_COHERE_ROUTES = ["/v2/chat"] self.assemblyai_passthrough_logging_handler = AssemblyAIPassthroughLoggingHandler() # Langfuse @@ -536,13 +537,7 @@ class PassThroughEndpointLogging: return False def is_cohere_route(self, url_route: str) -> bool: - for route in self.TRACKED_COHERE_ROUTES: - if route not in url_route: - continue - if route == "/v1/embed" and "/v1/embeddings" in url_route: - continue - return True - return False + return is_cohere_embed_route(url_route) or any(route in url_route for route in self.TRACKED_COHERE_ROUTES) def is_assemblyai_route(self, url_route: str): parsed_url: Final = urlparse(url_route) diff --git a/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 24f21751065..2d8ea27edb7 100644 --- a/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -1,5 +1,7 @@ import json +from collections.abc import Mapping from datetime import datetime +from typing import Final from unittest.mock import MagicMock import httpx @@ -519,6 +521,162 @@ def test_cohere_rerank_relay_is_costed_per_search_unit(): assert logging_obj.response_cost_calculator(result=result) == pytest.approx(2 * per_query) +COHERE_EMBED_BODY: Final[Mapping[str, object]] = { + "id": "embed-1", + "texts": ["hello world"], + "embeddings": {"float": [[0.1, 0.2]]}, + "meta": {"api_version": {"version": "2"}, "billed_units": {"input_tokens": 2, "image_tokens": 0}}, + "response_type": "embeddings_by_type", +} +COHERE_V1_EMBED_META: Final[Mapping[str, object]] = { + "api_version": {"version": "1"}, + "billed_units": {"input_tokens": 2, "image_tokens": 0}, +} +TWO_VECTORS: Final = [[0.1, 0.2], [0.3, 0.4]] + + +def _cohere_embed_relay_result( + body: Mapping[str, object], native_path: str = "providers/cohere/v2/embed" +) -> tuple[object, Logging]: + return _relay_logging_result(AzureAIPassthroughConfig(), "embed-v-4-0", native_path, body) + + +def _foundry_embed_price(price_key: str) -> float: + return litellm.get_model_info("azure_ai/embed-v-4-0")[price_key] + + +def test_native_cohere_embed_relay_is_costed_per_billed_input_token() -> None: + result, logging_obj = _cohere_embed_relay_result(COHERE_EMBED_BODY) + per_token: Final = _foundry_embed_price("input_cost_per_token") + + assert isinstance(result, EmbeddingResponse) + assert logging_obj.call_type == "aembedding" + assert per_token > 0 + assert logging_obj.response_cost_calculator(result=result) == pytest.approx(2 * per_token) + + +@pytest.mark.parametrize( + "body", + [ + { + "id": "embed-1", + "texts": ["hello world", "second text"], + "embeddings": {"float": TWO_VECTORS}, + "meta": COHERE_V1_EMBED_META, + "response_type": "embeddings_by_type", + }, + { + "id": "embed-1", + "texts": ["hello world", "second text"], + "embeddings": TWO_VECTORS, + "meta": COHERE_V1_EMBED_META, + "response_type": "embeddings_floats", + }, + ], + ids=["embeddings_by_type", "embeddings_floats"], +) +def test_native_cohere_v1_embed_relay_is_costed_per_billed_input_token(body: Mapping[str, object]) -> None: + result, logging_obj = _cohere_embed_relay_result(body, "providers/cohere/v1/embed") + per_token: Final = _foundry_embed_price("input_cost_per_token") + + assert isinstance(result, EmbeddingResponse) + assert [(row["index"], row["embedding"]) for row in result.data] == list(enumerate(TWO_VECTORS)) + assert logging_obj.call_type == "aembedding" + assert per_token > 0 + assert logging_obj.response_cost_calculator(result=result) == pytest.approx(2 * per_token) + + +def test_native_cohere_embed_relay_prices_an_image_at_the_input_tokens_foundry_already_folded_it_into() -> None: + body: Final = { + **COHERE_EMBED_BODY, + "texts": [], + "images": [{"width": 64, "height": 64, "format": "image/png", "bit_depth": 8}], + "meta": {"api_version": {"version": "2"}, "billed_units": {"input_tokens": 24, "image_tokens": 6}}, + } + result, logging_obj = _cohere_embed_relay_result(body) + per_token: Final = _foundry_embed_price("input_cost_per_token") + per_image_token: Final = _foundry_embed_price("input_cost_per_image_token") + + assert isinstance(result, EmbeddingResponse) + assert per_token > 0 + assert per_image_token > 0 + assert logging_obj.response_cost_calculator(result=result) == pytest.approx(24 * per_token) + + + +def test_native_cohere_embed_relay_rows_keep_their_embedding_type() -> None: + result, _ = _cohere_embed_relay_result({**COHERE_EMBED_BODY, "embeddings": {"float": TWO_VECTORS, "int8": [[1, 2]]}}) + + assert isinstance(result, EmbeddingResponse) + assert [(row["type"], row["index"], row["embedding"]) for row in result.data] == [ + ("float", 0, TWO_VECTORS[0]), + ("float", 1, TWO_VECTORS[1]), + ("int8", 0, [1, 2]), + ] + +@pytest.mark.parametrize("texts", [None, "hello world", [1, 2], ["hello", 3]]) +def test_native_cohere_embed_relay_with_billed_units_is_costed_whatever_texts_holds(texts: object) -> None: + result, logging_obj = _cohere_embed_relay_result({**COHERE_EMBED_BODY, "texts": texts}) + + assert isinstance(result, EmbeddingResponse) + assert logging_obj.call_type == "aembedding" + assert logging_obj.response_cost_calculator(result=result) == pytest.approx( + 2 * _foundry_embed_price("input_cost_per_token") + ) + + +def test_native_cohere_embed_relay_without_texts_is_costed_per_billed_input_token() -> None: + body: Final = {key: value for key, value in COHERE_EMBED_BODY.items() if key != "texts"} + result, logging_obj = _cohere_embed_relay_result(body) + + assert isinstance(result, EmbeddingResponse) + assert logging_obj.call_type == "aembedding" + assert logging_obj.response_cost_calculator(result=result) == pytest.approx( + 2 * _foundry_embed_price("input_cost_per_token") + ) + + +@pytest.mark.parametrize( + ("texts", "counted_texts"), + [ + (["hello world", "a longer second text"], ["hello world", "a longer second text"]), + (["before <|endoftext|> after"], ["before <|endoftext|> after"]), + ("hello world", []), + (["hello", 3], []), + (None, []), + ], +) +def test_native_cohere_embed_relay_without_billed_units_is_costed_from_the_echoed_texts( + texts: object, counted_texts: list[str] +) -> None: + result, logging_obj = _cohere_embed_relay_result({**COHERE_EMBED_BODY, "texts": texts, "meta": {}}) + per_token: Final = _foundry_embed_price("input_cost_per_token") + echoed_tokens: Final = sum(len(litellm.encoding.encode(text, disallowed_special=())) for text in counted_texts) + + assert isinstance(result, EmbeddingResponse) + assert logging_obj.call_type == "aembedding" + assert result.usage.prompt_tokens == echoed_tokens + assert logging_obj.response_cost_calculator(result=result) == pytest.approx(echoed_tokens * per_token) + + +@pytest.mark.parametrize( + "body", + [ + {"message": "no embeddings here"}, + {**COHERE_EMBED_BODY, "embeddings": None}, + {**COHERE_EMBED_BODY, "embeddings": "vectors"}, + {**COHERE_EMBED_BODY, "meta": "not an object"}, + ], +) +def test_native_cohere_embed_relay_with_a_non_embed_body_keeps_the_passthrough_object_and_call_type( + body: Mapping[str, object], +) -> None: + result, logging_obj = _cohere_embed_relay_result(body) + + assert result == {"response": body} + assert logging_obj.call_type == "allm_passthrough_route" + + def test_image_generation_relay_is_costed_per_image(): result, logging_obj = _relay_logging_result( AzureAIPassthroughConfig(), "FLUX.2-pro", "openai/deployments/FLUX.2-pro/images/generations", IMAGE_BODY diff --git a/tests/unit/llms/cohere/embed/test_cohere_embedding_transformation.py b/tests/unit/llms/cohere/embed/test_cohere_embedding_transformation.py index 516ea1f5221..245d30a6283 100644 --- a/tests/unit/llms/cohere/embed/test_cohere_embedding_transformation.py +++ b/tests/unit/llms/cohere/embed/test_cohere_embedding_transformation.py @@ -7,7 +7,12 @@ import respx import litellm from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.llms.cohere.embed.transformation import CohereEmbeddingConfig +from litellm.llms.cohere.embed.transformation import ( + CohereEmbeddingConfig, + CohereEmbeddingRowShape, + CohereImageTokenBilling, + cohere_embedding_response, +) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.utils import CallTypes, EmbeddingResponse @@ -108,18 +113,50 @@ def test_cohere_embedding_response_counts_tokens_when_provider_usage_is_missing( assert response.usage.total_tokens == response.usage.prompt_tokens -def test_cohere_embedding_usage_includes_text_and_image_token_details(): - config: Final = CohereEmbeddingConfig() - usage: Final = config._calculate_usage( - input=["hello"], - encoding=litellm.encoding, - meta={"billed_units": {"input_tokens": 5, "images": 2}}, +def test_cohere_embedding_fallback_counts_special_token_text_as_plain_text() -> None: + text: Final = "before <|endoftext|> after" + response: Final = CohereEmbeddingConfig().transform_embedding_response( + model="embed-english-v3.0", + raw_response=httpx.Response(200, json={"embeddings": {"float": [[0.1]]}}), + model_response=EmbeddingResponse(), + logging_obj=_logging_obj([text]), + api_key="test-key", + request_data={"texts": [text]}, + optional_params={}, + litellm_params={}, ) + assert response.usage.prompt_tokens == len(litellm.encoding.encode(text, disallowed_special=())) + + +@pytest.mark.parametrize("image_key", ["images", "image_tokens"]) +def test_cohere_embedding_usage_includes_text_and_image_token_details(image_key: str) -> None: + usage: Final = cohere_embedding_response( + {"embeddings": {"float": [[0.1]]}, "meta": {"billed_units": {"input_tokens": 5, image_key: 2}}}, + "embed-english-v3.0", + litellm.encoding, + ["hello"], + CohereImageTokenBilling.SEPARATE_FROM_INPUT_TOKENS, + CohereEmbeddingRowShape.OPENAI, + ).usage + assert usage.prompt_tokens == 7 assert (usage.prompt_tokens_details.text_tokens, usage.prompt_tokens_details.image_tokens) == (5, 2) +def test_cohere_embedding_usage_folded_image_tokens_count_input_tokens_only() -> None: + usage: Final = cohere_embedding_response( + {"embeddings": {"float": [[0.1]]}, "meta": {"billed_units": {"input_tokens": 24, "image_tokens": 6}}}, + "embed-v-4-0", + litellm.encoding, + [], + CohereImageTokenBilling.FOLDED_INTO_INPUT_TOKENS, + CohereEmbeddingRowShape.OPENAI, + ).usage + + assert (usage.prompt_tokens, usage.total_tokens, usage.prompt_tokens_details) == (24, 24, None) + + def test_cohere_embedding_error_class_preserves_status_and_message(): error: Final = CohereEmbeddingConfig().get_error_class( error_message="invalid embedding request", diff --git a/tests/unit/llms/cohere/embed/test_transformation.py b/tests/unit/llms/cohere/embed/test_transformation.py index c0f355cfb57..abae30013fe 100644 --- a/tests/unit/llms/cohere/embed/test_transformation.py +++ b/tests/unit/llms/cohere/embed/test_transformation.py @@ -1,6 +1,6 @@ import datetime import json -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from typing import Final import httpx @@ -70,8 +70,11 @@ def test_cohere_embedding_response_lists_every_embedding_type_in_provider_order( ("meta", "expected_counts", "expected_details"), [ ({"billed_units": {"input_tokens": 5}}, (5, 0, 5), (5, None)), + ({"billed_units": {"input_tokens": 5, "image_tokens": 0}}, (5, 0, 5), (5, None)), ({"billed_units": {"images": 3}}, (3, 0, 3), (None, 3)), ({"billed_units": {"input_tokens": 5, "images": 3}}, (8, 0, 8), (5, 3)), + ({"billed_units": {"input_tokens": 0, "image_tokens": 6}}, (6, 0, 6), (0, 6)), + ({"billed_units": {"input_tokens": 3, "image_tokens": 6}}, (9, 0, 9), (3, 6)), ({"billed_units": {"input_tokens": 0}}, (0, 0, 0), (0, None)), ( {"api_version": {"version": "2"}, "billed_units": {"input_tokens": 2, "output_tokens": 9}}, @@ -80,11 +83,15 @@ def test_cohere_embedding_response_lists_every_embedding_type_in_provider_order( ), ], ) -def test_cohere_embedding_usage_comes_from_billed_units(meta, expected_counts, expected_details): - response = _transform({"embeddings": {"float": FLOAT_VECTORS}, "meta": meta}) +def test_cohere_embedding_usage_comes_from_billed_units( + meta: Mapping[str, object], + expected_counts: tuple[int, int, int], + expected_details: tuple[int | None, int | None], +) -> None: + response: Final = _transform({"embeddings": {"float": FLOAT_VECTORS}, "meta": meta}) assert response.usage is not None - details = response.usage.prompt_tokens_details + details: Final = response.usage.prompt_tokens_details assert details is not None assert _token_counts(response) == expected_counts assert (details.text_tokens, details.image_tokens) == expected_details @@ -123,8 +130,20 @@ def test_cohere_embedding_rows_pass_through_whatever_each_type_lists(embeddings, assert [row["embedding"] for row in response.data] == expected_vectors -@pytest.mark.parametrize("embeddings", [FLOAT_VECTORS, [], None, "vectors", 3]) -def test_cohere_embeddings_that_are_not_keyed_by_type_are_rejected(embeddings): +@pytest.mark.parametrize("vectors", [FLOAT_VECTORS, []]) +def test_cohere_v1_float_embeddings_list_one_row_per_vector(vectors: list[list[float]]) -> None: + response: Final = _transform( + {"embeddings": vectors, "meta": {"billed_units": {"input_tokens": 1}}, "response_type": "embeddings_floats"} + ) + + assert response.data == [ + {"object": "embedding", "index": index, "embedding": vector} for index, vector in enumerate(vectors) + ] + assert _token_counts(response) == (1, 0, 1) + + +@pytest.mark.parametrize("embeddings", [None, "vectors", 3]) +def test_cohere_embeddings_that_are_neither_keyed_by_type_nor_a_list_are_rejected(embeddings: object) -> None: with pytest.raises(ValidationError): _transform({"embeddings": embeddings, "meta": {"billed_units": {"input_tokens": 1}}}) @@ -147,12 +166,15 @@ def test_cohere_embedding_body_that_is_not_an_object_is_rejected(payload): _transform(payload) -def test_cohere_embedding_body_without_embeddings_is_rejected(): - with pytest.raises(KeyError): +def test_cohere_embedding_body_without_embeddings_is_rejected() -> None: + with pytest.raises(ValidationError): _transform({"meta": {"billed_units": {"input_tokens": 1}}}) COHERE_V2_EMBED_URL: Final = "https://api.cohere.ai/v2/embed" +ONE_PIXEL_PNG: Final = ( + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC" +) @pytest.fixture @@ -362,13 +384,43 @@ async def test_cohere_embed_v4_with_optional_params(respx_mock: respx.MockRouter assert response.data == [{"object": "embedding", "index": 0, "embedding": vector}] +@pytest.mark.usefixtures("_cohere_httpx_transport") +@pytest.mark.respx(assert_all_called=True) +@pytest.mark.parametrize(("input_tokens", "image_tokens"), [(0, 6), (3, 6)]) +async def test_cohere_embed_v4_image_tokens_are_billed_apart_from_input_tokens( + respx_mock: respx.MockRouter, input_tokens: int, image_tokens: int +) -> None: + respx_mock.post(COHERE_V2_EMBED_URL).mock( + return_value=httpx.Response( + 200, + json={ + "id": "5807ee2e-0cda-445a-9ec8-864c60a06606", + "embeddings": {"float": [[0.125, -0.375]]}, + "texts": [], + "images": [{"width": 1, "height": 1, "format": "image/png", "bit_depth": 8}], + "meta": { + "api_version": {"version": "2"}, + "billed_units": {"input_tokens": input_tokens, "image_tokens": image_tokens}, + }, + "response_type": "embeddings_by_type", + }, + ) + ) + model_info: Final = litellm.get_model_info("cohere/embed-v4.0") + + response: Final = await _embed(True, input=[ONE_PIXEL_PNG], input_type="image") + + assert model_info["input_cost_per_image_token"] > 0 + assert response._hidden_params["response_cost"] == pytest.approx( + input_tokens * model_info["input_cost_per_token"] + image_tokens * model_info["input_cost_per_image_token"] + ) + + @pytest.mark.usefixtures("_cohere_httpx_transport") @pytest.mark.respx(assert_all_called=True) @pytest.mark.parametrize("sync_mode", [True, False]) async def test_cohere_embed_v4_image_embedding(respx_mock: respx.MockRouter, sync_mode: bool) -> None: - image: Final = ( - "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC" - ) + image: Final = ONE_PIXEL_PNG route: Final = respx_mock.post(COHERE_V2_EMBED_URL).mock( return_value=httpx.Response( 200, diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py index 32e87d93fdb..f2bb9eaa9a5 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py @@ -1,18 +1,23 @@ import json +import logging +from collections.abc import Mapping from datetime import datetime +from typing import Final from unittest.mock import MagicMock, patch import httpx import pytest - +import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.llm_provider_handlers.cohere_passthrough_logging_handler import ( CoherePassthroughLoggingHandler, ) +from litellm.proxy.pass_through_endpoints.success_handler import PassThroughEndpointLogging from litellm.types.passthrough_endpoints.pass_through_endpoints import ( PassthroughStandardLoggingPayload, ) +from litellm.types.utils import CallTypes class TestCoherePassthroughLoggingHandler: @@ -64,46 +69,25 @@ class TestCoherePassthroughLoggingHandler: ) @patch("litellm.completion_cost") - @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload") - @patch("litellm.llms.cohere.embed.v1_transformation.CohereEmbeddingConfig._transform_response") - def test_cohere_embed_passthrough_cost_tracking( - self, mock_transform_response, mock_get_standard_logging, mock_completion_cost - ): + def test_cohere_embed_passthrough_cost_tracking(self, mock_completion_cost: MagicMock) -> None: """Test successful cost tracking for Cohere embed passthrough""" - # Arrange - from litellm.types.utils import EmbeddingResponse + mock_completion_cost.return_value = 3.6e-07 - # Create a mock embedding response - mock_embedding_response = EmbeddingResponse() - mock_embedding_response.data = [ - {"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}, - {"object": "embedding", "index": 1, "embedding": [0.4, 0.5, 0.6]}, - ] - mock_embedding_response.model = "embed-english-v3.0" - mock_embedding_response.object = "list" - from litellm.types.utils import Usage + mock_httpx_response: Final = self._create_mock_httpx_response() + mock_logging_obj: Final = self._create_mock_logging_obj() + passthrough_payload: Final = self._create_passthrough_logging_payload() - mock_embedding_response.usage = Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3) - - mock_transform_response.return_value = mock_embedding_response - mock_completion_cost.return_value = 3.6e-07 # Expected cost for embed-v4.0 - mock_get_standard_logging.return_value = {"test": "logging_payload"} - - mock_httpx_response = self._create_mock_httpx_response() - mock_logging_obj = self._create_mock_logging_obj() - passthrough_payload = self._create_passthrough_logging_payload() - - kwargs = { + kwargs: Final = { "passthrough_logging_payload": passthrough_payload, } - request_body = { + request_body: Final = { "model": "embed-english-v3.0", "texts": ["test passthrough"], } # Act - result = self.handler.cohere_passthrough_handler( + result: Final = self.handler.cohere_passthrough_handler( httpx_response=mock_httpx_response, response_body=self.mock_cohere_embed_response, logging_obj=mock_logging_obj, @@ -126,7 +110,7 @@ class TestCoherePassthroughLoggingHandler: # Verify cost calculation was called with correct parameters mock_completion_cost.assert_called_once() - call_args = mock_completion_cost.call_args + call_args: Final = mock_completion_cost.call_args assert call_args.kwargs["model"] == "embed-english-v3.0" assert call_args.kwargs["custom_llm_provider"] == "cohere" assert call_args.kwargs["call_type"] == "aembedding" @@ -174,5 +158,167 @@ class TestCoherePassthroughLoggingHandler: assert result == {"result": None, "kwargs": {}} +COHERE_EMBED_V2_TEXT_BODY: Final[Mapping[str, object]] = { + "id": "da6e531f-83f4-4bcb-8e9a-f1a2c6e7a1b3", + "embeddings": {"float": [[0.1, 0.2], [0.3, 0.4]]}, + "texts": ["hello world", "second text"], + "meta": {"api_version": {"version": "2"}, "billed_units": {"input_tokens": 9}}, + "response_type": "embeddings_by_type", +} +COHERE_EMBED_V2_IMAGE_BODY: Final[Mapping[str, object]] = { + "id": "5807ee2e-0cda-445a-9ec8-864c60a06606", + "embeddings": {"float": [[0.125, -0.375]]}, + "texts": [], + "images": [{"width": 64, "height": 64, "format": "image/png", "bit_depth": 8}], + "meta": {"api_version": {"version": "2"}, "billed_units": {"input_tokens": 0, "image_tokens": 6}}, + "response_type": "embeddings_by_type", +} +COHERE_EMBED_V1_FLOATS_BODY: Final[Mapping[str, object]] = { + "id": "5d4e1f3a-7c1b-4f3e-9a3e-1b2c3d4e5f60", + "embeddings": [[0.1, 0.2], [0.3, 0.4]], + "texts": ["hello world", "second text"], + "meta": {"api_version": {"version": "1"}, "billed_units": {"input_tokens": 3}}, + "response_type": "embeddings_floats", +} + + +_CALL_START: Final = datetime(2026, 10, 1, 12, 0, 0) +_CALL_END: Final = datetime(2026, 10, 1, 12, 0, 1) + + +def _passthrough_logging_obj() -> LiteLLMLoggingObj: + return LiteLLMLoggingObj( + model="unknown", + messages=[{"role": "user", "content": "no-message-pass-through-endpoint"}], + stream=False, + call_type=CallTypes.pass_through.value, + start_time=_CALL_START, + litellm_call_id="cohere-embed-passthrough-call", + function_id="cohere-embed-passthrough-function", + ) + + +@pytest.mark.parametrize( + ("url_route", "body", "input_tokens", "image_tokens"), + [ + ("https://api.cohere.com/v2/embed", COHERE_EMBED_V2_TEXT_BODY, 9, 0), + ("https://api.cohere.com/v2/embed", COHERE_EMBED_V2_IMAGE_BODY, 0, 6), + ("https://api.cohere.com/v1/embed", COHERE_EMBED_V1_FLOATS_BODY, 3, 0), + ], + ids=["v2_text", "v2_image", "v1_floats"], +) +def test_cohere_embed_passthrough_is_costed_from_billed_units( + url_route: str, body: Mapping[str, object], input_tokens: int, image_tokens: int +) -> None: + request_body: Final = { + "model": "embed-v4.0", + "texts": ["hello world", "second text"], + "input_type": "search_document", + } + model_info: Final = litellm.get_model_info("cohere/embed-v4.0") + expected_cost: Final = ( + input_tokens * model_info["input_cost_per_token"] + image_tokens * model_info["input_cost_per_image_token"] + ) + + logged: Final = PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( + httpx_response=httpx.Response(200, json=body, request=httpx.Request("POST", url_route)), + response_body=dict(body), + request_body=request_body, + logging_obj=_passthrough_logging_obj(), + url_route=url_route, + result="", + start_time=_CALL_START, + end_time=_CALL_END, + cache_hit=False, + ) + + assert expected_cost > 0 + assert logged["kwargs"]["custom_llm_provider"] == "cohere" + assert logged["kwargs"]["response_cost"] == pytest.approx(expected_cost) + + + +def _normalize_cohere_embed(url_route: str, body: Mapping[str, object], model: str) -> Mapping[str, object]: + return PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( + httpx_response=httpx.Response(200, json=body, request=httpx.Request("POST", url_route)), + response_body=dict(body), + request_body={"model": model, "texts": ["hello world"], "input_type": "search_document"}, + logging_obj=_passthrough_logging_obj(), + url_route=url_route, + result="", + start_time=_CALL_START, + end_time=_CALL_END, + cache_hit=False, + ) + + +def test_cohere_v3_image_embed_passthrough_is_costed_per_billed_image() -> None: + body: Final = { + "id": "8a1b2c3d-4e5f-4a6b-8c7d-9e0f1a2b3c4d", + "embeddings": {"float": [[0.125, -0.375]]}, + "texts": [], + "images": [{"width": 64, "height": 64, "format": "image/png", "bit_depth": 8}], + "meta": {"api_version": {"version": "2"}, "billed_units": {"images": 1}}, + "response_type": "embeddings_by_type", + } + per_image: Final = litellm.get_model_info("embed-english-v3.0", custom_llm_provider="cohere")["input_cost_per_image"] + + logged: Final = _normalize_cohere_embed("https://api.cohere.com/v2/embed", body, "embed-english-v3.0") + + assert per_image is not None and per_image > 0 + assert logged["kwargs"]["response_cost"] == pytest.approx(per_image) + + +def test_cohere_embed_passthrough_rows_keep_their_embedding_type() -> None: + body: Final = { + **COHERE_EMBED_V2_TEXT_BODY, + "embeddings": {"float": [[0.1, 0.2], [0.3, 0.4]], "int8": [[1, 2], [3, 4]]}, + } + + logged: Final = _normalize_cohere_embed("https://api.cohere.com/v2/embed", body, "embed-v4.0") + + assert [(row["type"], row["index"], row["embedding"]) for row in logged["standard_logging_response_object"].data] == [ + ("float", 0, [0.1, 0.2]), + ("float", 1, [0.3, 0.4]), + ("int8", 0, [1, 2]), + ("int8", 1, [3, 4]), + ] + + +def test_unpriceable_cohere_embed_passthrough_falls_back_to_the_plain_passthrough_row( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level(logging.WARNING): + logged: Final = _normalize_cohere_embed( + "https://api.cohere.com/v2/embed", COHERE_EMBED_V2_TEXT_BODY, "embed-model-missing-from-the-cost-map" + ) + + assert logged["standard_logging_response_object"] is None + assert "response_cost" not in logged["kwargs"] + assert any( + "Could not price the Cohere passthrough embed response" in record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING + ) + + +@pytest.mark.parametrize( + ("url_route", "body"), + [ + ("https://api.cohere.com/v1/embed-jobs", {"job_id": "job-1", "meta": {"api_version": {"version": "1"}}}), + ( + "https://api.cohere.com/v1/embed-jobs/job-1", + {"job_id": "job-1", "status": "processing", "model": "embed-english-v3.0"}, + ), + ], +) +def test_cohere_embed_job_passthrough_logs_the_plain_passthrough_row( + url_route: str, body: Mapping[str, object] +) -> None: + logged: Final = _normalize_cohere_embed(url_route, body, "embed-english-v3.0") + + assert logged["standard_logging_response_object"] is None + assert "response_cost" not in logged["kwargs"] + if __name__ == "__main__": pytest.main([__file__])