diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.13-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.13-py3-none-any.whl new file mode 100644 index 00000000000..ff270dd9c37 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.13-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.13.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.13.tar.gz new file mode 100644 index 00000000000..92b6ab7ef2a Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.13.tar.gz differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251210205007_add_daily_agent_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251210205007_add_daily_agent_spend_table/migration.sql new file mode 100644 index 00000000000..964904c14c1 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251210205007_add_daily_agent_spend_table/migration.sql @@ -0,0 +1,45 @@ +-- AlterTable +ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "agent_id" TEXT; + +-- CreateTable +CREATE TABLE "LiteLLM_DailyAgentSpend" ( + "id" TEXT NOT NULL, + "agent_id" TEXT, + "date" TEXT NOT NULL, + "api_key" TEXT NOT NULL, + "model" TEXT, + "model_group" TEXT, + "custom_llm_provider" TEXT, + "mcp_namespaced_tool_name" TEXT, + "prompt_tokens" BIGINT NOT NULL DEFAULT 0, + "completion_tokens" BIGINT NOT NULL DEFAULT 0, + "cache_read_input_tokens" BIGINT NOT NULL DEFAULT 0, + "cache_creation_input_tokens" BIGINT NOT NULL DEFAULT 0, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "api_requests" BIGINT NOT NULL DEFAULT 0, + "successful_requests" BIGINT NOT NULL DEFAULT 0, + "failed_requests" BIGINT NOT NULL DEFAULT 0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_DailyAgentSpend_pkey" PRIMARY KEY ("id") +); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyAgentSpend_date_idx" ON "LiteLLM_DailyAgentSpend"("date"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyAgentSpend_agent_id_idx" ON "LiteLLM_DailyAgentSpend"("agent_id"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyAgentSpend_api_key_idx" ON "LiteLLM_DailyAgentSpend"("api_key"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyAgentSpend_model_idx" ON "LiteLLM_DailyAgentSpend"("model"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyAgentSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyAgentSpend"("mcp_namespaced_tool_name"); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key" ON "LiteLLM_DailyAgentSpend"("agent_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name"); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index f876d63520b..b76bb0401da 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -494,6 +494,34 @@ model LiteLLM_DailyEndUserSpend { @@index([mcp_namespaced_tool_name]) } +// Track daily agent spend metrics per model and key +model LiteLLM_DailyAgentSpend { + id String @id @default(uuid()) + agent_id String? + date String + api_key String + model String? + model_group String? + custom_llm_provider String? + mcp_namespaced_tool_name String? + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + @@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name]) + @@index([date]) + @@index([agent_id]) + @@index([api_key]) + @@index([model]) + @@index([mcp_namespaced_tool_name]) +} + // Track daily team spend metrics per model and key model LiteLLM_DailyTeamSpend { id String @id @default(uuid()) diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 908660f585d..4825b025fdb 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm-proxy-extras" -version = "0.4.12" +version = "0.4.13" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." authors = ["BerriAI"] readme = "README.md" @@ -22,7 +22,7 @@ requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "0.4.12" +version = "0.4.13" version_files = [ "pyproject.toml:version", "../requirements.txt:litellm-proxy-extras==", diff --git a/litellm/constants.py b/litellm/constants.py index 87e873e35bc..1dcbe073837 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -150,6 +150,7 @@ REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_spend_update_buffer" REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_team_spend_update_buffer" REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_org_spend_update_buffer" REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_end_user_spend_update_buffer" +REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_agent_spend_update_buffer" REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_tag_spend_update_buffer" MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)) MAX_SIZE_IN_MEMORY_QUEUE = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", 2000)) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a255168d747..a76482db464 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3654,6 +3654,9 @@ class DailyTagSpendTransaction(BaseDailySpendTransaction): request_id: Optional[str] tag: str +class DailyAgentSpendTransaction(BaseDailySpendTransaction): + agent_id: str + class DBSpendUpdateTransactions(TypedDict): """ diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 7b18a2380c0..4a8d615f0b3 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -8,7 +8,7 @@ Follows the A2A Spec. 3. Get specific agent via GET `/v1/agents/{agent_id}` """ -from typing import Any, List +from typing import Any, List, Optional from fastapi import APIRouter, Depends, HTTPException, Request @@ -24,6 +24,11 @@ from litellm.types.agents import ( PatchAgentRequest, ) +from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity +from litellm.types.proxy.management_endpoints.common_daily_activity import ( + SpendAnalyticsPaginatedResponse, +) + router = APIRouter() @@ -703,3 +708,64 @@ async def make_agents_public( except Exception as e: verbose_proxy_logger.exception(f"Error making agent public: {e}") raise HTTPException(status_code=500, detail=str(e)) + +@router.get( + "/agent/daily/activity", + tags=["Agent Management"], + dependencies=[Depends(user_api_key_auth)], + response_model=SpendAnalyticsPaginatedResponse, +) +async def get_agent_daily_activity( + agent_ids: Optional[str] = None, + start_date: Optional[str] = None, + end_date: Optional[str] = None, + model: Optional[str] = None, + api_key: Optional[str] = None, + page: int = 1, + page_size: int = 10, + exclude_agent_ids: Optional[str] = None, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Get daily activity for specific agents or all accessible agents. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + agent_ids_list = agent_ids.split(",") if agent_ids else None + exclude_agent_ids_list: Optional[List[str]] = None + if exclude_agent_ids: + exclude_agent_ids_list = ( + exclude_agent_ids.split(",") if exclude_agent_ids else None + ) + + where_condition = {} + if agent_ids_list: + where_condition["agent_id"] = {"in": list(agent_ids_list)} + + agent_records = await prisma_client.db.litellm_agentstable.find_many( + where=where_condition + ) + agent_metadata = { + agent.agent_id: {"agent_name": agent.agent_name} for agent in agent_records + } + + return await get_daily_activity( + prisma_client=prisma_client, + table_name="litellm_dailyagentspend", + entity_id_field="agent_id", + entity_id=agent_ids_list, + entity_metadata_field=agent_metadata, + exclude_entity_ids=exclude_agent_ids_list, + start_date=start_date, + end_date=end_date, + model=model, + api_key=api_key, + page=page, + page_size=page_size, + ) \ No newline at end of file diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index c8c1726c11a..55a6eee7b7a 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -28,6 +28,7 @@ from litellm.proxy._types import ( DailyTeamSpendTransaction, DailyEndUserSpendTransaction, DailyUserSpendTransaction, + DailyAgentSpendTransaction, DBSpendUpdateTransactions, Litellm_EntityType, LiteLLM_UserTable, @@ -68,6 +69,7 @@ class DBSpendUpdateWriter: self.daily_spend_update_queue = DailySpendUpdateQueue() self.daily_team_spend_update_queue = DailySpendUpdateQueue() self.daily_end_user_spend_update_queue = DailySpendUpdateQueue() + self.daily_agent_spend_update_queue = DailySpendUpdateQueue() self.daily_org_spend_update_queue = DailySpendUpdateQueue() self.daily_tag_spend_update_queue = DailySpendUpdateQueue() @@ -192,6 +194,13 @@ class DBSpendUpdateWriter: ) ) + asyncio.create_task( + self.add_spend_log_transaction_to_daily_agent_transaction( + payload=payload, + prisma_client=prisma_client, + ) + ) + asyncio.create_task( self.add_spend_log_transaction_to_daily_team_transaction( payload=copy.deepcopy(payload), @@ -488,6 +497,7 @@ class DBSpendUpdateWriter: daily_team_spend_update_queue=self.daily_team_spend_update_queue, daily_org_spend_update_queue=self.daily_org_spend_update_queue, daily_end_user_spend_update_queue=self.daily_end_user_spend_update_queue, + daily_agent_spend_update_queue=self.daily_agent_spend_update_queue, daily_tag_spend_update_queue=self.daily_tag_spend_update_queue, ) @@ -561,6 +571,16 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_end_user_spend_update_transactions, ) + daily_agent_spend_update_transactions = ( + await self.redis_update_buffer.get_all_daily_agent_spend_update_transactions_from_redis_buffer() + ) + if daily_agent_spend_update_transactions is not None: + await DBSpendUpdateWriter.update_daily_agent_spend( + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + daily_spend_transactions=daily_agent_spend_update_transactions, + ) except Exception as e: verbose_proxy_logger.error(f"Error committing spend updates: {e}") finally: @@ -664,6 +684,20 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_end_user_spend_update_transactions, ) + ################## Daily Agent Spend Update Transactions ################## + # Aggregate all in memory daily agent spend transactions and commit to db + daily_agent_spend_update_transactions = cast( + Dict[str, DailyAgentSpendTransaction], + await self.daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(), + ) + + await DBSpendUpdateWriter.update_daily_agent_spend( + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + daily_spend_transactions=daily_agent_spend_update_transactions, + ) + async def _commit_spend_updates_to_db( # noqa: PLR0915 self, prisma_client: PrismaClient, @@ -1041,6 +1075,20 @@ class DBSpendUpdateWriter: ) -> None: ... + @overload + @staticmethod + async def _update_daily_spend( + n_retry_times: int, + prisma_client: PrismaClient, + proxy_logging_obj: ProxyLogging, + daily_spend_transactions: Dict[str, DailyAgentSpendTransaction], + entity_type: Literal["agent"], + entity_id_field: str, + table_name: str, + unique_constraint_name: str, + ) -> None: + ... + @overload @staticmethod async def _update_daily_spend( @@ -1067,14 +1115,15 @@ class DBSpendUpdateWriter: Dict[str, DailyTagSpendTransaction], Dict[str, DailyOrganizationSpendTransaction], Dict[str, DailyEndUserSpendTransaction], + Dict[str, DailyAgentSpendTransaction], ], - entity_type: Literal["user", "team", "org", "tag", "end_user"], + entity_type: Literal["user", "team", "org", "tag", "end_user", "agent"], entity_id_field: str, table_name: str, unique_constraint_name: str, ) -> None: """ - Generic function to update daily spend for any entity type (user, team, org, tag, end_user) + Generic function to update daily spend for any entity type (user, team, org, tag, end_user, agent) """ from litellm.proxy.utils import _raise_failed_update_spend_exception @@ -1340,6 +1389,27 @@ class DBSpendUpdateWriter: unique_constraint_name="end_user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name", ) + @staticmethod + async def update_daily_agent_spend( + n_retry_times: int, + prisma_client: PrismaClient, + proxy_logging_obj: ProxyLogging, + daily_spend_transactions: Dict[str, DailyAgentSpendTransaction], + ): + """ + Batch job to update LiteLLM_DailyAgentSpend table using in-memory daily_spend_transactions + """ + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + daily_spend_transactions=daily_spend_transactions, + entity_type="agent", + entity_id_field="agent_id", + table_name="litellm_dailyagentspend", + unique_constraint_name="agent_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name", + ) + @staticmethod async def update_daily_tag_spend( n_retry_times: int, @@ -1365,7 +1435,7 @@ class DBSpendUpdateWriter: self, payload: Union[dict, SpendLogsPayload], prisma_client: PrismaClient, - type: Literal["user", "team", "org", "request_tags", "end_user"] = "user", + type: Literal["user", "team", "org", "request_tags", "end_user", "agent"] = "user", ) -> Optional[BaseDailySpendTransaction]: common_expected_keys = ["startTime", "api_key"] if type == "user": @@ -1378,6 +1448,8 @@ class DBSpendUpdateWriter: expected_keys = ["request_tags", *common_expected_keys] elif type == "end_user": expected_keys = ["end_user_id", *common_expected_keys] + elif type == "agent": + expected_keys = ["agent_id", *common_expected_keys] else: raise ValueError(f"Invalid type: {type}") if not all(key in payload for key in expected_keys): @@ -1591,6 +1663,50 @@ class DBSpendUpdateWriter: update={daily_transaction_key: daily_transaction} ) + async def add_spend_log_transaction_to_daily_agent_transaction( + self, + payload: SpendLogsPayload, + prisma_client: Optional[PrismaClient] = None, + ) -> None: + if prisma_client is None: + verbose_proxy_logger.debug( + "prisma_client is None. Skipping writing spend logs to db." + ) + return + base_daily_transaction = ( + await self._common_add_spend_log_transaction_to_daily_transaction( + payload, prisma_client, "agent" + ) + ) + if base_daily_transaction is None: + return + if payload["agent_id"] is None: + verbose_proxy_logger.debug( + "agent_id is None for request. Skipping incrementing agent spend." + ) + return + payload_with_agent_id = cast( + SpendLogsPayload, + { + **payload, + "agent_id": payload["agent_id"], + }, + ) + base_daily_transaction = ( + await self._common_add_spend_log_transaction_to_daily_transaction( + payload_with_agent_id, prisma_client, "agent" + ) + ) + if base_daily_transaction is None: + return + daily_transaction_key = f"{payload['agent_id']}_{base_daily_transaction['date']}_{payload_with_agent_id['api_key']}_{payload_with_agent_id['model']}_{payload_with_agent_id['custom_llm_provider']}" + daily_transaction = DailyAgentSpendTransaction( + agent_id=payload['agent_id'], **base_daily_transaction + ) + await self.daily_agent_spend_update_queue.add_update( + update={daily_transaction_key: daily_transaction} + ) + async def add_spend_log_transaction_to_daily_tag_transaction( self, payload: SpendLogsPayload, 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 e3b20d7266d..37b42e26bc9 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -17,6 +17,7 @@ from litellm.constants import ( REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, REDIS_UPDATE_BUFFER_KEY, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -27,6 +28,7 @@ from litellm.proxy._types import ( DailyOrganizationSpendTransaction, DailyEndUserSpendTransaction, DBSpendUpdateTransactions, + DailyAgentSpendTransaction, ) from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( @@ -110,6 +112,7 @@ class RedisUpdateBuffer: daily_team_spend_update_queue: DailySpendUpdateQueue, daily_org_spend_update_queue: DailySpendUpdateQueue, daily_end_user_spend_update_queue: DailySpendUpdateQueue, + daily_agent_spend_update_queue: DailySpendUpdateQueue, daily_tag_spend_update_queue: DailySpendUpdateQueue, ): """ @@ -178,6 +181,9 @@ class RedisUpdateBuffer: daily_end_user_spend_update_transactions = ( await daily_end_user_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) + daily_agent_spend_update_transactions = ( + await daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) daily_tag_spend_update_transactions = ( await daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) @@ -219,6 +225,12 @@ class RedisUpdateBuffer: service_type=ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE, ) + await self._store_transactions_in_redis( + transactions=daily_agent_spend_update_transactions, + redis_key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, + service_type=ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE, + ) + await self._store_transactions_in_redis( transactions=daily_tag_spend_update_transactions, redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, @@ -401,6 +413,30 @@ class RedisUpdateBuffer: ), ) + async def get_all_daily_agent_spend_update_transactions_from_redis_buffer( + self, + ) -> Optional[Dict[str, DailyAgentSpendTransaction]]: + """ + Gets all the daily agent spend update transactions from Redis + """ + if self.redis_cache is None: + return None + list_of_transactions = await self.redis_cache.async_lpop( + key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, + count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + ) + if list_of_transactions is None: + return None + list_of_daily_spend_update_transactions = [ + json.loads(transaction) for transaction in list_of_transactions + ] + return cast( + Dict[str, DailyAgentSpendTransaction], + DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( + list_of_daily_spend_update_transactions + ), + ) + async def get_all_daily_tag_spend_update_transactions_from_redis_buffer( self, ) -> Optional[Dict[str, DailyTagSpendTransaction]]: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index f876d63520b..b76bb0401da 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -494,6 +494,34 @@ model LiteLLM_DailyEndUserSpend { @@index([mcp_namespaced_tool_name]) } +// Track daily agent spend metrics per model and key +model LiteLLM_DailyAgentSpend { + id String @id @default(uuid()) + agent_id String? + date String + api_key String + model String? + model_group String? + custom_llm_provider String? + mcp_namespaced_tool_name String? + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + @@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name]) + @@index([date]) + @@index([agent_id]) + @@index([api_key]) + @@index([model]) + @@index([mcp_namespaced_tool_name]) +} + // Track daily team spend metrics per model and key model LiteLLM_DailyTeamSpend { id String @id @default(uuid()) diff --git a/pyproject.toml b/pyproject.toml index 8093ad98761..3a47ba7cd08 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,7 +59,7 @@ websockets = {version = "^15.0.1", optional = true} boto3 = {version = "1.36.0", optional = true} redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"} mcp = {version = "^1.21.2", optional = true, python = ">=3.10"} -litellm-proxy-extras = {version = "0.4.12", optional = true} +litellm-proxy-extras = {version = "0.4.13", optional = true} rich = {version = "13.7.1", optional = true} litellm-enterprise = {version = "0.1.25", optional = true} diskcache = {version = "^5.6.1", optional = true} diff --git a/requirements.txt b/requirements.txt index 0293d5579a5..d7458727ae3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -44,7 +44,7 @@ sentry_sdk==2.21.0 # for sentry error handling detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests cryptography==44.0.1 tzdata==2025.1 # IANA time zone database -litellm-proxy-extras==0.4.12 # for proxy extras - e.g. prisma migrations +litellm-proxy-extras==0.4.13 # for proxy extras - e.g. prisma migrations ### LITELLM PACKAGE DEPENDENCIES python-dotenv==1.0.1 # for env tiktoken==0.8.0 # for calculating usage diff --git a/schema.prisma b/schema.prisma index f876d63520b..b76bb0401da 100644 --- a/schema.prisma +++ b/schema.prisma @@ -494,6 +494,34 @@ model LiteLLM_DailyEndUserSpend { @@index([mcp_namespaced_tool_name]) } +// Track daily agent spend metrics per model and key +model LiteLLM_DailyAgentSpend { + id String @id @default(uuid()) + agent_id String? + date String + api_key String + model String? + model_group String? + custom_llm_provider String? + mcp_namespaced_tool_name String? + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + @@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name]) + @@index([date]) + @@index([agent_id]) + @@index([api_key]) + @@index([model]) + @@index([mcp_namespaced_tool_name]) +} + // Track daily team spend metrics per model and key model LiteLLM_DailyTeamSpend { id String @id @default(uuid()) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py new file mode 100644 index 00000000000..8cec3538077 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -0,0 +1,260 @@ +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.agent_endpoints import endpoints as agent_endpoints +from litellm.proxy.agent_endpoints.endpoints import ( + get_agent_daily_activity, + router, + user_api_key_auth, +) +from litellm.types.agents import AgentResponse + + +def _sample_agent_card_params() -> dict: + return { + "protocolVersion": "1.0", + "name": "Test Agent", + "description": "desc", + "url": "http://localhost", + "version": "1.0.0", + "capabilities": {"streaming": True}, + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "skills": [], + } + + +def _sample_agent_config() -> dict: + return { + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {"make_public": False}, + } + + +def _sample_agent_response( + agent_id: str = "agent-123", agent_name: str = "Test Agent" +) -> AgentResponse: + return AgentResponse( + agent_id=agent_id, + agent_name=agent_name, + agent_card_params=_sample_agent_card_params(), + litellm_params={"make_public": False}, + ) + + +app = FastAPI() +app.include_router(router) +app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user", user_role=LitellmUserRoles.PROXY_ADMIN +) +client = TestClient(app) + + +@pytest.fixture +def mock_prisma_client(): + with patch("litellm.proxy.proxy_server.prisma_client") as mock: + yield mock + + +@pytest.fixture +def mock_user_api_key_auth(): + with patch("litellm.proxy.agent_endpoints.endpoints.user_api_key_auth") as mock: + mock.return_value = UserAPIKeyAuth( + user_id="test-user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + yield mock + + +def test_update_agent_success(mock_prisma_client, mock_user_api_key_auth, monkeypatch): + existing_agent = { + "agent_id": "agent-123", + "agent_name": "Existing Agent", + "agent_card_params": _sample_agent_card_params(), + } + mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock( + return_value=existing_agent + ) + + mock_registry = MagicMock() + mock_registry.update_agent_in_db = AsyncMock( + return_value=_sample_agent_response(agent_id="agent-123") + ) + mock_registry.deregister_agent = MagicMock() + mock_registry.register_agent = MagicMock() + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry) + + response = client.put( + "/v1/agents/agent-123", + json=_sample_agent_config(), + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 200 + assert response.json()["agent_id"] == "agent-123" + assert response.json()["agent_name"] == "Test Agent" + + +def test_update_agent_not_found( + mock_prisma_client, mock_user_api_key_auth, monkeypatch +): + mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None) + + mock_registry = MagicMock() + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry) + + response = client.put( + "/v1/agents/missing-agent", + json=_sample_agent_config(), + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 404 + assert "Agent with ID missing-agent not found" in response.json()["detail"] + + +def test_get_agent_by_id_not_found( + mock_prisma_client, mock_user_api_key_auth, monkeypatch +): + mock_registry = MagicMock() + mock_registry.get_agent_by_id = MagicMock(return_value=None) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry) + mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None) + + response = client.get( + "/v1/agents/missing-agent", headers={"Authorization": "Bearer test-key"} + ) + + assert response.status_code == 404 + assert "Agent with ID missing-agent not found" in response.json()["detail"] + + +def test_delete_agent_not_found( + mock_prisma_client, mock_user_api_key_auth, monkeypatch +): + mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None) + mock_registry = MagicMock() + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry) + + response = client.delete( + "/v1/agents/missing-agent", headers={"Authorization": "Bearer test-key"} + ) + + assert response.status_code == 404 + assert "Agent with ID missing-agent not found in DB." in response.json()["detail"] + + +def test_agent_error_schema_consistency( + mock_prisma_client, mock_user_api_key_auth, monkeypatch +): + mock_registry = MagicMock() + mock_registry.get_agent_by_id = MagicMock(return_value=None) + mock_registry.update_agent_in_db = AsyncMock(side_effect=Exception("should not run")) + mock_registry.delete_agent_from_db = AsyncMock(side_effect=Exception("should not run")) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry) + + mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None) + + missing_agent_id = "missing-agent" + responses = [ + client.get( + f"/v1/agents/{missing_agent_id}", + headers={"Authorization": "Bearer test-key"}, + ), + client.put( + f"/v1/agents/{missing_agent_id}", + json=_sample_agent_config(), + headers={"Authorization": "Bearer test-key"}, + ), + client.delete( + f"/v1/agents/{missing_agent_id}", + headers={"Authorization": "Bearer test-key"}, + ), + ] + + for resp in responses: + assert resp.status_code == 404 + detail = resp.json()["detail"] + assert isinstance(detail, str) + assert missing_agent_id in detail + + +@pytest.mark.asyncio +async def test_get_agent_daily_activity_admin_param_passing(monkeypatch): + mock_prisma = AsyncMock() + mock_prisma.db.litellm_agentstable.find_many = AsyncMock(return_value=[]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse") + get_daily_activity_mock = AsyncMock(return_value=mocked_response) + monkeypatch.setattr(agent_endpoints, "get_daily_activity", get_daily_activity_mock) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1") + result = await get_agent_daily_activity( + agent_ids="agent-1,agent-2", + start_date="2024-01-01", + end_date="2024-01-31", + model="gpt-4", + api_key="test-key", + page=2, + page_size=5, + exclude_agent_ids="agent-3", + user_api_key_dict=auth, + ) + + get_daily_activity_mock.assert_awaited_once() + kwargs = get_daily_activity_mock.call_args.kwargs + assert kwargs["table_name"] == "litellm_dailyagentspend" + assert kwargs["entity_id_field"] == "agent_id" + assert kwargs["entity_id"] == ["agent-1", "agent-2"] + assert kwargs["exclude_entity_ids"] == ["agent-3"] + assert kwargs["start_date"] == "2024-01-01" + assert kwargs["end_date"] == "2024-01-31" + assert kwargs["model"] == "gpt-4" + assert kwargs["api_key"] == "test-key" + assert kwargs["page"] == 2 + assert kwargs["page_size"] == 5 + assert result is mocked_response + + +@pytest.mark.asyncio +async def test_get_agent_daily_activity_with_agent_names(monkeypatch): + mock_prisma = AsyncMock() + mock_agent1 = MagicMock() + mock_agent1.agent_id = "agent-1" + mock_agent1.agent_name = "First Agent" + mock_agent2 = MagicMock() + mock_agent2.agent_id = "agent-2" + mock_agent2.agent_name = "Second Agent" + + mock_prisma.db.litellm_agentstable.find_many = AsyncMock( + return_value=[mock_agent1, mock_agent2] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse") + get_daily_activity_mock = AsyncMock(return_value=mocked_response) + monkeypatch.setattr(agent_endpoints, "get_daily_activity", get_daily_activity_mock) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1") + await get_agent_daily_activity( + agent_ids="agent-1,agent-2", + start_date="2024-01-01", + end_date="2024-01-31", + model=None, + api_key=None, + page=1, + page_size=10, + exclude_agent_ids=None, + user_api_key_dict=auth, + ) + + kwargs = get_daily_activity_mock.call_args.kwargs + assert kwargs["entity_metadata_field"] == { + "agent-1": {"agent_name": "First Agent"}, + "agent-2": {"agent_name": "Second Agent"}, + } 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 db6c318357c..634b90ad9e1 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 @@ -645,4 +645,84 @@ async def test_add_spend_log_transaction_to_daily_end_user_transaction_skips_whe prisma_client=mock_prisma, ) - writer.daily_end_user_spend_update_queue.add_update.assert_not_called() \ No newline at end of file + writer.daily_end_user_spend_update_queue.add_update.assert_not_called() + + +@pytest.mark.asyncio +async def test_add_spend_log_transaction_to_daily_agent_transaction_injects_agent_id_and_queues_update(): + """ + Ensure agent_id is injected and queued for daily aggregation. + """ + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + mock_prisma.get_request_status = MagicMock(return_value="success") + + agent_id = "agent-123" + payload = { + "request_id": "req-123", + "agent_id": agent_id, + "user": "test-user", + "startTime": "2024-01-01T12:00:00", + "api_key": "test-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "model_group": "gpt-4-group", + "prompt_tokens": 20, + "completion_tokens": 10, + "spend": 0.3, + "metadata": '{"usage_object": {}}', + } + + writer.daily_agent_spend_update_queue.add_update = AsyncMock() + + await writer.add_spend_log_transaction_to_daily_agent_transaction( + payload=payload, + prisma_client=mock_prisma, + ) + + writer.daily_agent_spend_update_queue.add_update.assert_called_once() + + call_args = writer.daily_agent_spend_update_queue.add_update.call_args[1] + update_dict = call_args["update"] + assert len(update_dict) == 1 + for key, transaction in update_dict.items(): + assert key == f"{agent_id}_2024-01-01_test-key_gpt-4_openai" + assert transaction["agent_id"] == agent_id + assert transaction["date"] == "2024-01-01" + assert transaction["api_key"] == "test-key" + assert transaction["model"] == "gpt-4" + assert transaction["custom_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_add_spend_log_transaction_to_daily_agent_transaction_skips_when_agent_id_missing(): + """ + Do not queue agent spend updates when agent_id is None. + """ + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + mock_prisma.get_request_status = MagicMock(return_value="success") + + payload = { + "request_id": "req-456", + "agent_id": None, + "user": "test-user", + "startTime": "2024-01-01T12:00:00", + "api_key": "test-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "model_group": "gpt-4-group", + "prompt_tokens": 15, + "completion_tokens": 5, + "spend": 0.1, + "metadata": '{"usage_object": {}}', + } + + writer.daily_agent_spend_update_queue.add_update = AsyncMock() + + await writer.add_spend_log_transaction_to_daily_agent_transaction( + payload=payload, + prisma_client=mock_prisma, + ) + + writer.daily_agent_spend_update_queue.add_update.assert_not_called() \ No newline at end of file