test(logging): cover rust spend metadata

This commit is contained in:
Ishaan Jaff 2026-06-25 22:18:04 -07:00
parent 82d0a9f058
commit a9ec5637a8
No known key found for this signature in database

View file

@ -26,6 +26,7 @@ import pytest
import litellm
from litellm.proxy.spend_tracking.spend_tracking_utils import (
_get_spend_logs_metadata,
get_logging_payload,
_sanitize_request_body_for_spend_logs_payload,
)
@ -234,6 +235,58 @@ def test_spend_logs_payload(model_id: Optional[str]):
assert payload["custom_llm_provider"] == "azure"
def test_spend_logs_payload_includes_litellm_rust_metadata():
input_args: dict = {
"kwargs": {
"model": "mistral/mistral-ocr-latest",
"litellm_params": {
"metadata": {
"user_api_key": "sk-test-key",
}
},
"standard_logging_object": {
"request_tags": [],
"metadata": {
"user_api_key_end_user_id": "test-user",
},
"model_map_information": {
"model_map_key": "mistral/mistral-ocr-latest",
"model_map_value": None,
},
"hidden_params": {
"litellm_rust": True,
},
},
},
"response_obj": litellm.ModelResponse(
id="ocr-response-id",
choices=[
litellm.Choices(
finish_reason="stop",
index=0,
message=litellm.Message(content="OCR done", role="assistant"),
)
],
model="mistral/mistral-ocr-latest",
usage=litellm.Usage(completion_tokens=0, prompt_tokens=0, total_tokens=0),
),
"start_time": datetime.datetime.now(),
"end_time": datetime.datetime.now(),
}
payload: SpendLogsPayload = get_logging_payload(**input_args)
assert isinstance(payload["metadata"], str)
metadata = json.loads(payload["metadata"])
assert metadata["litellm_rust"] is True
def test_spend_logs_metadata_includes_litellm_rust_when_metadata_is_none():
metadata = _get_spend_logs_metadata(metadata=None, litellm_rust=True)
assert metadata["litellm_rust"] is True
def test_spend_logs_payload_whisper():
"""
Ensure we can write /transcription request/responses to spend logs