diff --git a/litellm/llms/openrouter/embedding/transformation.py b/litellm/llms/openrouter/embedding/transformation.py index 14b0e462ea7..fa826fd69cd 100644 --- a/litellm/llms/openrouter/embedding/transformation.py +++ b/litellm/llms/openrouter/embedding/transformation.py @@ -10,7 +10,9 @@ Docs: https://openrouter.ai/docs from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter +from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.types.llms.openai import AllEmbeddingInputValues from litellm.types.utils import EmbeddingResponse @@ -26,6 +28,30 @@ else: LiteLLMLoggingObj = Any +class _OpenRouterEmbeddingCostDetails(BaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + + upstream_inference_cost: float | None = None + upstream_inference_prompt_cost: float | None = None + upstream_inference_completions_cost: float | None = None + + +class _OpenRouterEmbeddingUsage(BaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + + cost: float | None = None + cost_details: _OpenRouterEmbeddingCostDetails | None = None + + +class _OpenRouterEmbeddingResponse(BaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + + usage: _OpenRouterEmbeddingUsage | None = None + + +_OPENROUTER_EMBEDDING_RESPONSE_ADAPTER: Final = TypeAdapter(dict[str, JsonValue]) + + class OpenrouterEmbeddingConfig(BaseEmbeddingConfig): """ Configuration for OpenRouter's Embedding API. @@ -136,13 +162,20 @@ class OpenrouterEmbeddingConfig(BaseEmbeddingConfig): logging_obj.post_call(original_response=raw_response.text) # OpenRouter returns standard OpenAI-compatible embedding response - response_json: Final = raw_response.json() - - return convert_to_model_response_object( + response_json: Final = _OPENROUTER_EMBEDDING_RESPONSE_ADAPTER.validate_json(raw_response.content) + usage: Final = _OpenRouterEmbeddingResponse.model_validate(response_json).usage or _OpenRouterEmbeddingUsage() + convert_to_model_response_object( response_object=response_json, model_response_object=model_response, response_type="embedding", + hidden_params=( + {"response_cost_details": usage.cost_details.model_dump(exclude_none=True)} + if usage.cost_details is not None + else None + ), ) + set_response_cost_in_hidden_params(model_response, usage.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 a9d4b14ef73..a1b8449fbc4 100644 --- a/tests/unit/llms/openrouter/test_openrouter_embedding_transformation.py +++ b/tests/unit/llms/openrouter/test_openrouter_embedding_transformation.py @@ -2,6 +2,12 @@ Unit tests for OpenRouter embedding transformation logic. """ +from unittest.mock import Mock + +import httpx +from litellm.cost_calculator import response_cost_calculator +from litellm.types.utils import EmbeddingResponse + from litellm.llms.openrouter.embedding.transformation import ( OpenrouterEmbeddingConfig, ) @@ -131,3 +137,115 @@ def test_openrouter_embedding_map_params(): assert result["timeout"] == 30 # Unsupported params should not be included assert "unsupported" not in result + + +def test_openrouter_embedding_preserves_provider_reported_cost(): + config = OpenrouterEmbeddingConfig() + raw_response = httpx.Response( + status_code=200, + json={ + "data": [{"embedding": [0.1, 0.2], "index": 0, "object": "embedding"}], + "id": "gen-emb-test", + "model": "text-embedding-3-small", + "object": "list", + "provider": "OpenAI", + "usage": { + "prompt_tokens": 3, + "total_tokens": 3, + "cost": 0.00000006, + "is_byok": False, + "cost_details": { + "upstream_inference_cost": 0.00000006, + "upstream_inference_prompt_cost": 0.00000006, + "upstream_inference_completions_cost": 0, + }, + }, + }, + ) + + response = config.transform_embedding_response( + model="openrouter/openai/text-embedding-3-small", + raw_response=raw_response, + model_response=EmbeddingResponse(), + logging_obj=Mock(), + api_key="test-api-key", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] == 0.00000006 + assert response._hidden_params["response_cost_details"] == { + "upstream_inference_cost": 0.00000006, + "upstream_inference_prompt_cost": 0.00000006, + "upstream_inference_completions_cost": 0, + } + assert ( + response_cost_calculator( + response_object=response, + model="openai/text-embedding-3-small", + custom_llm_provider="openrouter", + call_type="embedding", + optional_params={}, + ) + == 0.00000006 + ) + + +def test_openrouter_embedding_accepts_usage_without_cost(): + config = OpenrouterEmbeddingConfig() + raw_response = httpx.Response( + status_code=200, + json={ + "data": [{"embedding": [0.1, 0.2], "index": 0, "object": "embedding"}], + "model": "text-embedding-3-small", + "object": "list", + "usage": {"prompt_tokens": 3, "total_tokens": 3}, + }, + ) + + response = config.transform_embedding_response( + model="openrouter/openai/text-embedding-3-small", + raw_response=raw_response, + model_response=EmbeddingResponse(), + logging_obj=Mock(), + api_key="test-api-key", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + assert "llm_provider-x-litellm-response-cost" not in response._hidden_params.get("additional_headers", {}) + assert "response_cost_details" not in response._hidden_params + + +def test_openrouter_embedding_accepts_partial_cost_details(): + config = OpenrouterEmbeddingConfig() + raw_response = httpx.Response( + status_code=200, + json={ + "data": [{"embedding": [0.1, 0.2], "index": 0, "object": "embedding"}], + "model": "text-embedding-3-small", + "object": "list", + "usage": { + "prompt_tokens": 3, + "total_tokens": 3, + "cost": 0.00000006, + "cost_details": {"upstream_inference_cost": 0.00000005}, + }, + }, + ) + + response = config.transform_embedding_response( + model="openrouter/openai/text-embedding-3-small", + raw_response=raw_response, + model_response=EmbeddingResponse(), + logging_obj=Mock(), + api_key="test-api-key", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] == 0.00000006 + assert response._hidden_params["response_cost_details"] == {"upstream_inference_cost": 0.00000005}