diff --git a/litellm/llms/deepinfra/chat/transformation.py b/litellm/llms/deepinfra/chat/transformation.py index 311c17d49f6..690fead42d1 100644 --- a/litellm/llms/deepinfra/chat/transformation.py +++ b/litellm/llms/deepinfra/chat/transformation.py @@ -1,12 +1,22 @@ import json from collections.abc import Coroutine -from typing import Any, Final, Literal, cast, overload +from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload + +import httpx import litellm from litellm.constants import MIN_NON_ZERO_TEMPERATURE from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import ModelResponse + +if TYPE_CHECKING: + import tiktoken + + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) class DeepInfraConfig(OpenAIGPTConfig): @@ -197,3 +207,62 @@ class DeepInfraConfig(OpenAIGPTConfig): api_base = api_base or get_secret_str("DEEPINFRA_API_BASE") or "https://api.deepinfra.com/v1/openai" dynamic_api_key: Final = api_key or get_secret_str("DEEPINFRA_API_KEY") return api_base, dynamic_api_key + + def transform_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ModelResponse, + logging_obj: "LiteLLMLoggingObj", + request_data: dict, + messages: list[AllMessageValues], + optional_params: dict, + litellm_params: dict, + encoding: "tiktoken.Encoding | None", + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ModelResponse: + """ + Transform the response from DeepInfra. + + DeepInfra reports the charge it actually applied as + ``usage.estimated_cost``, which already reflects the delivered + ``service_tier``. Pricing from the static map instead under-reports + priority-tier traffic, so the reported value is treated as + authoritative, matching how OpenRouter's ``usage.cost`` is handled. + """ + model_response = super().transform_response( + model=model, + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data=request_data, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + encoding=encoding, + api_key=api_key, + json_mode=json_mode, + ) + + try: + response_json: Final = raw_response.json() + usage: Final = response_json.get("usage") or {} + reported_cost: Final = usage.get("estimated_cost") + service_tier: Final = response_json.get("service_tier") + + if reported_cost is not None: + if not hasattr(model_response, "_hidden_params"): + model_response._hidden_params = {} + if "additional_headers" not in model_response._hidden_params: + model_response._hidden_params["additional_headers"] = {} + model_response._hidden_params["additional_headers"][ + "llm_provider-x-litellm-response-cost" + ] = float(reported_cost) + + if service_tier is not None: + setattr(model_response, "service_tier", service_tier) + except Exception: + pass + + return model_response diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_chat_transformation.py b/tests/test_litellm/llms/deepinfra/test_deepinfra_chat_transformation.py index 0865a14fbd9..0d617f13baa 100644 --- a/tests/test_litellm/llms/deepinfra/test_deepinfra_chat_transformation.py +++ b/tests/test_litellm/llms/deepinfra/test_deepinfra_chat_transformation.py @@ -200,3 +200,94 @@ async def test_deepinfra_tool_message_content_transformation_async(): print(f"āœ“ Async test passed: {tool_message['content']}") print("\nāœ… DeepInfra async tool message transformation test passed!") + + +def test_provider_reported_cost_is_authoritative(): + from litellm.llms.deepinfra.chat.transformation import DeepInfraConfig + from litellm.types.utils import ModelResponse + + body = { + "usage": { + "prompt_tokens": 11, + "completion_tokens": 5, + "total_tokens": 16, + "estimated_cost": 8.775e-06, + }, + "service_tier": "priority", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "id": "x", + "created": 1, + "model": "deepinfra/Qwen", + "object": "chat.completion", + } + + raw = MagicMock() + raw.json.return_value = body + raw.text = json.dumps(body) + raw.headers = {} + raw.status_code = 200 + + response = DeepInfraConfig().transform_response( + model="deepinfra/Qwen", + raw_response=raw, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + headers = response._hidden_params["additional_headers"] + assert headers["llm_provider-x-litellm-response-cost"] == 8.775e-06 + assert response.service_tier == "priority" + + +def test_missing_reported_cost_is_not_fatal(): + from litellm.llms.deepinfra.chat.transformation import DeepInfraConfig + from litellm.types.utils import ModelResponse + + body = { + "usage": {"prompt_tokens": 11, "completion_tokens": 5, "total_tokens": 16}, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "id": "x", + "created": 1, + "model": "deepinfra/Qwen", + "object": "chat.completion", + } + + raw = MagicMock() + raw.json.return_value = body + raw.text = json.dumps(body) + raw.headers = {} + raw.status_code = 200 + + response = DeepInfraConfig().transform_response( + model="deepinfra/Qwen", + raw_response=raw, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + headers = (getattr(response, "_hidden_params", {}) or {}).get( + "additional_headers", {} + ) + assert "llm_provider-x-litellm-response-cost" not in headers