scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 178

feat(spend-tracking): show actual model used by azure_ai/model_router on usage page

Browse source
This commit is contained in:
Sameer Kankute 2026-03-20 16:22:40 +05:30
parent d7c419bfee
commit a7c2d25dcb
4 changed files with 504 additions and 77 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/proxy/_types.py
View file

@ -3255,6 +3255,9 @@ class SpendLogsMetadata(TypedDict):
cost_breakdown: Optional[
CostBreakdown
] # Detailed cost breakdown (input_cost, output_cost, margin, discount, etc.)
response_model: Optional[
str
] # Actual model used by model router (e.g., azure_ai/gpt-oss-120b)
class SpendLogsPayload(TypedDict):

263
litellm/proxy/db/db_spend_update_writer.py
View file

@ -20,6 +20,7 @@ from typing import (
List,
Literal,
Optional,
Tuple,
Union,
cast,
overload,
@ -30,6 +31,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.caching import DualCache, RedisCache
from litellm.constants import DB_SPEND_UPDATE_JOB_NAME
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.llms.azure_ai.cost_calculator import _is_azure_model_router
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
BaseDailySpendTransaction,
@ -306,6 +308,147 @@ class DBSpendUpdateWriter:
"_enqueue_tool_registry_upsert error (non-blocking): %s", e
)
@staticmethod
def _split_model_router_payload(
payload: SpendLogsPayload,
) -> Optional[Tuple[SpendLogsPayload, SpendLogsPayload]]:
"""
Split a model router payload into two: one for the router and one for the actual model.
Returns (router_payload, actual_model_payload) or None if split is not possible.
router_payload: model=original router model, spend=router flat cost, tokens=0
actual_model_payload: model=actual model from response, spend=base model cost, tokens=original
"""
metadata_str = payload.get("metadata")
if not metadata_str:
return None
metadata: SpendLogsMetadata = json.loads(metadata_str)
response_model = metadata.get("response_model")
if not response_model:
return None
cost_breakdown = metadata.get("cost_breakdown")
router_flat_cost = 0.0
if cost_breakdown:
additional_costs = cost_breakdown.get("additional_costs") or {}
router_flat_cost = additional_costs.get("azure_model_router_flat_cost", 0.0)
total_spend = payload.get("spend", 0) or 0
base_model_cost = total_spend - router_flat_cost
# Router entry: flat cost only, no tokens
router_payload = cast(
SpendLogsPayload,
{
**payload,
"spend": router_flat_cost,
"prompt_tokens": 0,
"completion_tokens": 0,
"total_tokens": 0,
},
)
# Actual model entry: base model cost, original tokens
# Ensure the model has the provider prefix
actual_model_name = response_model
custom_llm_provider = payload.get("custom_llm_provider", "")
if custom_llm_provider and "/" not in actual_model_name:
actual_model_name = f"{custom_llm_provider}/{actual_model_name}"
actual_payload = cast(
SpendLogsPayload,
{
**payload,
"model": actual_model_name,
"spend": base_model_cost,
},
)
return router_payload, actual_payload
async def _enqueue_all_daily_transactions(
self,
payload: SpendLogsPayload,
prisma_client: Optional[PrismaClient],
org_id: Optional[str],
api_requests: Optional[int] = None,
) -> None:
"""Enqueue daily spend transactions for all 6 entity types."""
try:
await self.add_spend_log_transaction_to_daily_user_transaction(
payload=payload,
prisma_client=prisma_client,
api_requests=api_requests,
)
except Exception:
verbose_proxy_logger.debug(
"_enqueue_all_daily_transactions: daily_user failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_end_user_transaction(
payload=payload,
prisma_client=prisma_client,
api_requests=api_requests,
)
except Exception:
verbose_proxy_logger.debug(
"_enqueue_all_daily_transactions: daily_end_user failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_agent_transaction(
payload=payload,
prisma_client=prisma_client,
api_requests=api_requests,
)
except Exception:
verbose_proxy_logger.debug(
"_enqueue_all_daily_transactions: daily_agent failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_team_transaction(
payload=payload,
prisma_client=prisma_client,
api_requests=api_requests,
)
except Exception:
verbose_proxy_logger.debug(
"_enqueue_all_daily_transactions: daily_team failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_org_transaction(
payload=payload,
org_id=org_id,
prisma_client=prisma_client,
api_requests=api_requests,
)
except Exception:
verbose_proxy_logger.debug(
"_enqueue_all_daily_transactions: daily_org failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_tag_transaction(
payload=payload,
prisma_client=prisma_client,
api_requests=api_requests,
)
except Exception:
verbose_proxy_logger.debug(
"_enqueue_all_daily_transactions: daily_tag failed: %s",
traceback.format_exc(),
)
async def _batch_database_updates(
self,
*,
@ -403,71 +546,27 @@ class DBSpendUpdateWriter:
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_user_transaction(
payload=payload_copy,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: add_spend_log_transaction_to_daily_user_transaction failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_end_user_transaction(
payload=payload_copy,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: add_spend_log_transaction_to_daily_end_user_transaction failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_agent_transaction(
payload=payload_copy,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: add_spend_log_transaction_to_daily_agent_transaction failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_team_transaction(
payload=payload_copy,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: add_spend_log_transaction_to_daily_team_transaction failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_org_transaction(
payload=payload_copy,
org_id=org_id,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: add_spend_log_transaction_to_daily_org_transaction failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_tag_transaction(
payload=payload_copy,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: add_spend_log_transaction_to_daily_tag_transaction failed: %s",
traceback.format_exc(),
# For Azure AI Model Router requests, split into two daily spend entries:
# one for the router (flat cost, no tokens) and one for the actual model (base cost, tokens).
# This ensures both models appear on the usage page without double-counting.
if _is_azure_model_router(payload_copy.get("model", "") or ""):
split_result = self._split_model_router_payload(payload_copy)
if split_result is not None:
router_payload, actual_payload = split_result
await self._enqueue_all_daily_transactions(
router_payload, prisma_client, org_id
)
await self._enqueue_all_daily_transactions(
actual_payload, prisma_client, org_id, api_requests=0
)
else:
# Split not possible (missing metadata), fall back to single entry
await self._enqueue_all_daily_transactions(
payload_copy, prisma_client, org_id
)
else:
await self._enqueue_all_daily_transactions(
payload_copy, prisma_client, org_id
)
async def _update_key_db(
@ -1859,6 +1958,7 @@ class DBSpendUpdateWriter:
type: Literal[
"user", "team", "org", "request_tags", "end_user", "agent"
] = "user",
api_requests: Optional[int] = None,
) -> Optional[BaseDailySpendTransaction]:
common_expected_keys = ["startTime", "api_key"]
if type == "user":
@ -1918,6 +2018,8 @@ class DBSpendUpdateWriter:
if call_type:
endpoint = ROUTE_ENDPOINT_MAPPING.get(call_type, None)
_api_requests = api_requests if api_requests is not None else 1
_counts_as_request = _api_requests > 0
daily_transaction = BaseDailySpendTransaction(
date=date,
api_key=payload["api_key"],
@ -1929,9 +2031,13 @@ class DBSpendUpdateWriter:
prompt_tokens=payload["prompt_tokens"],
completion_tokens=payload["completion_tokens"],
spend=payload["spend"],
api_requests=1,
successful_requests=1 if request_status == "success" else 0,
failed_requests=1 if request_status != "success" else 0,
api_requests=_api_requests,
successful_requests=(
1 if _counts_as_request and request_status == "success" else 0
),
failed_requests=(
1 if _counts_as_request and request_status != "success" else 0
),
cache_read_input_tokens=usage_obj.get("cache_read_input_tokens", 0)
or 0,
cache_creation_input_tokens=usage_obj.get(
@ -1947,6 +2053,7 @@ class DBSpendUpdateWriter:
self,
payload: Union[dict, SpendLogsPayload],
prisma_client: Optional[PrismaClient] = None,
api_requests: Optional[int] = None,
):
"""
Add a spend log transaction to the `daily_spend_update_queue`
@ -1963,7 +2070,7 @@ class DBSpendUpdateWriter:
base_daily_transaction = (
await self._common_add_spend_log_transaction_to_daily_transaction(
payload, prisma_client, "user"
payload, prisma_client, "user", api_requests=api_requests
)
)
if base_daily_transaction is None:
@ -1982,6 +2089,7 @@ class DBSpendUpdateWriter:
self,
payload: SpendLogsPayload,
prisma_client: Optional[PrismaClient] = None,
api_requests: Optional[int] = None,
) -> None:
if prisma_client is None:
verbose_proxy_logger.debug(
@ -1991,7 +2099,7 @@ class DBSpendUpdateWriter:
base_daily_transaction = (
await self._common_add_spend_log_transaction_to_daily_transaction(
payload, prisma_client, "team"
payload, prisma_client, "team", api_requests=api_requests
)
)
if base_daily_transaction is None:
@ -2016,6 +2124,7 @@ class DBSpendUpdateWriter:
payload: SpendLogsPayload,
prisma_client: Optional[PrismaClient] = None,
org_id: Optional[str] = None,
api_requests: Optional[int] = None,
) -> None:
if prisma_client is None:
verbose_proxy_logger.debug(
@ -2039,7 +2148,7 @@ class DBSpendUpdateWriter:
base_daily_transaction = (
await self._common_add_spend_log_transaction_to_daily_transaction(
payload_with_org, prisma_client, "org"
payload_with_org, prisma_client, "org", api_requests=api_requests
)
)
if base_daily_transaction is None:
@ -2058,6 +2167,7 @@ class DBSpendUpdateWriter:
self,
payload: SpendLogsPayload,
prisma_client: Optional[PrismaClient] = None,
api_requests: Optional[int] = None,
) -> None:
if prisma_client is None:
verbose_proxy_logger.debug(
@ -2082,7 +2192,10 @@ class DBSpendUpdateWriter:
base_daily_transaction = (
await self._common_add_spend_log_transaction_to_daily_transaction(
payload_with_end_user_id, prisma_client, "end_user"
payload_with_end_user_id,
prisma_client,
"end_user",
api_requests=api_requests,
)
)
if base_daily_transaction is None:
@ -2101,6 +2214,7 @@ class DBSpendUpdateWriter:
self,
payload: SpendLogsPayload,
prisma_client: Optional[PrismaClient] = None,
api_requests: Optional[int] = None,
) -> None:
if prisma_client is None:
verbose_proxy_logger.debug(
@ -2118,7 +2232,7 @@ class DBSpendUpdateWriter:
)
base_daily_transaction = (
await self._common_add_spend_log_transaction_to_daily_transaction(
payload_with_agent_id, prisma_client, "agent"
payload_with_agent_id, prisma_client, "agent", api_requests=api_requests
)
)
if base_daily_transaction is None:
@ -2136,6 +2250,7 @@ class DBSpendUpdateWriter:
self,
payload: SpendLogsPayload,
prisma_client: Optional[PrismaClient] = None,
api_requests: Optional[int] = None,
) -> None:
if prisma_client is None:
verbose_proxy_logger.debug(
@ -2145,7 +2260,7 @@ class DBSpendUpdateWriter:
base_daily_transaction = (
await self._common_add_spend_log_transaction_to_daily_transaction(
payload, prisma_client, "request_tags"
payload, prisma_client, "request_tags", api_requests=api_requests
)
)
if base_daily_transaction is None:

18
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

@ -18,12 +18,15 @@ from litellm.constants import (
from litellm.constants import (
MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB,
)
from litellm.constants import REDACTED_BY_LITELM_STRING
from litellm.constants import (
REDACTED_BY_LITELM_STRING,
)
from litellm.litellm_core_utils.core_helpers import (
get_litellm_metadata_from_kwargs,
reconstruct_model_name,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.azure_ai.cost_calculator import _is_azure_model_router
from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload
from litellm.proxy.utils import PrismaClient, hash_token
from litellm.types.utils import (
@ -83,6 +86,7 @@ def _get_spend_logs_metadata(
cold_storage_object_key: Optional[str] = None,
litellm_overhead_time_ms: Optional[float] = None,
cost_breakdown: Optional[CostBreakdown] = None,
response_model: Optional[str] = None,
) -> SpendLogsMetadata:
if metadata is None:
return SpendLogsMetadata(
@ -111,6 +115,7 @@ def _get_spend_logs_metadata(
attempted_retries=None,
max_retries=None,
cost_breakdown=None,
response_model=response_model,
)
verbose_proxy_logger.debug(
"getting payload for SpendLogs, available keys in metadata: "
@ -135,6 +140,7 @@ def _get_spend_logs_metadata(
clean_metadata["cold_storage_object_key"] = cold_storage_object_key
clean_metadata["litellm_overhead_time_ms"] = litellm_overhead_time_ms
clean_metadata["cost_breakdown"] = cost_breakdown
clean_metadata["response_model"] = response_model
return clean_metadata
@ -337,6 +343,15 @@ def get_logging_payload( # noqa: PLR0915
hidden_params = standard_logging_payload.get("hidden_params", {})
litellm_overhead_time_ms = hidden_params.get("litellm_overhead_time_ms")
# For Azure AI Model Router, capture the actual model used from the response
_raw_model = cast(str, kwargs.get("model") or "")
response_model: Optional[str] = None
if _raw_model and _is_azure_model_router(_raw_model):
if standard_logging_payload is not None:
response_model = standard_logging_payload.get("model")
if not response_model:
response_model = response_obj_dict.get("model")
# clean up litellm metadata
clean_metadata = _get_spend_logs_metadata(
metadata,
@ -392,6 +407,7 @@ def get_logging_payload( # noqa: PLR0915
if standard_logging_payload is not None
else None
),
response_model=response_model,
)
special_usage_fields = ["completion_tokens", "prompt_tokens", "total_tokens"]

297
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -13,6 +13,7 @@ from unittest.mock import AsyncMock, MagicMock, call, patch
import pytest
from litellm.llms.azure_ai.cost_calculator import _is_azure_model_router
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
@ -1398,8 +1399,8 @@ async def test_commit_spend_updates_uses_pipeline():
mock_redis_update_buffer = AsyncMock()
mock_redis_update_buffer.store_in_memory_spend_updates_in_redis = AsyncMock()
# Return all-None tuple (no data to commit)
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
return_value=(None, None, None, None, None, None, None)
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = (
AsyncMock(return_value=(None, None, None, None, None, None, None))
)
db_writer.redis_update_buffer = mock_redis_update_buffer
@ -1428,3 +1429,295 @@ async def test_commit_spend_updates_uses_pipeline():
mock_redis_update_buffer.get_all_daily_end_user_spend_update_transactions_from_redis_buffer.assert_not_called()
mock_redis_update_buffer.get_all_daily_agent_spend_update_transactions_from_redis_buffer.assert_not_called()
mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer.assert_not_called()
# ===== Azure AI Model Router Dual Spend Tracking Tests =====
def test_is_azure_model_router_detects_model_router():
"""Test _is_azure_model_router correctly identifies model router requests."""
# Should detect model_router
assert _is_azure_model_router("azure_ai/model_router") is True
assert _is_azure_model_router("azure_ai/model-router") is True
assert _is_azure_model_router("azure_ai/model_router/my-deployment") is True
# Should NOT detect non-router models
assert _is_azure_model_router("azure_ai/gpt-4") is False
assert _is_azure_model_router("openai/gpt-4") is False
assert _is_azure_model_router("") is False
def test_split_model_router_payload_splits_correctly():
"""Test _split_model_router_payload correctly splits cost and tokens."""
writer = DBSpendUpdateWriter()
metadata = {
"response_model": "azure_ai/gpt-oss-120b",
"cost_breakdown": {
"additional_costs": {"azure_model_router_flat_cost": 0.00014}
},
}
payload = {
"model": "azure_ai/model_router",
"custom_llm_provider": "azure_ai",
"spend": 0.01014, # 0.01 base + 0.00014 flat
"prompt_tokens": 1000,
"completion_tokens": 500,
"total_tokens": 1500,
"metadata": json.dumps(metadata),
"api_key": "test-key",
"request_id": "req-123",
}
result = writer._split_model_router_payload(payload)
assert result is not None
router_payload, actual_payload = result
# Router payload: flat cost only, no tokens
assert router_payload["model"] == "azure_ai/model_router"
assert router_payload["spend"] == pytest.approx(0.00014)
assert router_payload["prompt_tokens"] == 0
assert router_payload["completion_tokens"] == 0
assert router_payload["total_tokens"] == 0
# Actual model payload: base cost, original tokens
assert actual_payload["model"] == "azure_ai/gpt-oss-120b"
assert actual_payload["spend"] == pytest.approx(0.01)
assert actual_payload["prompt_tokens"] == 1000
assert actual_payload["completion_tokens"] == 500
def test_split_model_router_payload_returns_none_without_response_model():
"""Test _split_model_router_payload returns None when response_model is missing."""
writer = DBSpendUpdateWriter()
# No response_model in metadata
payload = {
"model": "azure_ai/model_router",
"spend": 0.01,
"metadata": json.dumps({}),
}
assert writer._split_model_router_payload(payload) is None
# No metadata at all
payload_no_meta = {
"model": "azure_ai/model_router",
"spend": 0.01,
}
assert writer._split_model_router_payload(payload_no_meta) is None
def test_split_model_router_payload_adds_provider_prefix():
"""Test that response model gets provider prefix if missing."""
writer = DBSpendUpdateWriter()
metadata = {
"response_model": "gpt-oss-120b", # No provider prefix
"cost_breakdown": {
"additional_costs": {"azure_model_router_flat_cost": 0.00014}
},
}
payload = {
"model": "azure_ai/model_router",
"custom_llm_provider": "azure_ai",
"spend": 0.01014,
"prompt_tokens": 100,
"completion_tokens": 50,
"total_tokens": 150,
"metadata": json.dumps(metadata),
}
result = writer._split_model_router_payload(payload)
assert result is not None
_, actual_payload = result
assert actual_payload["model"] == "azure_ai/gpt-oss-120b"
def test_split_model_router_payload_no_cost_breakdown():
"""Test split when cost_breakdown is missing - router flat cost defaults to 0."""
writer = DBSpendUpdateWriter()
metadata = {
"response_model": "azure_ai/gpt-oss-120b",
# No cost_breakdown
}
payload = {
"model": "azure_ai/model_router",
"custom_llm_provider": "azure_ai",
"spend": 0.01,
"prompt_tokens": 100,
"completion_tokens": 50,
"total_tokens": 150,
"metadata": json.dumps(metadata),
}
result = writer._split_model_router_payload(payload)
assert result is not None
router_payload, actual_payload = result
# Router gets 0 cost (no flat cost found)
assert router_payload["spend"] == 0.0
# Actual model gets all cost
assert actual_payload["spend"] == pytest.approx(0.01)
@pytest.mark.asyncio
async def test_common_daily_transaction_respects_api_requests_override():
"""Test that _common_add_spend_log_transaction_to_daily_transaction uses api_requests override."""
writer = DBSpendUpdateWriter()
mock_prisma = MagicMock()
mock_prisma.get_request_status = MagicMock(return_value="success")
payload = {
"user": "test-user",
"startTime": "2024-01-15T10:00:00Z",
"api_key": "test-key",
"model": "azure_ai/gpt-oss-120b",
"custom_llm_provider": "azure_ai",
"model_group": "test-group",
"prompt_tokens": 100,
"completion_tokens": 50,
"spend": 0.01,
"metadata": json.dumps({"usage_object": {}}),
}
# Default api_requests (should be 1)
result = await writer._common_add_spend_log_transaction_to_daily_transaction(
payload, mock_prisma, "user"
)
assert result is not None
assert result["api_requests"] == 1
assert result["successful_requests"] == 1
# Override api_requests to 0
result_zero = await writer._common_add_spend_log_transaction_to_daily_transaction(
payload, mock_prisma, "user", api_requests=0
)
assert result_zero is not None
assert result_zero["api_requests"] == 0
assert result_zero["successful_requests"] == 0
assert result_zero["failed_requests"] == 0
@pytest.mark.asyncio
async def test_batch_database_updates_dual_transactions_for_model_router():
"""Test that _batch_database_updates emits two daily transaction sets for model router."""
writer = DBSpendUpdateWriter()
# Mock the helper to track calls
writer._enqueue_all_daily_transactions = AsyncMock()
writer._update_user_db = AsyncMock()
writer._update_key_db = AsyncMock()
writer._update_team_db = AsyncMock()
writer._update_org_db = AsyncMock()
writer._update_tag_db = AsyncMock()
writer._update_agent_db = AsyncMock()
metadata = {
"response_model": "azure_ai/gpt-oss-120b",
"cost_breakdown": {
"additional_costs": {"azure_model_router_flat_cost": 0.00014}
},
}
payload = {
"model": "azure_ai/model_router",
"custom_llm_provider": "azure_ai",
"spend": 0.01014,
"prompt_tokens": 1000,
"completion_tokens": 500,
"total_tokens": 1500,
"metadata": json.dumps(metadata),
"api_key": "test-key",
"request_id": "req-123",
"user": "test-user",
"team_id": "test-team",
"agent_id": None,
"request_tags": None,
}
mock_prisma = MagicMock()
await writer._batch_database_updates(
response_cost=0.01014,
user_id="test-user",
hashed_token="test-token",
team_id="test-team",
org_id="test-org",
end_user_id=None,
prisma_client=mock_prisma,
user_api_key_cache=MagicMock(),
litellm_proxy_budget_name=None,
payload_copy=payload,
request_tags=None,
)
# Should be called twice: once for router, once for actual model
assert writer._enqueue_all_daily_transactions.call_count == 2
# First call: router payload (api_requests defaults to None → 1)
first_call = writer._enqueue_all_daily_transactions.call_args_list[0]
router_payload = first_call[0][0]
assert router_payload["model"] == "azure_ai/model_router"
assert router_payload["spend"] == pytest.approx(0.00014)
assert router_payload["prompt_tokens"] == 0
assert first_call[1].get("api_requests") is None # default (1)
# Second call: actual model payload (api_requests=0)
second_call = writer._enqueue_all_daily_transactions.call_args_list[1]
actual_payload = second_call[0][0]
assert actual_payload["model"] == "azure_ai/gpt-oss-120b"
assert actual_payload["spend"] == pytest.approx(0.01)
assert actual_payload["prompt_tokens"] == 1000
assert second_call[1].get("api_requests") == 0
@pytest.mark.asyncio
async def test_batch_database_updates_single_transaction_for_non_router():
"""Test that _batch_database_updates emits one daily transaction set for non-router models."""
writer = DBSpendUpdateWriter()
writer._enqueue_all_daily_transactions = AsyncMock()
writer._update_user_db = AsyncMock()
writer._update_key_db = AsyncMock()
writer._update_team_db = AsyncMock()
writer._update_org_db = AsyncMock()
writer._update_tag_db = AsyncMock()
writer._update_agent_db = AsyncMock()
payload = {
"model": "azure_ai/gpt-4",
"custom_llm_provider": "azure_ai",
"spend": 0.05,
"prompt_tokens": 1000,
"completion_tokens": 500,
"total_tokens": 1500,
"metadata": json.dumps({}),
"api_key": "test-key",
"request_id": "req-456",
"user": "test-user",
"team_id": "test-team",
"agent_id": None,
"request_tags": None,
}
mock_prisma = MagicMock()
await writer._batch_database_updates(
response_cost=0.05,
user_id="test-user",
hashed_token="test-token",
team_id="test-team",
org_id=None,
end_user_id=None,
prisma_client=mock_prisma,
user_api_key_cache=MagicMock(),
litellm_proxy_budget_name=None,
payload_copy=payload,
request_tags=None,
)
# Should be called once for non-router model
assert writer._enqueue_all_daily_transactions.call_count == 1
first_call = writer._enqueue_all_daily_transactions.call_args_list[0]
assert first_call[0][0]["model"] == "azure_ai/gpt-4"
assert first_call[1].get("api_requests") is None # default