From 05a6e2246be7d74b0fe7d34c05dccf3655596d6e Mon Sep 17 00:00:00 2001 From: Hector Date: Fri, 14 Aug 2026 12:44:23 +0800 Subject: [PATCH] fix(openrouter): preserve exact embedding costs --- .../openrouter/embedding/transformation.py | 42 ++++++- ...est_openrouter_embedding_transformation.py | 118 ++++++++++++++++++ 2 files changed, 157 insertions(+), 3 deletions(-) diff --git a/litellm/llms/openrouter/embedding/transformation.py b/litellm/llms/openrouter/embedding/transformation.py index 29d0c8c1c56..032cfd65b0f 100644 --- a/litellm/llms/openrouter/embedding/transformation.py +++ b/litellm/llms/openrouter/embedding/transformation.py @@ -10,6 +10,7 @@ Docs: https://openrouter.ai/docs from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.types.llms.openai import AllEmbeddingInputValues @@ -26,6 +27,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 = TypeAdapter(dict[str, JsonValue]) + + class OpenrouterEmbeddingConfig(BaseEmbeddingConfig): """ Configuration for OpenRouter's Embedding API. @@ -136,13 +161,24 @@ 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) + provider_response: Final = _OpenRouterEmbeddingResponse.model_validate(response_json) + provider_cost: Final = provider_response.usage.cost if provider_response.usage is not None else None + cost_details: Final = provider_response.usage.cost_details if provider_response.usage is not None else None + hidden_params: Final = {} # mutable-ok: response conversion attaches headers in place + if provider_cost is not None: + hidden_params["additional_headers"] = { # mutable-ok: the cost calculator reads a header mapping + "llm_provider-x-litellm-response-cost": provider_cost, + } + if cost_details is not None: + hidden_params["response_cost_details"] = cost_details.model_dump(exclude_none=True) + convert_to_model_response_object( response_object=response_json, model_response_object=model_response, response_type="embedding", + hidden_params=hidden_params, ) + return model_response def get_supported_openai_params(self, model: str) -> list: """ diff --git a/tests/test_litellm/llms/openrouter/test_openrouter_embedding_transformation.py b/tests/test_litellm/llms/openrouter/test_openrouter_embedding_transformation.py index a9d4b14ef73..a1b8449fbc4 100644 --- a/tests/test_litellm/llms/openrouter/test_openrouter_embedding_transformation.py +++ b/tests/test_litellm/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}