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__])