diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304184926_baseline_diff/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304184926_baseline_diff/migration.sql new file mode 100644 index 00000000000..2f725d83806 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304184926_baseline_diff/migration.sql @@ -0,0 +1,2 @@ +-- This is an empty migration. + diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 65674d01be7..40480bec3e2 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/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, - ) \ No newline at end of file + )