diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d62bfbb7d50..cb54019267a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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): diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index a305d5be1e6..0caec0a5d1d 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -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: diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 3eacc19a6df..ff8bc361001 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -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"] diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index b5f82ef04c5..ac01c19b633 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -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