diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b47b43411c5..d796e389802 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4505,6 +4505,7 @@ class DBSpendUpdateTransactions(TypedDict): team_list_transactions: Optional[Dict[str, float]] team_member_list_transactions: Optional[Dict[str, float]] org_list_transactions: Optional[Dict[str, float]] + project_list_transactions: Dict[str, float] | None tag_list_transactions: Optional[Dict[str, float]] agent_list_transactions: Optional[Dict[str, float]] diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 54a4c2dad91..40ad194ba8d 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -131,6 +131,7 @@ class DBSpendUpdateWriter: start_time: Optional[datetime], end_time: Optional[datetime], response_cost: Optional[float], + project_id: str | None = None, ): from litellm.proxy.proxy_server import ( disable_spend_logs, @@ -191,6 +192,7 @@ class DBSpendUpdateWriter: hashed_token=hashed_token, team_id=team_id, org_id=org_id, + project_id=project_id, end_user_id=end_user_id, prisma_client=prisma_client, litellm_proxy_budget_name=litellm_proxy_budget_name, @@ -210,12 +212,13 @@ class DBSpendUpdateWriter: spend_log_error( "Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue " "may not have completed for this request. " - "response_cost=%s, token=%s, user_id=%s, team_id=%s, org_id=%s, end_user_id=%s", + "response_cost=%s, token=%s, user_id=%s, team_id=%s, org_id=%s, project_id=%s, end_user_id=%s", response_cost, token, user_id, team_id, org_id, + project_id, end_user_id, ) @@ -325,6 +328,7 @@ class DBSpendUpdateWriter: prisma_client: Optional[PrismaClient], litellm_proxy_budget_name: Optional[str], payload: SpendLogsPayload, + project_id: str | None = None, ): """ Runs all 11 spend-update helpers sequentially inside a single asyncio task. @@ -387,6 +391,12 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) + await self._update_project_db( + response_cost=response_cost, + project_id=project_id, + prisma_client=prisma_client, + ) + try: await self._update_tag_db( response_cost=response_cost, @@ -631,6 +641,35 @@ class DBSpendUpdateWriter: ) raise e + async def _update_project_db( + self, + response_cost: float | None, + project_id: str | None, + prisma_client: PrismaClient | None, + ): + try: + if project_id is None or prisma_client is None: + verbose_proxy_logger.debug( + "track_cost_callback: project_id is None or prisma_client is None. Not tracking spend for project" + ) + return + + await self.spend_update_queue.add_update( + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.PROJECT, + entity_id=project_id, + response_cost=response_cost, + ) + ) + except Exception as e: # noqa: BLE001 # background spend enqueue must not break request flow or block sibling helpers + spend_log_error( + "Spend tracking - failed to enqueue project spend update. project_id=%s, response_cost=%s - %s", + project_id, + response_cost, + str(e), + exc=e, + ) + async def _update_agent_db( self, response_cost: Optional[float], @@ -1278,6 +1317,18 @@ class DBSpendUpdateWriter: e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj ) + ### UPDATE PROJECT TABLE ### + project_list_transactions = db_spend_update_transactions.get("project_list_transactions") + await DBSpendUpdateWriter._update_entity_spend_in_db( + entity_name="Project", + transactions=project_list_transactions, + table_accessor="litellm_projecttable", + where_field="project_id", + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + ### UPDATE TAG TABLE ### tag_list_transactions = db_spend_update_transactions["tag_list_transactions"] await DBSpendUpdateWriter._update_entity_spend_in_db( diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index c924448669d..6c1a4f97906 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -339,6 +339,10 @@ class RedisUpdateBuffer: Litellm_EntityType.ORGANIZATION, db_spend_update_transactions.get("org_list_transactions"), ), + ( + Litellm_EntityType.PROJECT, + db_spend_update_transactions.get("project_list_transactions"), + ), ( Litellm_EntityType.TAG, db_spend_update_transactions.get("tag_list_transactions"), @@ -717,6 +721,7 @@ class RedisUpdateBuffer: team_list_transactions={}, team_member_list_transactions={}, org_list_transactions={}, + project_list_transactions={}, tag_list_transactions={}, agent_list_transactions={}, ) @@ -729,6 +734,7 @@ class RedisUpdateBuffer: "team_list_transactions", "team_member_list_transactions", "org_list_transactions", + "project_list_transactions", "tag_list_transactions", "agent_list_transactions", ] diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index 0689fc00b02..5a55a9729cb 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -138,6 +138,7 @@ class SpendUpdateQueue(BaseUpdateQueue): team_list_transactions={}, team_member_list_transactions={}, org_list_transactions={}, + project_list_transactions={}, tag_list_transactions={}, agent_list_transactions={}, ) @@ -150,6 +151,7 @@ class SpendUpdateQueue(BaseUpdateQueue): Litellm_EntityType.TEAM: "team_list_transactions", Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions", Litellm_EntityType.ORGANIZATION: "org_list_transactions", + Litellm_EntityType.PROJECT: "project_list_transactions", Litellm_EntityType.TAG: "tag_list_transactions", Litellm_EntityType.AGENT: "agent_list_transactions", } @@ -187,6 +189,8 @@ class SpendUpdateQueue(BaseUpdateQueue): transactions_dict = db_spend_update_transactions["team_member_list_transactions"] elif dict_key == "org_list_transactions": transactions_dict = db_spend_update_transactions["org_list_transactions"] + elif dict_key == "project_list_transactions": + transactions_dict = db_spend_update_transactions["project_list_transactions"] elif dict_key == "tag_list_transactions": transactions_dict = db_spend_update_transactions["tag_list_transactions"] elif dict_key == "agent_list_transactions": diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index b839426fcda..8a80adebbaf 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -174,6 +174,7 @@ class _ProxyDBLogger(CustomLogger): start_time=actual_start_time, end_time=datetime.now(), org_id=user_api_key_dict.org_id, + project_id=user_api_key_dict.project_id, ) @log_db_metrics @@ -208,6 +209,7 @@ class _ProxyDBLogger(CustomLogger): user_id = cast(Optional[str], metadata.get("user_api_key_user_id", None)) team_id = cast(Optional[str], metadata.get("user_api_key_team_id", None)) org_id = cast(Optional[str], metadata.get("user_api_key_org_id", None)) + project_id = cast("str | None", metadata.get("user_api_key_project_id", None)) # cast-ok: untyped metadata key_alias = cast(Optional[str], metadata.get("user_api_key_alias", None)) end_user_max_budget = metadata.get("user_api_end_user_max_budget", None) sl_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) @@ -245,6 +247,7 @@ class _ProxyDBLogger(CustomLogger): end_user_id=end_user_id, team_id=team_id, org_id=org_id, + project_id=project_id, kwargs=kwargs, completion_response=completion_response, start_time=start_time, @@ -488,6 +491,7 @@ async def _update_database_and_spend_counters( response_cost: float, budget_reservation: Optional[dict], request_tags: Optional[List[str]] = None, + project_id: str | None = None, ) -> None: try: await proxy_logging_obj.db_spend_update_writer.update_database( @@ -501,6 +505,7 @@ async def _update_database_and_spend_counters( start_time=start_time, end_time=end_time, org_id=org_id, + project_id=project_id, ) except Exception: if budget_reservation is not None: diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py index 0ed5940dd75..ad90bdf48cd 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py @@ -34,6 +34,31 @@ async def test_add_update(spend_queue): assert spend_queue.update_queue.qsize() == 1 +@pytest.mark.asyncio +async def test_project_entity_type_aggregates_into_project_transactions(spend_queue): + # PROJECT updates must aggregate into project_list_transactions so project spend is committed. + await spend_queue.add_update( + { + "entity_type": Litellm_EntityType.PROJECT, + "entity_id": "project123", + "response_cost": 0.5, + } + ) + await spend_queue.add_update( + { + "entity_type": Litellm_EntityType.PROJECT, + "entity_id": "project123", + "response_cost": 0.25, + } + ) + + aggregated = ( + await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() + ) + + assert aggregated["project_list_transactions"]["project123"] == 0.75 + + @pytest.mark.asyncio async def test_missing_response_cost(spend_queue): # Test with missing response_cost - should default to 0 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 4c17c5d3482..e3bdb2ea3ea 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 @@ -627,6 +627,161 @@ async def test_update_agent_db_skips_when_prisma_client_none(): writer.spend_update_queue.add_update.assert_not_called() +@pytest.mark.asyncio +async def test_update_project_db_enqueues_project_spend(): + """ + Test that _update_project_db enqueues a SpendUpdateQueueItem with entity_type=PROJECT. + """ + from litellm.proxy._types import Litellm_EntityType + + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + project_id = "project-123" + response_cost = 0.1 + + writer.spend_update_queue.add_update = AsyncMock() + + await writer._update_project_db( + response_cost=response_cost, + project_id=project_id, + prisma_client=mock_prisma, + ) + + writer.spend_update_queue.add_update.assert_called_once() + call_args = writer.spend_update_queue.add_update.call_args[1] + assert call_args["update"]["entity_type"] == Litellm_EntityType.PROJECT + assert call_args["update"]["entity_id"] == project_id + assert call_args["update"]["response_cost"] == response_cost + + +@pytest.mark.asyncio +async def test_update_project_db_skips_when_project_id_none(): + """_update_project_db does not enqueue when project_id is None.""" + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + writer.spend_update_queue.add_update = AsyncMock() + + await writer._update_project_db( + response_cost=0.05, + project_id=None, + prisma_client=mock_prisma, + ) + + writer.spend_update_queue.add_update.assert_not_called() + + +@pytest.mark.asyncio +async def test_update_project_db_skips_when_prisma_client_none(): + """_update_project_db does not enqueue when prisma_client is None.""" + writer = DBSpendUpdateWriter() + writer.spend_update_queue.add_update = AsyncMock() + + await writer._update_project_db( + response_cost=0.05, + project_id="project-456", + prisma_client=None, + ) + + writer.spend_update_queue.add_update.assert_not_called() + + +@pytest.mark.asyncio +async def test_batch_database_updates_forwards_project_id_to_project_helper(): + """ + _batch_database_updates must route the request's project_id into _update_project_db + so project spend is attributed alongside user/key/team/org. + """ + db_writer = DBSpendUpdateWriter() + + db_writer._update_user_db = AsyncMock() + db_writer._update_key_db = AsyncMock() + db_writer._update_team_db = AsyncMock() + db_writer._update_org_db = AsyncMock() + db_writer._update_project_db = AsyncMock() + db_writer._update_tag_db = AsyncMock() + db_writer._update_agent_db = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock() + + await db_writer._batch_database_updates( + response_cost=0.1, + user_id="u1", + hashed_token="t1", + team_id="team1", + org_id="org1", + project_id="project1", + end_user_id="eu1", + prisma_client=MagicMock(), + litellm_proxy_budget_name="budget", + payload={"key": "value"}, + ) + + db_writer._update_project_db.assert_awaited_once() + call_args = db_writer._update_project_db.call_args[1] + assert call_args["project_id"] == "project1" + assert call_args["response_cost"] == 0.1 + + +@pytest.mark.asyncio +async def test_commit_spend_updates_to_db_increments_project_spend(): + """ + Test that _commit_spend_updates_to_db calls litellm_projecttable.update_many + with spend increment when project_list_transactions is present. + """ + db_writer = DBSpendUpdateWriter() + + mock_batcher = MagicMock() + mock_batcher.litellm_projecttable = MagicMock() + mock_batcher.litellm_projecttable.update_many = MagicMock() + + mock_transaction = AsyncMock() + mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction) + mock_transaction.__aexit__ = AsyncMock(return_value=False) + mock_transaction.batch_ = MagicMock( + return_value=AsyncMock( + __aenter__=AsyncMock(return_value=mock_batcher), + __aexit__=AsyncMock(return_value=False), + ) + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction) + + mock_proxy_logging = MagicMock() + + project_id = "project-789" + response_cost = 0.25 + db_spend_update_transactions = { + "user_list_transactions": {}, + "end_user_list_transactions": {}, + "key_list_transactions": {}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "project_list_transactions": {project_id: response_cost}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + + with patch("litellm.proxy.utils._raise_failed_update_spend_exception"): + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=mock_proxy_logging, + db_spend_update_transactions=db_spend_update_transactions, + ) + + mock_batcher.litellm_projecttable.update_many.assert_called_once() + call_kwargs = mock_batcher.litellm_projecttable.update_many.call_args[1] + assert call_kwargs["where"] == {"project_id": project_id} + assert call_kwargs["data"] == {"spend": {"increment": response_cost}} + + @pytest.mark.asyncio async def test_commit_spend_updates_to_db_increments_agent_spend(): """ @@ -1644,6 +1799,15 @@ async def test_commit_spend_updates_uses_pipeline(): ["agent_a", "agent_b", "agent_c"], id="agent", ), + pytest.param( + "project_list_transactions", + {"project_c": 0.1, "project_a": 0.2, "project_b": 0.3}, + "litellm_projecttable", + "update_many", + "project_id", + ["project_a", "project_b", "project_c"], + id="project", + ), ], ) @pytest.mark.asyncio @@ -1693,6 +1857,7 @@ async def test_commit_spend_updates_iterates_in_sorted_order( "team_list_transactions": {}, "team_member_list_transactions": {}, "org_list_transactions": {}, + "project_list_transactions": {}, "tag_list_transactions": {}, "agent_list_transactions": {}, } diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index f289148101a..12b014a46d4 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -32,6 +32,7 @@ async def test_async_post_call_failure_hook(): user_id="test_user_id", team_id="test_team_id", org_id="test_org_id", + project_id="test_project_id", team_alias="test_team_alias", end_user_id="test_end_user_id", ) @@ -70,6 +71,7 @@ async def test_async_post_call_failure_hook(): assert call_args["end_user_id"] == "test_end_user_id" assert call_args["team_id"] == "test_team_id" assert call_args["org_id"] == "test_org_id" + assert call_args["project_id"] == "test_project_id" assert call_args["completion_response"] == original_exception # Check that metadata was properly updated @@ -452,6 +454,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda end_user_id="test_end_user_id", team_id="test_team_id", org_id="test_org_id", + project_id="test_project_id", kwargs={}, completion_response=None, start_time=datetime.now(), @@ -462,6 +465,10 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda ) proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once() + assert ( + proxy_logging_obj.db_spend_update_writer.update_database.call_args[1]["project_id"] + == "test_project_id" + ) increment_spend_counters.assert_awaited_once_with( token="test_api_key", team_id="test_team_id",