fix(proxy): log spend for OpenAI passthrough embeddings with unmapped models

This commit is contained in:
mateo-berri 2026-08-18 19:48:02 -07:00
parent ec6f1c1a56
commit 81914ebc31
2 changed files with 96 additions and 3 deletions

View file

@ -229,6 +229,25 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
verbose_proxy_logger.warning("Error calculating image editing cost: %s", e)
return 0.0
@staticmethod
def _calculate_embeddings_cost(
litellm_model_response: EmbeddingResponse,
model: str,
custom_llm_provider: str,
) -> float:
try:
return litellm.completion_cost(
completion_response=litellm_model_response,
model=model,
custom_llm_provider=custom_llm_provider,
call_type="aembedding",
)
except Exception as e: # noqa: BLE001 # completion_cost raises bare Exception for unmapped models; cost failure must never drop the spend log
verbose_proxy_logger.warning(
"Error calculating embeddings cost for model %s, logging spend with cost 0: %s", model, e
)
return 0.0
@staticmethod
def _build_responses_api_response_and_cost(
model: str,
@ -351,11 +370,10 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
model_response_object=EmbeddingResponse(),
response_type="embedding",
)
response_cost = litellm.completion_cost(
completion_response=litellm_model_response,
response_cost = OpenAIPassthroughLoggingHandler._calculate_embeddings_cost(
litellm_model_response=litellm_model_response,
model=model,
custom_llm_provider=custom_llm_provider,
call_type="aembedding",
)
litellm_model_response._hidden_params["response_cost"] = response_cost
elif is_image_generation:
@ -471,6 +489,12 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
except Exception as e:
verbose_proxy_logger.error("Error in OpenAI passthrough cost tracking: %s", e)
if not is_chat_completions:
unbilled_result: Final[PassThroughEndpointLoggingTypedDict] = {
"result": None,
"kwargs": kwargs,
}
return unbilled_result
# Fall back to base handler without cost tracking
base_handler = OpenAIPassthroughLoggingHandler()
return base_handler.passthrough_chat_handler(

View file

@ -1286,6 +1286,75 @@ class TestOpenAIPassthroughIntegration:
mock_chat_handler.assert_called_once()
assert result == {"result": None, "kwargs": {}}
def test_openai_passthrough_handler_embeddings_unmapped_model_logs_zero_cost(self):
response_body = {
"object": "list",
"model": "lit5787-unmapped-embeddings-deployment",
"data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}],
"usage": {"prompt_tokens": 9, "total_tokens": 9},
}
mock_logging_obj = self._create_mock_logging_obj()
result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler(
httpx_response=self._create_mock_httpx_response(response_body),
response_body=response_body,
logging_obj=mock_logging_obj,
url_route="https://my-resource.openai.azure.com/openai/v1/embeddings",
result="",
start_time=self.start_time,
end_time=self.end_time,
cache_hit=False,
request_body={
"model": "lit5787-unmapped-embeddings-deployment",
"input": "spend probe",
},
passthrough_logging_payload=PassthroughStandardLoggingPayload(
url="https://my-resource.openai.azure.com/openai/v1/embeddings",
request_body={
"model": "lit5787-unmapped-embeddings-deployment",
"input": "spend probe",
},
request_method="POST",
),
litellm_params={},
)
assert result["result"] is not None
assert result["result"].usage.prompt_tokens == 9
assert result["kwargs"]["response_cost"] == 0.0
assert result["kwargs"]["model"] == "lit5787-unmapped-embeddings-deployment"
assert result["result"]._hidden_params["response_cost"] == 0.0
assert mock_logging_obj.model_call_details["response_cost"] == 0.0
def test_openai_passthrough_handler_embeddings_error_skips_chat_fallback(self):
response_body = {
"object": "list",
"model": "text-embedding-3-small",
"usage": {"prompt_tokens": 9, "total_tokens": 9},
}
kwargs_in = {
"passthrough_logging_payload": PassthroughStandardLoggingPayload(
url="https://api.openai.com/v1/embeddings",
request_body={"model": "text-embedding-3-small", "input": "spend probe"},
request_method="POST",
),
"litellm_params": {},
}
result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler(
httpx_response=self._create_mock_httpx_response(response_body),
response_body=response_body,
logging_obj=self._create_mock_logging_obj(),
url_route="https://api.openai.com/v1/embeddings",
result="",
start_time=self.start_time,
end_time=self.end_time,
cache_hit=False,
request_body={"model": "text-embedding-3-small", "input": "spend probe"},
**kwargs_in,
)
assert result["result"] is None
assert result["kwargs"]["passthrough_logging_payload"] == kwargs_in["passthrough_logging_payload"]
@patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.openai_passthrough_handler"
)