Merge 05a6e2246b into 22a349ee70

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:
         """

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}