feat: return spend on agent endpoints

litellm-proxy-extras/litellm_proxy_extras/migrations/20260304184926_baseline_diff/migration.sql

@@ -0,0 +1,2 @@
+-- This is an empty migration.
+

litellm/proxy/agent_endpoints/endpoints.py

@@ -14,19 +14,16 @@ from fastapi import APIRouter, Depends, HTTPException, Request
 
 import litellm
 from litellm._logging import verbose_proxy_logger
-from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
+from litellm.proxy._types import (CommonProxyErrors, LitellmUserRoles,
+                                  UserAPIKeyAuth)
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
-from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity
-from litellm.types.agents import (
-    AgentConfig,
-    AgentMakePublicResponse,
-    AgentResponse,
-    MakeAgentsPublicRequest,
-    PatchAgentRequest,
-)
-from litellm.types.proxy.management_endpoints.common_daily_activity import (
-    SpendAnalyticsPaginatedResponse,
-)
+from litellm.proxy.management_endpoints.common_daily_activity import \
+    get_daily_activity
+from litellm.types.agents import (AgentConfig, AgentMakePublicResponse,
+                                  AgentResponse, MakeAgentsPublicRequest,
+                                  PatchAgentRequest)
+from litellm.types.proxy.management_endpoints.common_daily_activity import \
+    SpendAnalyticsPaginatedResponse
 
 router = APIRouter()
 
@@ -69,14 +66,14 @@ async def get_agents(
     Returns: List[AgentResponse]
 
     """
-    from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
-    from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
-        AgentRequestHandler,
-    )
+    from litellm.proxy.agent_endpoints.agent_registry import \
+        global_agent_registry
+    from litellm.proxy.agent_endpoints.auth.agent_permission_handler import \
+        AgentRequestHandler
 
     try:
         returned_agents: List[AgentResponse] = []
-        
+
         # Admin users get all agents
         if (
             user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
@@ -88,7 +85,7 @@ async def get_agents(
             allowed_agent_ids = await AgentRequestHandler.get_allowed_agents(
                 user_api_key_auth=user_api_key_dict
             )
-            
+
             # If no restrictions (empty list), return all agents
             if len(allowed_agent_ids) == 0:
                 returned_agents = global_agent_registry.get_agent_list()
@@ -96,10 +93,23 @@ async def get_agents(
                 # Filter agents by allowed IDs
                 all_agents = global_agent_registry.get_agent_list()
                 returned_agents = [
-                    agent for agent in all_agents
-                    if agent.agent_id in allowed_agent_ids
+                    agent for agent in all_agents if agent.agent_id in allowed_agent_ids
                 ]
 
+        # Fetch current spend from DB for all returned agents
+        from litellm.proxy.proxy_server import prisma_client
+
+        if prisma_client is not None:
+            agent_ids = [agent.agent_id for agent in returned_agents]
+            if agent_ids:
+                db_agents = await prisma_client.db.litellm_agentstable.find_many(
+                    where={"agent_id": {"in": agent_ids}},
+                )
+                spend_map = {a.agent_id: a.spend for a in db_agents}
+                for agent in returned_agents:
+                    if agent.agent_id in spend_map:
+                        agent.spend = spend_map[agent.agent_id]
+
         # add is_public field to each agent - we do it this way, to allow setting config agents as public
         for agent in returned_agents:
             if agent.litellm_params is None:
@@ -125,9 +135,8 @@ async def get_agents(
 
 #### CRUD ENDPOINTS FOR AGENTS ####
 
-from litellm.proxy.agent_endpoints.agent_registry import (
-    global_agent_registry as AGENT_REGISTRY,
-)
+from litellm.proxy.agent_endpoints.agent_registry import \
+    global_agent_registry as AGENT_REGISTRY
 
 
 @router.post(
@@ -259,10 +268,21 @@ async def get_agent_by_id(agent_id: str):
                 agent_dict = agent_row.model_dump()
                 if agent_row.object_permission is not None:
                     try:
-                        agent_dict["object_permission"] = agent_row.object_permission.model_dump()
+                        agent_dict["object_permission"] = (
+                            agent_row.object_permission.model_dump()
+                        )
                     except Exception:
-                        agent_dict["object_permission"] = agent_row.object_permission.dict()
+                        agent_dict["object_permission"] = (
+                            agent_row.object_permission.dict()
+                        )
                 agent = AgentResponse(**agent_dict)  # type: ignore
+        else:
+            # Agent found in memory — refresh spend from DB
+            db_row = await prisma_client.db.litellm_agentstable.find_unique(
+                where={"agent_id": agent_id}
+            )
+            if db_row is not None:
+                agent.spend = db_row.spend
 
         if agent is None:
             raise HTTPException(
@@ -564,9 +584,8 @@ async def make_agent_public(
     try:
         # Update the public model groups
         import litellm
-        from litellm.proxy.agent_endpoints.agent_registry import (
-            global_agent_registry as AGENT_REGISTRY,
-        )
+        from litellm.proxy.agent_endpoints.agent_registry import \
+            global_agent_registry as AGENT_REGISTRY
         from litellm.proxy.proxy_server import proxy_config
 
         # Check if user has admin permissions
@@ -681,9 +700,8 @@ async def make_agents_public(
     try:
         # Update the public model groups
         import litellm
-        from litellm.proxy.agent_endpoints.agent_registry import (
-            global_agent_registry as AGENT_REGISTRY,
-        )
+        from litellm.proxy.agent_endpoints.agent_registry import \
+            global_agent_registry as AGENT_REGISTRY
         from litellm.proxy.proxy_server import proxy_config
 
         # Load existing config
@@ -743,6 +761,7 @@ async def make_agents_public(
         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"],
@@ -802,4 +821,4 @@ async def get_agent_daily_activity(
         api_key=api_key,
         page=page,
         page_size=page_size,
-    )
+    )