diff --git a/litellm/llms/openrouter/embedding/transformation.py b/litellm/llms/openrouter/embedding/transformation.py index d1530a29515..61d6574cdd6 100644 --- a/litellm/llms/openrouter/embedding/transformation.py +++ b/litellm/llms/openrouter/embedding/transformation.py @@ -7,10 +7,12 @@ OpenRouter is OpenAI-compatible and supports embeddings via the /v1/embeddings e Docs: https://openrouter.ai/docs """ +import math from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.types.llms.openai import AllEmbeddingInputValues @@ -131,11 +133,20 @@ class OpenrouterEmbeddingConfig(BaseEmbeddingConfig): # OpenRouter returns standard OpenAI-compatible embedding response response_json: Final = raw_response.json() - return convert_to_model_response_object( + convert_to_model_response_object( response_object=response_json, model_response_object=model_response, response_type="embedding", ) + response_cost: Final[object] = (response_json.get("usage") or {}).get("cost") + if ( + isinstance(response_cost, int | float) + and not isinstance(response_cost, bool) + and response_cost >= 0 + and math.isfinite(response_cost) + ): + set_response_cost_in_hidden_params(model_response, response_cost) + return model_response def get_supported_openai_params(self, model: str) -> list: """ diff --git a/tests/unit/llms/openrouter/test_openrouter_embedding_transformation.py b/tests/unit/llms/openrouter/test_openrouter_embedding_transformation.py index cb60eff5325..7c56c3c798d 100644 --- a/tests/unit/llms/openrouter/test_openrouter_embedding_transformation.py +++ b/tests/unit/llms/openrouter/test_openrouter_embedding_transformation.py @@ -4,16 +4,105 @@ Unit tests for OpenRouter embedding transformation logic. import json from typing import Final +from unittest.mock import Mock, patch import httpx import pytest import respx import litellm +from litellm.cost_calculator import response_cost_calculator from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.openrouter.embedding.transformation import ( OpenrouterEmbeddingConfig, ) +from litellm.types.utils import EmbeddingResponse + + +def _embedding_response(usage: dict[str, object], model: str = "qwen/qwen3-embedding-8b") -> EmbeddingResponse: + raw_response: Final = httpx.Response( + 200, + text=json.dumps( + { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "model": model, + "usage": usage, + } + ), + ) + return OpenrouterEmbeddingConfig().transform_embedding_response( + model=f"openrouter/{model}", + raw_response=raw_response, + model_response=EmbeddingResponse(), + logging_obj=Mock(), + api_key="test-api-key", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + +@pytest.mark.parametrize("cost", [2e-8, 0]) +def test_openrouter_embedding_preserves_cost_without_price_map(cost: float) -> None: + with patch.dict(litellm.model_cost, {}, clear=True): + response: Final = _embedding_response({"prompt_tokens": 2, "total_tokens": 2, "cost": cost}) + assert response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] == cost + assert ( + response_cost_calculator( + response_object=response, + model="openrouter/qwen/qwen3-embedding-8b", + custom_llm_provider="openrouter", + call_type="embedding", + optional_params={}, + ) + == cost + ) + assert response.data[0]["embedding"] == [0.1, 0.2] + assert response.usage.prompt_tokens == response.usage.total_tokens == 2 + + +@pytest.mark.parametrize( + "cost_fields", + [ + pytest.param({}, id="missing"), + pytest.param({"cost": None}, id="null"), + pytest.param({"cost": "invalid"}, id="string"), + pytest.param({"cost": True}, id="boolean"), + pytest.param({"cost": -1}, id="negative"), + pytest.param({"cost": {}}, id="object"), + pytest.param({"cost": float("nan")}, id="nan"), + pytest.param({"cost": float("inf")}, id="infinity"), + ], +) +@patch.dict( + litellm.model_cost, + { + "openrouter/openai/text-embedding-3-small": { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0, + "litellm_provider": "openrouter", + "mode": "embedding", + } + }, +) +def test_openrouter_embedding_missing_or_invalid_cost_uses_price_map(cost_fields: dict[str, object]) -> None: + response: Final = _embedding_response( + {"prompt_tokens": 2, "total_tokens": 2, **cost_fields}, model="openai/text-embedding-3-small" + ) + assert "llm_provider-x-litellm-response-cost" not in response._hidden_params.get("additional_headers", {}) + assert response.data[0]["embedding"] == [0.1, 0.2] + assert response.usage.prompt_tokens == response.usage.total_tokens == 2 + assert ( + response_cost_calculator( + response_object=response, + model="openai/text-embedding-3-small", + custom_llm_provider="openrouter", + call_type="embedding", + optional_params={}, + ) + == 2 * 0.001 + ) def test_openrouter_embedding_supported_params():