fix: handle null daily activity token rollups

This commit is contained in:
oppnc 2026-05-26 20:53:07 +08:00
parent 96a2e8b16d
commit 2626e727a3
2 changed files with 99 additions and 9 deletions

View file

@ -696,16 +696,35 @@ _GROUP_DATE_ENDPOINT_API_KEY = 30 # 0b0011110
def _record_to_spend_metrics(record: Any) -> SpendMetrics:
"""Build a SpendMetrics directly from one already-aggregated rollup row."""
spend = 0 if record.spend is None else record.spend
prompt_tokens = 0 if record.prompt_tokens is None else record.prompt_tokens
completion_tokens = (
0 if record.completion_tokens is None else record.completion_tokens
)
cache_read_input_tokens = (
0 if record.cache_read_input_tokens is None else record.cache_read_input_tokens
)
cache_creation_input_tokens = (
0
if record.cache_creation_input_tokens is None
else record.cache_creation_input_tokens
)
api_requests = 0 if record.api_requests is None else record.api_requests
successful_requests = (
0 if record.successful_requests is None else record.successful_requests
)
failed_requests = 0 if record.failed_requests is None else record.failed_requests
return SpendMetrics(
spend=record.spend,
prompt_tokens=record.prompt_tokens,
completion_tokens=record.completion_tokens,
total_tokens=record.prompt_tokens + record.completion_tokens,
cache_read_input_tokens=record.cache_read_input_tokens,
cache_creation_input_tokens=record.cache_creation_input_tokens,
api_requests=record.api_requests,
successful_requests=record.successful_requests,
failed_requests=record.failed_requests,
spend=spend,
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
cache_read_input_tokens=cache_read_input_tokens,
cache_creation_input_tokens=cache_creation_input_tokens,
api_requests=api_requests,
successful_requests=successful_requests,
failed_requests=failed_requests,
)

View file

@ -230,6 +230,77 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown():
mock_prisma.db.query_raw.assert_called_once()
@pytest.mark.asyncio
async def test_get_daily_activity_aggregated_handles_null_token_rollups():
"""Test that aggregated activity treats NULL token sums as zero."""
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
base = {
"api_key": None,
"model": None,
"model_group": None,
"custom_llm_provider": None,
"mcp_namespaced_tool_name": None,
"cache_read_input_tokens": None,
"cache_creation_input_tokens": None,
"api_requests": 1,
"successful_requests": 1,
"failed_requests": 0,
"spend": 0.0,
"prompt_tokens": None,
"completion_tokens": None,
}
mock_rows = [
{
**base,
"date": "2024-01-01",
"endpoint": "/v1/chat/completions",
"group_level": 62,
},
{
**base,
"date": "2024-01-01",
"endpoint": None,
"group_level": 63,
},
{
**base,
"date": None,
"endpoint": None,
"group_level": 127,
},
]
mock_prisma.db.query_raw = AsyncMock(return_value=mock_rows)
result = await get_daily_activity_aggregated(
prisma_client=mock_prisma,
table_name="litellm_dailyuserspend",
entity_id_field="user_id",
entity_id=None,
entity_metadata_field=None,
start_date="2024-01-01",
end_date="2024-01-01",
model=None,
api_key=None,
)
assert result.metadata.total_prompt_tokens == 0
assert result.metadata.total_completion_tokens == 0
assert result.metadata.total_tokens == 0
daily_data = result.results[0]
assert daily_data.metrics.prompt_tokens == 0
assert daily_data.metrics.completion_tokens == 0
assert daily_data.metrics.total_tokens == 0
endpoint = daily_data.breakdown.endpoints["/v1/chat/completions"]
assert endpoint.metrics.prompt_tokens == 0
assert endpoint.metrics.completion_tokens == 0
assert endpoint.metrics.total_tokens == 0
@pytest.mark.asyncio
async def test_get_api_key_metadata_returns_active_key_metadata():
"""Test that get_api_key_metadata should return metadata for active keys."""